How to MatPlotLib plot and then add different axes?

Question

I want to plot the solution of a PDE from (0, 0) to (10, 10). The solution is given in a 20 by 20 matrix.

Here is my code:

plt.figure()
plt.title(f"Temperature at t = 100")
plt.xlabel("x")
plt.ylabel("y")

plt.pcolormesh(U[-1], cmap=plt.cm.jet)
plt.colorbar()

So I would like the same plot, but the axis should be from 0 to 10. Can I add a second axis that goes from 0 to 10 and then hide the current axis? Is it possible to achieve this without plt.subplots() because I would like to animate this figure (animation.FuncAnimation(plt.figure(), animate, frames=range(0, max_iter)), where animate is a function containing the code above)?

Answer 1

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?