The solution U contains the color values and array shape information, but don't explicitly define the bounded x/y-values as you've pointed out. To do this with pcolormesh, we just need to change the values according to the axes() class, as so (using a random dataset):
import numpy as np
import matplotlib.pyplot as plt
U = np.random.rand(20,20)
print(U.shape)
plt.figure()
ax=plt.axes()
plt.title(f"Temperature at t = 100")
plt.xlabel("x")
plt.ylabel("y")
plt.pcolormesh(U, cmap=plt.cm.jet)
x=y=np.linspace(0,10,11) # Define a min and max with regular spacing
ax.set_xticks(np.linspace(0,20,11)) # Ensure same number of x,y ticks
ax.set_yticks(np.linspace(0,20,11))
ax.set_xticklabels(x.astype(int)) # Change values according to the grid
ax.set_yticklabels(y.astype(int))
plt.colorbar()
plt.show()
Alternatively, we can explicitly define these using the extent=[<xmin>,<xmax>,<ymin>,<ymax>] option in imshow, as seen below:
import numpy as np
import matplotlib.pyplot as plt
U = np.random.rand(20,20)
print(U.shape) # Displays (20,20)
plt.figure()
plt.title(f"Temperature at t = 100")
plt.xlabel("x")
plt.ylabel("y")
plt.imshow(U, origin='lower', aspect='auto', extent = [0,10,0,10], cmap=plt.cm.jet)
plt.colorbar()
plt.show()
For further reading on which to use, pcolormesh or imshow, there is a good thread about it here: When to use imshow over pcolormesh?
animation.FuncAnimation(plt.figure(), animate, frames=range(0, max_iter)), where animate is a function containing the code in my question.pcolormeshfor this? UsingimshowI can do exactly what you are asking.