Matplotlib绘制2D热力图：全面指南与实例|极客教程

Matplotlib绘制2D热力图：全面指南与实例

参考：How to draw 2D Heatmap using Matplotlib

Matplotlib是Python中强大的数据可视化库，其中绘制2D热力图是一项常用且重要的功能。热力图可以直观地展示二维数据的分布情况，广泛应用于数据分析、科学研究和商业决策等领域。本文将全面介绍如何使用Matplotlib绘制2D热力图，包括基础绘图、自定义颜色、添加标签、调整布局等多个方面，并提供丰富的示例代码，帮助读者掌握这一重要技能。

1. 热力图基础

热力图是一种用色彩来表示数值大小的图表。在2D热力图中，数据通常以矩阵的形式呈现，每个单元格的颜色深浅代表其对应的数值大小。Matplotlib提供了imshow()函数，这是绘制热力图的主要工具。

让我们从一个简单的例子开始：

import matplotlib.pyplot as plt
import numpy as np

# 创建一个简单的数据矩阵
data = np.random.rand(10, 10)

# 绘制热力图
plt.imshow(data)
plt.colorbar()
plt.title('How to draw 2D Heatmap using Matplotlib - how2matplotlib.com')
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们首先导入必要的库：Matplotlib的pyplot模块和NumPy。然后，我们创建一个10×10的随机数据矩阵。imshow()函数将这个矩阵可视化为热力图。colorbar()函数添加了一个颜色条，显示颜色与数值的对应关系。最后，我们添加了一个标题并显示图像。

2. 自定义颜色映射

Matplotlib提供了多种内置的颜色映射（colormap），可以根据数据的特性和个人喜好选择合适的配色方案。以下是一个使用不同颜色映射的例子：

import matplotlib.pyplot as plt
import numpy as np

data = np.random.rand(10, 10)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

im1 = ax1.imshow(data, cmap='viridis')
ax1.set_title('Viridis Colormap - how2matplotlib.com')
plt.colorbar(im1, ax=ax1)

im2 = ax2.imshow(data, cmap='hot')
ax2.set_title('Hot Colormap - how2matplotlib.com')
plt.colorbar(im2, ax=ax2)

plt.tight_layout()
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

这个例子展示了两种不同的颜色映射：’viridis’和’hot’。我们使用subplots()函数创建了两个并排的子图，分别使用不同的颜色映射绘制相同的数据。tight_layout()函数用于自动调整子图之间的间距。

3. 添加标签和刻度

为了使热力图更具信息量，我们通常需要添加轴标签和刻度。以下是一个添加自定义标签和刻度的例子：

import matplotlib.pyplot as plt
import numpy as np

data = np.random.rand(5, 5)
row_labels = ['A', 'B', 'C', 'D', 'E']
col_labels = ['V', 'W', 'X', 'Y', 'Z']

fig, ax = plt.subplots()
im = ax.imshow(data)

# 设置刻度和标签
ax.set_xticks(np.arange(len(col_labels)))
ax.set_yticks(np.arange(len(row_labels)))
ax.set_xticklabels(col_labels)
ax.set_yticklabels(row_labels)

# 旋转x轴标签
plt.setp(ax.get_xticklabels(), rotation=45, ha="right", rotation_mode="anchor")

# 添加颜色条和标题
cbar = ax.figure.colorbar(im)
ax.set_title("Custom Labels Heatmap - how2matplotlib.com")

plt.tight_layout()
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们为行和列分别定义了标签。使用set_xticks()和set_yticks()函数设置刻度位置，然后用set_xticklabels()和set_yticklabels()函数设置对应的标签。我们还通过plt.setp()函数旋转了x轴的标签，以防止它们重叠。

4. 在单元格中显示数值

有时，我们可能希望直接在热力图的每个单元格中显示具体的数值。以下是一个实现这一功能的例子：

import matplotlib.pyplot as plt
import numpy as np

data = np.random.randint(0, 100, size=(5, 5))

fig, ax = plt.subplots()
im = ax.imshow(data)

# 在每个单元格中显示数值
for i in range(5):
    for j in range(5):
        text = ax.text(j, i, data[i, j], ha="center", va="center", color="w")

ax.set_title("Heatmap with Values - how2matplotlib.com")
plt.colorbar(im)
plt.tight_layout()
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们使用嵌套的for循环遍历数据矩阵的每个元素，并使用ax.text()函数在对应的位置添加文本。文本的颜色设置为白色（”w”），以确保在深色背景上也能清晰可见。

5. 调整热力图的纵横比

默认情况下，imshow()函数会将热力图调整为正方形。但有时我们可能需要保持数据的原始纵横比。以下是一个保持原始纵横比的例子：

import matplotlib.pyplot as plt
import numpy as np

data = np.random.rand(8, 5)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

im1 = ax1.imshow(data)
ax1.set_title('Default Aspect Ratio - how2matplotlib.com')
plt.colorbar(im1, ax=ax1)

im2 = ax2.imshow(data, aspect='auto')
ax2.set_title('Auto Aspect Ratio - how2matplotlib.com')
plt.colorbar(im2, ax=ax2)

plt.tight_layout()
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们创建了两个子图。左边的子图使用默认设置，右边的子图通过设置aspect='auto'来保持数据的原始纵横比。这在处理非方形数据时特别有用。

6. 使用掩码数据

有时，我们可能需要在热力图中隐藏某些特定的数据点。这可以通过使用掩码数组来实现。以下是一个使用掩码数据的例子：

import matplotlib.pyplot as plt
import numpy as np
import numpy.ma as ma

data = np.random.rand(10, 10)
mask = np.zeros_like(data)
mask[2:8, 3:7] = True
masked_data = ma.masked_array(data, mask)

fig, ax = plt.subplots()
im = ax.imshow(masked_data)
ax.set_title('Masked Heatmap - how2matplotlib.com')
plt.colorbar(im)
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们首先创建了一个随机数据矩阵和一个掩码矩阵。掩码矩阵中的True值对应于我们想要隐藏的数据点。然后，我们使用ma.masked_array()函数创建一个掩码数组，并将其传递给imshow()函数。结果是一个部分数据被隐藏的热力图。

7. 使用对数刻度

当数据范围跨越多个数量级时，使用对数刻度可能会更有效地展示数据。以下是一个使用对数刻度的例子：

import matplotlib.pyplot as plt
import numpy as np

data = np.random.rand(10, 10) * 1000
data[0, 0] = 0.1  # 添加一个很小的值

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

im1 = ax1.imshow(data)
ax1.set_title('Linear Scale - how2matplotlib.com')
plt.colorbar(im1, ax=ax1)

im2 = ax2.imshow(data, norm=plt.LogNorm())
ax2.set_title('Log Scale - how2matplotlib.com')
plt.colorbar(im2, ax=ax2)

plt.tight_layout()
plt.show()

在这个例子中，我们创建了两个子图。左边的子图使用线性刻度，右边的子图通过设置norm=plt.LogNorm()来使用对数刻度。注意，我们在数据中添加了一个很小的值（0.1），因为对数刻度不能处理零或负值。

8. 自定义颜色范围

有时，我们可能需要手动设置颜色映射的范围，而不是使用数据的最小值和最大值。以下是一个自定义颜色范围的例子：

import matplotlib.pyplot as plt
import numpy as np

data = np.random.rand(10, 10)

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

im1 = ax1.imshow(data, vmin=0, vmax=1)
ax1.set_title('Full Range [0, 1] - how2matplotlib.com')
plt.colorbar(im1, ax=ax1)

im2 = ax2.imshow(data, vmin=0.2, vmax=0.8)
ax2.set_title('Custom Range [0.2, 0.8] - how2matplotlib.com')
plt.colorbar(im2, ax=ax2)

plt.tight_layout()
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，左边的子图使用数据的完整范围（0到1），而右边的子图通过设置vmin和vmax参数来限制颜色范围为0.2到0.8。这可以帮助突出显示数据中的特定范围。

9. 使用离散颜色映射

有时，我们可能希望使用离散的颜色而不是连续的颜色渐变。这可以通过使用BoundaryNorm来实现。以下是一个使用离散颜色映射的例子：

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import BoundaryNorm
from matplotlib.ticker import MaxNLocator

data = np.random.randint(0, 100, size=(10, 10))

# 定义颜色边界和规范化
levels = MaxNLocator(nbins=5).tick_values(data.min(), data.max())
cmap = plt.get_cmap('viridis')
norm = BoundaryNorm(levels, ncolors=cmap.N, clip=True)

fig, ax = plt.subplots()
im = ax.imshow(data, cmap=cmap, norm=norm)
ax.set_title('Discrete Colormap - how2matplotlib.com')
plt.colorbar(im, ax=ax, ticks=levels)
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们首先使用MaxNLocator定义了颜色边界，然后创建了一个BoundaryNorm对象。这个对象被传递给imshow()函数，结果是一个使用离散颜色的热力图。

10. 添加网格线

为了更清晰地区分热力图中的单元格，我们可以添加网格线。以下是一个添加网格线的例子：

import matplotlib.pyplot as plt
import numpy as np

data = np.random.rand(5, 5)

fig, ax = plt.subplots()
im = ax.imshow(data)

# 添加网格线
ax.set_xticks(np.arange(data.shape[1]+1)-.5, minor=True)
ax.set_yticks(np.arange(data.shape[0]+1)-.5, minor=True)
ax.grid(which="minor", color="w", linestyle='-', linewidth=3)
ax.tick_params(which="minor", bottom=False, left=False)

ax.set_title('Heatmap with Grid - how2matplotlib.com')
plt.colorbar(im)
plt.tight_layout()
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们使用set_xticks()和set_yticks()函数设置了次要刻度的位置，然后使用grid()函数添加网格线。我们将网格线的颜色设置为白色，以便在热力图上清晰可见。

11. 使用自定义颜色映射

除了使用Matplotlib提供的内置颜色映射，我们还可以创建自定义的颜色映射。以下是一个创建和使用自定义颜色映射的例子：

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import LinearSegmentedColormap

data = np.random.rand(10, 10)

# 定义自定义颜色映射
colors = ['blue', 'white', 'red']
n_bins = 100
cmap = LinearSegmentedColormap.from_list('custom_cmap', colors, N=n_bins)

fig, ax = plt.subplots()
im = ax.imshow(data, cmap=cmap)
ax.set_title('Custom Colormap - how2matplotlib.com')
plt.colorbar(im)
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们使用LinearSegmentedColormap.from_list()函数创建了一个从蓝色到白色再到红色的自定义颜色映射。这种自定义颜色映射可以帮助我们更好地突出数据中的特定特征。

12. 添加轮廓线

在某些情况下，添加轮廓线可以帮助我们更好地理解热力图中的数据分布。以下是一个在热力图上添加轮廓线的例子：

import matplotlib.pyplot as plt
import numpy as np

data = np.random.rand(20, 20)

fig, ax = plt.subplots()
im = ax.imshow(data)
cs = ax.contour(data, colors='black', linewidths=1)
ax.clabel(cs, inline=True, fontsize=10)

ax.set_title('Heatmap with Contour Lines - how2matplotlib.com')
plt.colorbar(im)
plt.tight_layout()
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们首先使用imshow()函数绘制热力图，然后使用contour()函数添加轮廓线。clabel()函数用于在轮廓线上添加标签。这种方法可以同时展示数据的整体分布和局部变化。

13. 使用极坐标系

虽然大多数热力图都使用笛卡尔坐标系，但在某些情况下，使用极坐标系可能更合适。以下是一个在极坐标系中绘制热力图的例子：

import matplotlib.pyplot as plt
import numpy as np

r = np.linspace(0, 2, 100)
theta = np.linspace(0, 2*np.pi, 100)
r, theta = np.meshgrid(r, theta)

data = np.sin(5*theta) * (r**2)

fig, ax = plt.subplots(subplot_kw=dict(projection='polar'))
im = ax.pcolormesh(theta, r, data)
ax.set_title('Polar Heatmap - how2matplotlib.com')
plt.colorbar(im)
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们使用pcolormesh()函数在极坐标系中绘制热力图。这种方法特别适合展示具有周期性或径向分布特征的数据。

14. 使用三角形网格

有时，我们的数据可能不是规则的矩形网格，而是不规则的三角形网格。Matplotlib也提供了处理这种情况的工具。以下是一个使用三角形网格绘制热力图的例子：

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.tri import Triangulation

# 创建一些随机数据点
n_angles = 36
n_radii = 8
min_radius = 0.25
radii = np.linspace(min_radius, 0.95, n_radii)

angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False)
angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1)
angles[:, 1::2] += np.pi/n_angles

x = (radii*np.cos(angles)).flatten()
y = (radii*np.sin(angles)).flatten()
z = (np.cos(radii)*np.cos(3*angles)).flatten()

# 创建 Triangulation
triang = Triangulation(x, y)

# 绘制热力图
fig, ax = plt.subplots()
im = ax.tricontourf(triang, z)
ax.set_title('Triangular Mesh Heatmap - how2matplotlib.com')
plt.colorbar(im)
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们首先创建了一些随机的数据点，然后使用Triangulation类创建了一个三角形网格。最后，我们使用tricontourf()函数绘制热力图。这种方法可以处理不规则分布的数据点。

15. 添加注释

有时，我们可能需要在热力图上添加注释来突出显示某些特定的区域或值。以下是一个在热力图上添加注释的例子：

import matplotlib.pyplot as plt
import numpy as np

data = np.random.rand(10, 10)

fig, ax = plt.subplots()
im = ax.imshow(data)

# 添加注释
ax.annotate('Interesting\nRegion', xy=(2, 3), xytext=(3, 8),
            arrowprops=dict(facecolor='black', shrink=0.05))

ax.set_title('Heatmap with Annotation - how2matplotlib.com')
plt.colorbar(im)
plt.tight_layout()
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们使用annotate()函数添加了一个带箭头的注释。这个注释指向数据矩阵中的特定位置，并在图上的另一个位置显示文本。这种方法可以帮助我们强调热力图中的重要特征。

16. 使用多个子图

在某些情况下，我们可能需要在同一个图形中比较多个热力图。以下是一个创建多个子图的例子：

import matplotlib.pyplot as plt
import numpy as np

data1 = np.random.rand(10, 10)
data2 = np.random.rand(10, 10)
data3 = np.random.rand(10, 10)
data4 = np.random.rand(10, 10)

fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(12, 10))

im1 = ax1.imshow(data1)
ax1.set_title('Heatmap 1 - how2matplotlib.com')
plt.colorbar(im1, ax=ax1)

im2 = ax2.imshow(data2)
ax2.set_title('Heatmap 2 - how2matplotlib.com')
plt.colorbar(im2, ax=ax2)

im3 = ax3.imshow(data3)
ax3.set_title('Heatmap 3 - how2matplotlib.com')
plt.colorbar(im3, ax=ax3)

im4 = ax4.imshow(data4)
ax4.set_title('Heatmap 4 - how2matplotlib.com')
plt.colorbar(im4, ax=ax4)

plt.tight_layout()
plt.show()

Output:

Matplotlib绘制2D热力图：全面指南与实例

在这个例子中，我们创建了一个2×2的子图网格，每个子图都包含一个热力图。这种方法允许我们在同一个图形中比较多个相关的数据集。

17. 使用交互式热力图

Matplotlib还支持创建交互式的热力图，这在探索大型数据集时特别有用。以下是一个使用pcolormesh()函数创建可缩放的交互式热力图的例子：

import matplotlib.pyplot as plt
import numpy as np

data = np.random.rand(1000, 1000)

fig, ax = plt.subplots()
im = ax.pcolormesh(data)
ax.set_title('Interactive Heatmap - how2matplotlib.com')
plt.colorbar(im)

def on_xlims_change(event_ax):
    xlim = event_ax.get_xlim()
    ylim = event_ax.get_ylim()
    im.set_data(data[int(ylim[0]):int(ylim[1]), int(xlim[0]):int(xlim[1])])
    im.autoscale()

ax.callbacks.connect('xlim_changed', on_xlims_change)
ax.callbacks.connect('ylim_changed', on_xlims_change)

plt.show()

在这个例子中，我们定义了一个回调函数on_xlims_change，它在图表的x轴或y轴限制发生变化时被调用。这个函数更新热力图的数据，只显示当前视图中的数据。这样，用户可以通过缩放和平移来探索大型数据集的不同部分。

18. 使用动画热力图

对于随时间变化的数据，我们可以创建动画热力图。以下是一个创建简单动画热力图的例子：

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation

fig, ax = plt.subplots()
data = np.random.rand(10, 10)
im = ax.imshow(data, animated=True)
ax.set_title('Animated Heatmap - how2matplotlib.com')

def update(frame):
    data = np.random.rand(10, 10)
    im.set_array(data)
    return [im]

ani = FuncAnimation(fig, update, frames=200, interval=50, blit=True)
plt.colorbar(im)
plt.show()