Gaussian roll - matplotlib animation for 3d line plots with colored segments

gaussian_roll.py

from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Line3DCollection

fig, ax = plt.subplots(subplot_kw = dict(projection = '3d'))

# generate data
x = np.linspace(-5, 5, 500)
y = np.linspace(-5, 5, 500)
z = np.exp(-(x - 2)**2)

# uggly
segs = np.array([[(x1,y2), (x2, y2), (z1, z2)] for x1, x2, y1, y2, z1, z2 in zip(x[:-1], x[1:], y[:-1], y[1:], z[:-1], z[1:])])
segs = np.moveaxis(segs, 1, 2)

# setup segments

# get bounds
bounds_min = segs.reshape(-1, 3).min(0)
bounds_max = segs.reshape(-1, 3).max(0)

# setup colorbar stuff
# get bounds of colors
norm = plt.cm.colors.Normalize(bounds_min[2], bounds_max[2])
cmap = plt.cm.plasma
# setup scalar mappable for colorbar
sm   = plt.cm.ScalarMappable(norm, plt.cm.plasma)

# get average of segment
avg = segs.mean(1)[..., -1]
# get colors
colors = cmap(norm(avg))
# generate colors
lc = Line3DCollection(segs, norm = norm, cmap = cmap, colors = colors)
ax.add_collection(lc)

def update(idx):
    segs[..., -1] = np.roll(segs[..., -1], idx)
    lc.set_offsets(segs)
    return lc

ax.set_xlim(bounds_min[0], bounds_max[0])
ax.set_ylim(bounds_min[1], bounds_max[1])
ax.set_zlim(bounds_min[2], bounds_max[2])
fig.colorbar(sm)

from matplotlib import animation
frames = np.linspace(0, 30, 10, 0).astype(int)
ani = animation.FuncAnimation(fig, update, frames = frames)
ani.save("./test_roll.gif", savefig_kwargs = dict(transparent = False))

fig.show()