animating 2d heat map

[schniefen]

Consider the following heat map:

from scipy import special
import numpy as np
import matplotlib.pyplot as plt

u0=200
r0x=25
r0y=25
rmax=2.5
alpha=2
t=0.575

y, x = np.meshgrid(np.linspace(0, 50, 100), np.linspace(0, 50, 100))
r=np.sqrt((x-r0x)**2+(y-r0y)**2)
z=u0*(special.erf((rmax-r)/(4*alpha*t))-special.erf((-rmax-r)/(4*alpha*t)))

plt.pcolormesh(x, y, z, cmap='jet')
plt.xlim(x.min(), x.max())
plt.ylim(y.min(), y.max())
plt.xlabel('x')
plt.ylabel('y')
plt.title(f'Temperature at {t} unit time')
cbar=plt.colorbar()
cbar.set_label('Temperature')
plt.clim(0,200)

   

I would like to animate this heat map for different t, say t=np.linspace(0.0001,3). My attempt so far is the following:

import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation

u0=200
r0x=25
r0y=25
rmax=2.5
alpha=2

def pulse_maker(x,y,t):
    x,y,t =np.meshgrid(x,y,t)
    r=np.sqrt((x-r0x)**2+(y-r0y)**2)
    return u0*(special.erf((rmax-r)/np.sqrt(4*alpha*t))-special.erf((-rmax-r)/np.sqrt(4*alpha*t)))

def plot_pulse(xvals, yvals, pulse, t, i):
    plt.clf()
    plt.title(f"Temperature at t = {t[i]:.2f}")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.xlim(xvals.min(), xvals.max())
    plt.ylim(yvals.min(), yvals.max())
    plt.pcolormesh(xvals, yvals, pulse[i], cmap='jet') 
    cbar=plt.colorbar()
    cbar.set_label('Temperature')
    plt.clim(0,200)
    return plt

def animate(i):
    plt = plot_pulse(xrange,yrange, pulse, times, i);

max_iterations = 200
times = np.linspace(0.0001, 3, max_iterations)
xrange = np.linspace(0,50, max_iterations)
yrange = np.linspace(0,50, max_iterations)
pulse = pulse_maker(x=xrange, y=yrange, t=times)
anim = animation.FuncAnimation(plt.figure(), animate, interval=1, 
                               frames=max_iterations, repeat=False)
anim.save("heat.gif",writer='pillow')

However, there is something wrong with this animation if you run it. It does not give back the still image as above, rather the heat is very off-centered, although I have used the same initial values. Grateful for any feedback.