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.