How can I plot a 2D heatmap from a NumPy array using Matplotlib?

I’m trying to visualize a 2D NumPy array as a heatmap using Matplotlib. My array is n × n, and each element has a value between 0 and 1. I want each element (i, j) to be represented as a square in the heatmap, with its color corresponding to the value in the array.

In other words, I want a simple matplotlib heatmap where higher values are displayed with more intense colors and lower values with lighter colors.

What’s the best approach to create this kind of heatmap? Are there built-in functions in Matplotlib for this, or do I need to manually draw each square?

Example of the data:

import numpy as np

data = np.random.rand(5, 5)  # 5x5 array with values between 0 and 1

How can I turn this into a visually clear heatmap?

Hey! The easiest is imshow():

import matplotlib.pyplot as plt
import numpy as np

data = np.random.rand(5, 5)
plt.imshow(data, cmap='viridis', interpolation='nearest')
plt.colorbar()
plt.show()

cmap controls the color scheme, and colorbar() adds a reference scale.

I like pcolormesh when I want more control over each cell’s boundaries:

plt.pcolormesh(data, cmap='coolwarm', edgecolors='k', linewidth=0.5)
plt.colorbar()
plt.show()

It makes the grid visually clear.

Seaborn is also very convenient for heatmaps:

import seaborn as sns sns.heatmap(data, annot=True, cmap=“YlGnBu”) plt.show()

You get annotations inside cells and nicer default styling with less code.