Animate A Collection Of Patches In Matplotlib

I'm trying to animate a set of particles that follow trajectories in x, y and z. Each object has a specific radius which is relevant in axis units, which is why I want each object

Solution 1:

The idea would of course be to add the collection only once to the axes. Then change the artists inside the collection via collection.set_paths.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as anm
import matplotlib.collections as clt

fig, ax = plt.subplots(1,1,figsize=(7,7))

ax.set_xlim(-1,1)
ax.set_ylim(-1,1)

n_of_particles = 3
frames = 10

radius = 0.05
x = 0.5*np.random.randn(frames,n_of_particles)
y = 0.5*np.random.randn(frames,n_of_particles)
z = 0.5*np.random.randn(frames,n_of_particles)

patches = []
for p in range(n_of_particles):
    circle = plt.Circle((x[0,p], y[0,p]), radius)
    patches.append(circle)

collection = clt.PatchCollection(patches, cmap=plt.cm.jet, alpha=0.4)
collection.set_array(z[0,:])
collection.set_clim([-1, 1])
fig.colorbar(collection)

ax.add_collection(collection)

def animate(frame):
    patches = []
    for p in range(n_of_particles):
        circle = plt.Circle((x[frame,p], y[frame,p]), radius)
        patches.append(circle)

    collection.set_paths(patches)
    collection.set_array(z[frame,:])

anim = anm.FuncAnimation(fig, animate,
                               frames=10, interval=1000, blit=False)
plt.show()

The other option could be to use the solution provided in this answer: matplotlib change a Patch in PatchCollection
and create an UpdatablePatchCollection. This would allow to just update the properties of the original patches inside the loop.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as anm
import matplotlib.collections as clt

fig, ax = plt.subplots(1,1,figsize=(7,7))

ax.set_xlim(-1,1)
ax.set_ylim(-1,1)

n_of_particles = 3
frames = 10

radius = 0.05
x = 0.5*np.random.randn(frames,n_of_particles)
y = 0.5*np.random.randn(frames,n_of_particles)
z = 0.5*np.random.randn(frames,n_of_particles)

patches = []
for p in range(n_of_particles):
    circle = plt.Circle((x[0,p], y[0,p]), radius)
    patches.append(circle)

class UpdatablePatchCollection(clt.PatchCollection):
    def __init__(self, patches, *args, **kwargs):
        self.patches = patches
        clt.PatchCollection.__init__(self, patches, *args, **kwargs)

    def get_paths(self):
        self.set_paths(self.patches)
        return self._paths

collection = UpdatablePatchCollection(patches, cmap=plt.cm.jet, alpha=0.4)
collection.set_array(z[0,:])
collection.set_clim([-1, 1])
fig.colorbar(collection)

ax.add_collection(collection)

def animate(frame):
    for p in range(n_of_particles):
        patches[p].center = x[frame,p], y[frame,p]   
    collection.set_array(z[frame,:])

anim = anm.FuncAnimation(fig, animate,
                               frames=10, interval=1000, blit=False)

plt.show()