Why are my subplots not positioned properly after animating them?

I'm answering module on image processing and I wanted to show how the values of a given 2D image change based on the intensity of the values of each pixel. The only thing is that I can't seem to animate them properly together, animating them individually works fine but I want to show the simultaneous changes side-by-side. I followed this guide on animating 3D plots using plot_surface (which I already used before I got the idea of animating the whole thing): https://pythonmatplotlibtips.blogspot.com/2018/11/animation-3d-surface-plot-funcanimation-matplotlib.html

Then I used the same idea for the 2D image. The problem is that the 2D image seems to be displaced someplace else but I can't seem to find a workaround. I tried including fig.add_subplot() to the animation function used in updating the data with FuncAnimation, but the result is the same.

Here's the code I used:

N = 256

inner_radius_ratio = 0.4
outer_radius_ratio = 0.8

xx = np.linspace(-1, 1, N)
X, Y = np.meshgrid(xx, xx) 
R = np.sqrt(X**2 + Y**2)

fig = plt.figure()
canvas = np.zeros(np.shape(R))                                            
canvas[np.where((R>inner_radius_ratio) & (R<=outer_radius_ratio))] = 0
zarray = np.zeros((N, N, N))

for i in range(N):
    canvas[np.where((R>inner_radius_ratio) & (R<=outer_radius_ratio))] = i/256
    zarray[:,:,i] = canvas

def update_plot(frame_number,zarray, im1, im2):
    im1[0].remove()
    im1[0] = ax.imshow(zarray[:,:,frame_number],cmap='magma', vmin=0, vmax=1, animated=True)
    im2[0].remove()  
    im2[0] = ax.plot_surface(X, Y, zarray[:,:,frame_number], cmap='magma')

#Plot 2D Image
ax = fig.add_subplot(121)
im1 = [ax.imshow(zarray[:,:,0],cmap='magma', vmin=0, vmax=1, animated=True)]
ax.grid(False)

#Plot 3D Surface
ax = fig.add_subplot(122, projection='3d')
im2 = [ax.plot_surface(X, Y , zarray[:,:,0], cmap='magma')]#,rstride=1, cstride=1)
ax.set_zlim(0,1.0)
ax.grid(True)

plt.tight_layout()

ani = animation.FuncAnimation(fig, update_plot, fargs=(zarray, im1, im2), frames=N, repeat=True, interval=10)
plt.close()
HTML(ani.to_html5_video())

What ends up showing is the 3D animation working fine, but the 2D plot not showing at all. My guess is that it's behind the 3D plot since there seems to be some text behind it, but I'm not completely sure since it's position was already set up with fig.add_subplot(122). Also, the gif looks slow but its much faster in from the video (i just converted it to gif online): animation

You use the same variable ax for both plots - but you should use ax1, ax2.

That's all.

def update_plot(frame_number,zarray, im1, im2):
    im1[0].remove()
    im1[0] = ax1.imshow(zarray[:,:,frame_number],cmap='magma', vmin=0, vmax=1, animated=True)
    im2[0].remove()  
    im2[0] = ax2.plot_surface(X, Y, zarray[:,:,frame_number], cmap='magma')

#Plot 2D Image
ax1 = fig.add_subplot(121)
im1 = [ax1.imshow(zarray[:,:,0],cmap='magma', vmin=0, vmax=1, animated=True)]
ax1.grid(False)

#Plot 3D Surface
ax2 = fig.add_subplot(122, projection='3d')
im2 = [ax2.plot_surface(X, Y , zarray[:,:,0], cmap='magma')]#,rstride=1, cstride=1)
ax2.set_zlim(0,1.0)
ax2.grid(True)