Matplotlib FuncAnimation runs in VS Code but not in Google Colab

I have a diffusion code that runs flawlessly in VS Code using Matplotlib's FuncAnimation function. When I copy the code over to a Jupyter notebook (my boss for my internship wants all my code to be in Google Colab), it compiles but it doesn't produce an animated plot. Here is the code:

%matplotlib notebook
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation
from numpy import copy,empty,zeros
class diffuse():
def __init__(self,N,Tex,Tin,D,circ,squr):
self.N = N       #grid side length
self.Tex = Tex
self.Tin = Tin
self.D = D
self.dq = 0.1
self.dt = self.dq**4/4/D/self.dq
if circ == True:
self.grido = empty((N,N)) #initialize grid
m = N//2
r = N//4
e = 0.01
#loop through every position
for i in range(N):
for j in range(N):
#check to see if inside of desired circle
#got help from https://stackoverflow.com/questions/22777049
#/how-can-i-draw-a-circle-in-a-data-array-map-in-python
R = (i-m)**2 + (j-m)**2 - r**2
if R < e:
self.grido[i,j] = Tin
else:
self.grido[i,j] = Tex
elif squr == True:
self.grido = empty((N,N))
l = N//4
u = 3*N//4
for i in range(N):
for j in range(N):
#check to see if inside of desired square
if i > l and i < u and j > l and j < u:
self.grido[i,j] = self.Tin
else:
self.grido[i,j] = self.Tex
else:
print("error")
def timestep(self):
self.grid = zeros((self.N,self.N))
self.grid[1:-1,1:-1] = self.grido[1:-1,1:-1] + \
self.D*self.dt/self.dq**2*(
self.grido[2:,1:-1] + self.grido[:-2,1:-1] + 
self.grido[1:-1,2:] + self.grido[1:-1,:-2] -
4*self.grido[1:-1,1:-1])
self.grido = copy(self.grid)
def animate(self,i):
self.timestep()
self.ax.set_title(i)#f"{i*self.dt:.001} ms")
self.im.set_data(copy(self.grido))
def cool(self):
n = 5000
fig = plt.figure()
self.ax = fig.add_subplot()
self.im = self.ax.imshow(self.grido)#,vmin=self.Tex,vmax=self.Tin)
self.ax.set_axis_off()
self.ax.set_title(-1)#"0.0 ms")
#fig.subplots_adjust(right=0.85)
anim = FuncAnimation(fig,self.animate,
frames=n,
repeat=False,
interval=1)
plt.show()
N = 101
Tex = 293
Tin = 400
D = 4.25e-6      #thermal diffusivity constant
muffin = diffuse(N,Tex,Tin,D,False,True)
muffin.cool()

I know that when you animate in Jupyter notebooks, you need to use a magic command, namely %matplotlib notebook in order to get an interactive plot. However in Google Colab I cannot get any of the magic commands to do what I want them to do. I am unsure if I need to install something in order to get these commands to work. When using the same code in Anaconda's environment for Jupyter Notebooks, the magic commands work, but the plot still doesn't animate. I have had problems in the past where classes don't quite work as intended, so I'm not sure, there could be a problem there. Any help is greatly appreciated! Thank you!

I'm seeing animation.FuncAnimation stuff being clunky in current sessions launched from my 'animated-matplotlib' testing grounds here (Direct launch URL: https://mybinder.org/v2/gh/fomightez/animated_matplotlib-binder/master?filepath=index.ipynb ), but I am seeing some things work in actual Jupyter Notebook with small additions to your code. (technically Google Colab is a vastly customized very old now version of Jupyter Notebook and so to reference them as one-in-the-same, like you do in your post ("When I copy the code over to a Jupyter notebook (my boss for my internship wants all my code to be in Google Colab)") is not fully accurate and will cause you issues finding help and problems if you expect Jupyter Notebook-related stuff to work in Google Colab. And hence your title would probably be better specified, too. Something like 'Matplotlib FuncAnimation runs in VS Code but not in GoogleColab'). If your internship isn't with Google, you may want to encourage your boss to be working in and supporting an open source system that won't just change radically overnight out of your control or disappear one day, like the many other items in the Google Graveyard.
That being said these address the current title of your post, 'Matplotlib FuncAnimation runs in VS Code but not in a Jupyter Notebook', offering versions I can demonstrate work (at least somewhat) in Jupyter Notebook. I cannot say whether that extends to Google Colab that you only bring up parenthetically in the early part of your post.

One option to get something:

%matplotlib notebook
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation
from numpy import copy,empty,zeros
from IPython.display import HTML, display

class diffuse():
  def __init__(self,N,Tex,Tin,D,circ,squr):
    self.N = N       #grid side length
    self.Tex = Tex
    self.Tin = Tin
    self.D = D
    self.dq = 0.1
    self.dt = self.dq**4/4/D/self.dq
    
    if circ == True:
      self.grido = empty((N,N)) #initialize grid
      m = N//2
      r = N//4
      e = 0.01
      #loop through every position
      for i in range(N):
        for j in range(N):
        #check to see if inside of desired circle
        #got help from https://stackoverflow.com/questions/22777049
        #/how-can-i-draw-a-circle-in-a-data-array-map-in-python
          R = (i-m)**2 + (j-m)**2 - r**2
          if R < e:
            self.grido[i,j] = Tin
          else:
            self.grido[i,j] = Tex

    elif squr == True:
      self.grido = empty((N,N))
      l = N//4
      u = 3*N//4
      for i in range(N):
        for j in range(N):
          #check to see if inside of desired square
          if i > l and i < u and j > l and j < u:
            self.grido[i,j] = self.Tin
          else:
            self.grido[i,j] = self.Tex

    else:
      print("error")

  def timestep(self):
    self.grid = zeros((self.N,self.N))
    self.grid[1:-1,1:-1] = self.grido[1:-1,1:-1] + \
                            self.D*self.dt/self.dq**2*(
                            self.grido[2:,1:-1] + self.grido[:-2,1:-1] + 
                            self.grido[1:-1,2:] + self.grido[1:-1,:-2] -
                            4*self.grido[1:-1,1:-1])

    self.grido = copy(self.grid)

  def animate(self,i):
    self.timestep()
    self.ax.set_title(i)#f"{i*self.dt:.001} ms")
    self.im.set_data(copy(self.grido))

  def cool(self):
    n = 15
    fig = plt.figure()
    self.ax = fig.add_subplot()
    self.im = self.ax.imshow(self.grido)#,vmin=self.Tex,vmax=self.Tin)
    self.ax.set_axis_off()
    self.ax.set_title(-1)#"0.0 ms")
    #fig.subplots_adjust(right=0.85)
    anim = FuncAnimation(fig,self.animate,
                              frames=n,
                              repeat=False,
                              interval=1)
    display(HTML(anim.to_jshtml()))
    #plt.close(fig)
N = 101
Tex = 293
Tin = 400
D = 4.25e-6      #thermal diffusivity constant
muffin = diffuse(N,Tex,Tin,D,False,True)
muffin.cool();

It works but shows a non-interactive final frame in the 'interactive' window above the display where the animation works using the controller.
You can test that in launches from where I reference up above (my 'animated-matplotlib' testing grounds) and use the widget controller to run the animation.
(Interestingly that doesn't happen in the example I adapted and so maybe you can remove a double invoking somewhere in your code that I'm not seeing on quick glance. But the code that did work fine before from there is currently being a bit clunky and so something was changed somewhere. I'm at a loss as to what is causing it right now and usually these things sort themselves out as more consistent versions come along. Another aside: the notebook with the code I reference in 'static' form there is linked to from inside the manin notebook that comes up from launches from here. Note I use 'static' in quotes because if you try the widget controller you'll see the animations still work in nbviewer where you are viewing them because of the underlying frames & javascript saved in them, whereas they don't work if you preview that notebook directly on GitHub where nbviewer is rendering it from because presently GitHub limits 'static' notebooks running javascript.)

I mildly suspect part of the problem with your code is burying the anim = line inside a function inside a class. But maybe not? Things are generally clunky that weren't before, and so it makes it hard to tell what is 'off'. I don't have an actual example to compare for that and I don't generally use object-oriented coding for plot endeavors. You'll not the example at the bottom there returns a plot object to the animation function that is in the main block of the code. Maybe you could try that in yours and see if it improves handling in Colab?

Another option

This variation also 'works' imperfectly in actual Jupyter Notebook presently:

%matplotlib notebook
from matplotlib import pyplot as plt
from matplotlib.animation import FuncAnimation
from numpy import copy,empty,zeros
from IPython.display import HTML, display

class diffuse():
  def __init__(self,N,Tex,Tin,D,circ,squr):
    self.N = N       #grid side length
    self.Tex = Tex
    self.Tin = Tin
    self.D = D
    self.dq = 0.1
    self.dt = self.dq**4/4/D/self.dq
    
    if circ == True:
      self.grido = empty((N,N)) #initialize grid
      m = N//2
      r = N//4
      e = 0.01
      #loop through every position
      for i in range(N):
        for j in range(N):
        #check to see if inside of desired circle
        #got help from https://stackoverflow.com/questions/22777049
        #/how-can-i-draw-a-circle-in-a-data-array-map-in-python
          R = (i-m)**2 + (j-m)**2 - r**2
          if R < e:
            self.grido[i,j] = Tin
          else:
            self.grido[i,j] = Tex

    elif squr == True:
      self.grido = empty((N,N))
      l = N//4
      u = 3*N//4
      for i in range(N):
        for j in range(N):
          #check to see if inside of desired square
          if i > l and i < u and j > l and j < u:
            self.grido[i,j] = self.Tin
          else:
            self.grido[i,j] = self.Tex

    else:
      print("error")

  def timestep(self):
    self.grid = zeros((self.N,self.N))
    self.grid[1:-1,1:-1] = self.grido[1:-1,1:-1] + \
                            self.D*self.dt/self.dq**2*(
                            self.grido[2:,1:-1] + self.grido[:-2,1:-1] + 
                            self.grido[1:-1,2:] + self.grido[1:-1,:-2] -
                            4*self.grido[1:-1,1:-1])

    self.grido = copy(self.grid)

  def animate(self,i):
    self.timestep()
    self.ax.set_title(i)#f"{i*self.dt:.001} ms")
    self.im.set_data(copy(self.grido))

  def cool(self):
    n = 15
    fig = plt.figure()
    self.ax = fig.add_subplot()
    self.im = self.ax.imshow(self.grido)#,vmin=self.Tex,vmax=self.Tin)
    self.ax.set_axis_off()
    self.ax.set_title(-1)#"0.0 ms")
    #fig.subplots_adjust(right=0.85)
    anim = FuncAnimation(fig,self.animate,
                              frames=n,
                              repeat=False,
                              interval=1)
    display(HTML(anim.to_jshtml()))
    plt.close(fig)
N = 101
Tex = 293
Tin = 400
D = 4.25e-6      #thermal diffusivity constant
muffin = diffuse(N,Tex,Tin,D,False,True)
muffin.cool();

When I run that I see the animation get made and have a widget controller to allow me to control it. However, it also throws an error. That looks sort of bad but doesn't seem to interfere with the animation actually working for now.