在Matplotlib中叠加contourf和surface_plot

什么是叠加contourf和surface_plot？

在Matplotlib中，我们常常需要绘制各种二维和三维图形，而contourf和surface_plot是两个常用的绘图函数。contourf函数用于绘制等高线填充图，而surface_plot函数用于绘制三维曲面图。在某些情况下，我们需要叠加这两种图形来更好地表达数据。例如，在地形图中，我们可以用surface_plot绘制地形表面，同时用contourf来表示高程信息。本文将介绍如何在Matplotlib中叠加contourf和surface_plot。

如何叠加contourf和surface_plot？

我们可以通过在一个绘图坐标系中依次调用contourf和surface_plot函数来实现叠加。例如，考虑以下示例数据：

import numpy as np

X, Y = np.meshgrid(np.linspace(-5, 5, 50), np.linspace(-5, 5, 50))
Z = np.sin(np.sqrt(X**2 + Y**2))

这是一个在[-5, 5] x [-5, 5]范围内以原点为中心的二维正弦函数。

我们可以先以surface_plot方式绘制曲面：

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X, Y, Z, cmap=plt.cm.coolwarm, linewidth=0, antialiased=False)

注意这里传入的参数包括X、Y、Z表示数据，cmap表示使用的颜色映射，linewidth表示曲面线宽，antialiased表示是否抗锯齿。

然后以contourf方式绘制等高线：

cset = ax.contourf(X, Y, Z, zdir='z', offset=-2, cmap=plt.cm.coolwarm)

这里zdir=’z’表示沿着z方向投影，offset=-2表示等高线图离z轴的高度。

至此，我们就完成了两种图形的叠加。完整的代码如下：

import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt

X, Y = np.meshgrid(np.linspace(-5, 5, 50), np.linspace(-5, 5, 50))
Z = np.sin(np.sqrt(X**2 + Y**2))

fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(X, Y, Z, cmap=plt.cm.coolwarm, linewidth=0, antialiased=False)
cset = ax.contourf(X, Y, Z, zdir='z', offset=-2, cmap=plt.cm.coolwarm)

# Add a color bar which maps values to colors.
fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()

高级用法：使用Shapely进行轮廓剪裁

在实际应用中，我们可能需要将contourf图形剪裁为特定形状，以表达某种含义。例如，在地形图中，我们可以用圆形轮廓剪裁出一个由高程信息构成的山头。

为了实现轮廓剪裁，我们可以使用Python的Shapely库。Shapely是一个用于二维几何运算的Python库，可以用于计算点、线、面等几何对象的空间关系。

首先，我们需要将contourf图形转换为Shapely的多边形（polygon）对象。我们可以使用Matplotlib的QuadContourSet对象从contourf图形中提取多边形轮廓。

from matplotlib import path
from shapely.geometry import Polygon

cs = plt.contour(X, Y, Z, levels=[-1, 0, 1], colors='k')
ps = cs.collections[0]

polygons = []
for contour in ps:
    for path in contour.get_paths():
        polygons.append(Polygon(path.vertices))

注意这里levels=[-1, 0, 1]表示我们只提取三个等高线，即高度为-1、0、1的轮廓。

然后，我们可以定义一个剪裁区域，例如一个圆形：

from shapely.geometry import Point

clip_radius = 3
clip_center = Point(0, 0)
clip_area = clip_center.buffer(clip_radius)

这里buffer函数表示将点周围的区域扩展为圆形。

最后，我们可以使用Shapely的intersection函数对多边形轮廓进行剪裁：

from shapely.ops import cascaded_union

clipped_polygons = []
for polygon in polygons:
    clipped_poly = polygon.intersection(clip_area)
    if clipped_poly.geom_type == 'Polygon':
        clipped_polygons.append(clipped_poly)
    elif clipped_poly.geom_type == 'MultiPolygon':
        clipped_polygons += list(clipped_poly)
clip_path = mplPath.Path(np.concatenate([np.asarray(p.exterior.coords) for p in clipped_polygons]))

# plot clipped contourf
ax.contourf(X, Y, Z, zdir='z', offset=-2, cmap=plt.cm.coolwarm, levels=[-1, 0, 1], clip_path=clip_path)

注意这里使用了list函数将MultiPolygon对象的各个部分拆分为多个Polygon对象。定义剪裁路径后，我们可以将其作为参数传入contourf函数中，绘制剪裁后的图形。

完整的代码如下：

import numpy as np
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import path as mplPath
from shapely.geometry import Polygon, Point
from shapely.ops import cascaded_union

X, Y = np.meshgrid(np.linspace(-5, 5, 50), np.linspace(-5, 5, 50))
Z = np.sin(np.sqrt(X**2 + Y**2))

fig = plt.figure()
ax = fig.gca(projection='3d')

# surface plot
surf = ax.plot_surface(X, Y, Z, cmap=plt.cm.coolwarm, linewidth=0, antialiased=False)

# contourf plot
cs = plt.contour(X, Y, Z, levels=[-1, 0, 1], colors='k')
ps = cs.collections[0]

polygons = []
for contour in ps:
    for path in contour.get_paths():
        polygons.append(Polygon(path.vertices))

clip_radius = 3
clip_center = Point(0, 0)
clip_area = clip_center.buffer(clip_radius)

clipped_polygons = []
for polygon in polygons:
    clipped_poly = polygon.intersection(clip_area)
    if clipped_poly.geom_type == 'Polygon':
        clipped_polygons.append(clipped_poly)
    elif clipped_poly.geom_type == 'MultiPolygon':
        clipped_polygons += list(clipped_poly)
clip_path = mplPath.Path(np.concatenate([np.asarray(p.exterior.coords) for p in clipped_polygons]))

cset = ax.contourf(X, Y, Z, zdir='z', offset=-2, cmap=plt.cm.coolwarm, levels=[-1, 0, 1], clip_path=clip_path)

# Add a color bar which maps values to colors.
fig.colorbar(surf, shrink=0.5, aspect=5)

plt.show()