Matplotlib 3D scatter animations

Question

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.

Comments

[RDS]

Thats the thing I needed--i knew redrawing the whole thing wasn't right, but it was the only thing I could get working. As a "teach a man to fish" follow up, I scoured the matplot docs and could not find any mention of of _offsets3d. Do I need to be including github and the code itself for future reference?

[ImportanceOfBeingErnest]

This is a good point. The upper solution uses _offsets3d which is a private method (as indicated by the _ in front). Those are not documented. One will only find them by looking through the source code or by finding usage examples online. I often find it helpful to look at the source if the documentation does not provide a solution. Not to forget, asking a specific question like you did also helps finding hidden gems like _offsets3d. ;-) However, as private functions might also change between versions, their use should in general be avoided.

[whiletrue]

how can additional data points over time be added to the 3d offset?

[ImportanceOfBeingErnest]

@FábioFerreira The length of the arrays is not restricted. You can supply different length arrays from frame to frame, just within each frame x,y,z need to have the same length of course.