Matplotlib 3D scatter animations

I am graphing out positions in a star cluster, my data is in a dataframe with x,y,z positions as well as a time index.

I am able to produce a 3d scatter plot and was trying to produce a rotating plot--I have been somewhat successful, but struggling through the animation API.

If my "update_graph" function just returns a new ax.scatter(), the old one stays plotted unless i rebuild the entire graph. That seems inefficient. As well, I have to set my interval rather high or my animation "skips" every other frame, so it says my performance is rather bad. Finally I am forced to use the "blit=False" as I cannot get an iterator for a 3d scatter plot. Apparently the "graph.set_data()" doesn't work, and I can use the "graph.set_3d_properties" but that allows me new z coordinates only.

So i have cobbled together a cluuge-- (data i used is at https://www.kaggle.com/mariopasquato/star-cluster-simulations scroll to bottom)

Also I am only plotting 100 points (data=data[data.id<100])

My (working) code is as follows:

def update_graph(num):
     ax = p3.Axes3D(fig)
     ax.set_xlim3d([-5.0, 5.0])
     ax.set_xlabel('X')
     ax.set_ylim3d([-5.0, 5.0])
     ax.set_ylabel('Y')
     ax.set_zlim3d([-5.0, 5.0])
     ax.set_zlabel('Z')
     title='3D Test, Time='+str(num*100)
     ax.set_title(title)
     sample=data0[data0['time']==num*100]
     x=sample.x
     y=sample.y
     z=sample.z
     graph=ax.scatter(x,y,z)
     return(graph)

fig = plt.figure()
ax = p3.Axes3D(fig)

# Setting the axes properties
ax.set_xlim3d([-5.0, 5.0])
ax.set_xlabel('X')
ax.set_ylim3d([-5.0, 5.0])
ax.set_ylabel('Y')
ax.set_zlim3d([-5.0, 5.0])
ax.set_zlabel('Z')
ax.set_title('3D Test')
data=data0[data0['time']==0]
x=data.x
y=data.y
z=data.z
graph=ax.scatter(x,y,z)

# Creating the Animation object
line_ani = animation.FuncAnimation(fig, update_graph, 19, 
                               interval=350, blit=False)
plt.show()

Solution 1:

The scatter plot in 3D is a mpl_toolkits.mplot3d.art3d.Path3DCollection object. This provides an attribute _offsets3d which hosts a tuple (x,y,z) and can be used to update the scatter points' coordinates. Therefore it may be beneficial not to create the whole plot on every iteration of the animation, but instead only update its points.

The following is a working example on how to do this.

import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation
import pandas as pd


a = np.random.rand(2000, 3)*10
t = np.array([np.ones(100)*i for i in range(20)]).flatten()
df = pd.DataFrame({"time": t ,"x" : a[:,0], "y" : a[:,1], "z" : a[:,2]})

def update_graph(num):
    data=df[df['time']==num]
    graph._offsets3d = (data.x, data.y, data.z)
    title.set_text('3D Test, time={}'.format(num))


fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
title = ax.set_title('3D Test')

data=df[df['time']==0]
graph = ax.scatter(data.x, data.y, data.z)

ani = matplotlib.animation.FuncAnimation(fig, update_graph, 19, 
                               interval=40, blit=False)

plt.show()

This solution does not allow for blitting. However, depending on the usage case, it may not be necessary to use a scatter plot at all; using a normal plot might be equally possible, which allows for blitting - as seen in the following example.

import numpy as np
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation
import pandas as pd


a = np.random.rand(2000, 3)*10
t = np.array([np.ones(100)*i for i in range(20)]).flatten()
df = pd.DataFrame({"time": t ,"x" : a[:,0], "y" : a[:,1], "z" : a[:,2]})

def update_graph(num):
    data=df[df['time']==num]
    graph.set_data (data.x, data.y)
    graph.set_3d_properties(data.z)
    title.set_text('3D Test, time={}'.format(num))
    return title, graph, 


fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
title = ax.set_title('3D Test')

data=df[df['time']==0]
graph, = ax.plot(data.x, data.y, data.z, linestyle="", marker="o")

ani = matplotlib.animation.FuncAnimation(fig, update_graph, 19, 
                               interval=40, blit=True)

plt.show()

Solution 2:

If using Jupyter Notebook remember to use %matplotlib notebook don't use %matplotlib inline.