Make heatmap have the same axes size

Question

This is a sample code. Here, xlim = (0,2) and ylim = (0,5). I want to keep it that way (different limits) while figuring out a way to make the x-axis and y-axis the same size.

import matplotlib.pyplot as plt
import numpy as np

# Create a heatmap dataset.
x = np.linspace(0, 2, 100)
y = np.linspace(0, 5, 100)
Z = np.random.rand(100, 100)

# Create a figure and axes object.
fig, ax = plt.subplots()

# Set the x-limits to 2
ax.set_xlim([0, 2])

# Plot the heatmap.
im = ax.imshow(Z, extent=[0, 2, 0, 5])

# Set the labels and title.
ax.set_xlabel('X-Axis')
ax.set_ylabel('Y-Axis')
ax.set_title('Heatmap with Custom X-Limits')

# Show the plot.
plt.show()

Result:

I tried using ax.set_box_aspect(1). This makes the figure a square with equal axes, however, sets the heatmap in the middle while leaving blank spaces on either side.

Answer 1

Since your data is not an image, you should use pcolormesh, which allows you to set the x- and y-axes. Then, ax.set_box_aspect(1) works fine.

import matplotlib.pyplot as plt
import numpy as np

x = np.linspace(0, 2, 100)
y = np.linspace(0, 5, 100)
Z = np.random.rand(100, 100)

fig, ax = plt.subplots()

p = ax.pcolormesh(x, y, Z)
ax.set_box_aspect(1)
fig.show()