Making animations work in Google Colaboratory — new home for ML prototyping

I always felt Jupyter was such an amazing platform to get up and running quickly in Machine Learning — with focus on quick results/content, minimal installation drama.
Well, that was until I found Google Colab.
Change minimal installation to zero installation drama, add some bells and whistles to the UX, add Google Docs like real-time collaboration features, and add a 1-click option to switch to GPU when needed —Google Colab is definitely a 1-up on Jupyter notebooks.
That said, as with anything new and shiny, there are gaps, and understandably so.
1 big blocker which hits soon enough is — animating your plots. The matplotlib
animation APIs which work smoothly out of the box with Jupyter notebooks, are not that trivial to get working, on Colab.
Here’s the best way I have seen to get your plots animating
from matplotlib import rc
rc('animation', html='jshtml')
Let’s go through step by step details
Task: plot y= sin(x)
and animate a tangent line showing its gradient at each point
Let’s start with the imports
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation
get a range of x
to plot our function sin(x)
x = np.linspace(-4, 4, 100)
y = np.sin(x)
create the required plot objects
fig, ax = plt.subplots()
ax.set_xlim(( -4, 4))
ax.set_ylim((-2, 2))
create 2 line plots, 1 each for our function and its gradient line
line1, = ax.plot([], [], lw=2)
line2, = ax.plot([], [], lw=2)
next define the init()
function, which is responsible for the background of each frame in the animation
def init():
line1.set_data(x, y)
return (line1,)
now define the animate()
function, which is called sequentially and is responsible for the change we see in each frame
def animate(i):
at_x = x[i]
# gradient_line will have the form m*x + b
m = np.cos(at_x)
b = np.sin(at_x) - np.cos(at_x)*at_x
gradient_line = m*x + b
line2.set_data(x, gradient_line)
return (line2,)
finally, call the FuncAnimation()
of matplotlib.animation
to kick off the animation
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=100, interval=100, blit=True)
At this stage, we have our animation object anim
— and to animate it inline, we want to set the default HTML representation of our plots as interactive javascript widget.
This is done by setting the rc
parameter ofanimation.html
to jshtml
rc('animation', html='jshtml')
Once the above is set, we can just reference our animation object anim
to display the animated plot inline.
Here’s what the combined code block looks like:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation, rc
from IPython.display import HTML
# animate over some set of x, y
x = np.linspace(-4, 4, 100)
y = np.sin(x)
# First set up the figure, the axes, and the plot element
fig, ax = plt.subplots()
plt.close()
ax.set_xlim(( -4, 4))
ax.set_ylim((-2, 2))
line1, = ax.plot([], [], lw=2)
line2, = ax.plot([], [], lw=2)
# initialization function: plot the background of each frame
def init():
line1.set_data(x, y)
return (line1,)
# animation function: this is called sequentially
def animate(i):
at_x = x[i]
# gradient_line will have the form m*x + b
m = np.cos(at_x)
b = np.sin(at_x) - np.cos(at_x)*at_x
gradient_line = m*x + b
line2.set_data(x, gradient_line)
return (line2,)
anim = animation.FuncAnimation(fig, animate, init_func=init, frames=100, interval=100, blit=True)
rc('animation', html='jshtml')
anim
Running the above in Google Colab will give the inline animated plot we were after.
