Trying to get the cmap on a surface in 3D to match the cmap of its sublevel sets

I am trying to get the colors of the pictured three dimensional surface to match the colors of the corresponding sublevel sets on the right (image below). In particular I want the small values in the middle of the image on the right to be blue and the large values to be red. Here is my code:

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d

# Define the function to plot
def f(x, y):
    return x**2 + y**2

# Generate data for the x, y, and z coordinates
x = np.linspace(-6, 6, 100)
y = np.linspace(-6, 6, 100)
X, Y = np.meshgrid(x, y)
Z = f(X, Y)

# Create a 3D figure and a contour plot side by side
fig = plt.figure(figsize=(10, 4))
ax1 = fig.add_subplot(121, projection='3d')
ax2 = fig.add_subplot(122)

# Plot the surface on the left subplot
surf = ax1.plot_surface(X, Y, Z, cmap='jet')

# Plot the contour on the right subplot
for a in range(1,30,5):
    contour_levels = np.arange(-a, a, 0.1)
    ax2.contourf(X, Y, Z, levels=contour_levels, cmap='jet')
    plt.pause(.001)

# Add labels and title to both subplots
ax1.set_xlabel('x')
ax1.set_ylabel('y')
ax1.set_zlabel('z')
ax1.set_title('3D Plot of f(x, y) = sin(sqrt(x^2 + y^2))')
ax2.set_xlabel('x')
ax2.set_ylabel('y')
ax2.set_title('Level Set Contour Plot of f(x, y) = x^2 + y^2')
surf.set_edgecolors(surf.to_rgba(surf._A))

# Show the plot
plt.show()

And here is the output:

Anyone spot any bugs or see how I can fix this? I have a suspicion that the contourf function may not actually be giving me the sublevel sets that I want. Just for completeness sake, by a sublevel set I mean:

${xinmathbb{R}: f(x)le a}$

Also, PS will this site really not accept Latex?