使用Matplotlib axes3d绘制带投影等高线的3D曲面时z维度报错解决

使用axes3d绘制以a、b为变量的多元函数f_ab时，抛出TypeError: Input z must be 2D, not 1D错误，但检查发现f_ab形状为(50,50)。

原代码如下：

import numpy as np

a = np.linspace(5,10,50)
b = np.linspace(3,8,50)

f_ab = np.zeros((len(a),len(b)))

for i in range(len(a)):
    for j in range(len(b)):
        f_ab[i,j] = (a[i]**0.5)*(b[j])
        
        
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt


ax = plt.figure().add_subplot(projection='3d')
X, Y, Z = a, b, f_ab

# Plot the 3D surface
ax.plot_surface(X, Y, Z, edgecolor='royalblue', lw=0.5, rstride=8, cstride=8,
                alpha=0.3)

# Plot projections of the contours for each dimension.  By choosing offsets
# that match the appropriate axes limits, the projected contours will sit on
# the 'walls' of the graph.
ax.contour(X, Y, Z, zdir='z', offset=np.min(f_ab), cmap='coolwarm')
ax.contour(X, Y, Z, zdir='x', offset=50, cmap='coolwarm')
ax.contour(X, Y, Z, zdir='y', offset=50, cmap='coolwarm')

ax.set(xlim=(0,50), ylim=(0,50), zlim=(np.min(f_ab),np.max(f_ab)),
       xlabel='a', ylabel='b', zlabel='z')

错误信息：

File ~\AppData\Roaming\Python\Python310\site-packages\matplotlib\contour.py:1501, in QuadContourSet._check_xyz(self, args, kwargs)
   1498 z = ma.asarray(args[2], dtype=np.float64)
   1500 if z.ndim != 2:
-> 1501     raise TypeError(f"Input z must be 2D, not {z.ndim}D")
   1502 if z.shape[0] < 2 or z.shape[1] < 2:
   1503     raise TypeError(f"Input z must be at least a (2, 2) shaped array, "
   1504                     f"but has shape {z.shape}")

TypeError: Input z must be 2D, not 1D

错误核心是ax.contour要求X、Y必须是和Z同形状的二维网格数组，而原代码中X、Y是一维数组，导致contour无法正确解析Z维度。

修复步骤：

• 用np.meshgrid将一维的a、b转换为二维网格数组X、Y，确保与Z（f_ab）形状匹配
• 修正坐标轴范围，匹配a、b的实际取值区间，避免投影等高线位置错误

修正后的完整代码：

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

a = np.linspace(5,10,50)
b = np.linspace(3,8,50)
# 生成二维网格数组
X, Y = np.meshgrid(a, b)
# 直接用网格数组计算Z，替代双层循环（更高效）
Z = (X**0.5) * Y

ax = plt.figure().add_subplot(projection='3d')

# 绘制3D曲面，现在X、Y是二维网格，兼容plot_surface和contour
ax.plot_surface(X, Y, Z, edgecolor='royalblue', lw=0.5, rstride=8, cstride=8,
                alpha=0.3)

# 绘制各维度的投影等高线
z_min, z_max = np.min(Z), np.max(Z)
ax.contour(X, Y, Z, zdir='z', offset=z_min, cmap='coolwarm')
ax.contour(X, Y, Z, zdir='x', offset=np.min(a), cmap='coolwarm')  # 用a的最小值作为x轴偏移
ax.contour(X, Y, Z, zdir='y', offset=np.min(b), cmap='coolwarm')  # 用b的最小值作为y轴偏移

# 设置正确的坐标轴范围
ax.set(xlim=(np.min(a), np.max(a)), ylim=(np.min(b), np.max(b)), zlim=(z_min, z_max),
       xlabel='a', ylabel='b', zlabel='z')

plt.show()

额外优化：用网格数组直接计算Z，替代双层循环，代码更简洁高效。

内容的提问来源于stack exchange，提问作者numpy