Python绘制二元函数v=cos(2πz)sin(πy)时数组维度不匹配报错求解

这是一个Notebook相关问题，我目前无法得到正确结果，具体问题如下：

• 变量y的取值范围为0到5.3，共130个分段
• 变量z的取值范围为0到0.3，共48个分段
• 我需要绘制函数v，表达式为：
  $v = cos(2\pi z) \times sin(\pi y)$

我尝试用Python编写了如下代码：

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import math

y = np.arange(0, 5.3, 0.007692)
z = np.arange(0,0.3,0.021)
v = np.cos(2.0*math.pi*z)*np.sin(math.pi*y)

fig = plt.figure(figsize=(6, 6))
ax = fig.add_subplot(111, projection='3d')
ax.scatter( y, z, v,
           linewidths=1, alpha=.7,
           edgecolor='k',
           s = 200,
           c=v)
plt.show()

运行得到报错：

v = np.cos(2.0*math.pi*z)*np.sin(math.pi*y)
ValueError: operands could not be broadcast together with shapes (15,) (690,) 

后续我参考建议使用np.meshgrid生成了等高线图，和预期结果接近，但还想要得到y、z、v的3D散点图，修改后的代码如下：

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import math

y = np.arange(0, 5.3, 0.007692)
z = np.arange(0,0.3,0.021)
xx, yy = np.meshgrid(y, z, sparse=True)
v = np.cos(2.0*math.pi*xx)*np.sin(math.pi*yy)

h = plt.contourf(y,z,v)
plt.colorbar()
plt.show()

fig = plt.figure(figsize=(6, 6))
ax = fig.add_subplot(111, projection='3d')
ax.scatter( y, z, v,
           linewidths=1, alpha=.7,
           edgecolor='k',
           s = 200,
           c=v)
plt.show()

运行结果示意图：

但运行时又出现了如下报错：

ValueError: shape mismatch: objects cannot be broadcast to a single shape

错误原因 #

两次报错均为数组维度不匹配导致的广播异常：

1. 首次报错是由于y、z均为一维数组但长度不同，直接进行元素级运算不符合numpy广播规则。
2. 二次报错是因为经np.meshgrid计算得到的v为二维数组，而传入scatter方法的y、z仍为原始一维数组，三者维度无法匹配。

修正方案 #

1. 生成y、z数组时使用np.linspace方法直接指定采样点数，无需手动计算步长，可避免精度误差和点数不符的问题。
2. 调用np.meshgrid时移除sparse=True参数，生成与v维度完全匹配的二维坐标网格。
3. 3D散点图要求输入的x、y、z参数均为一维数组，调用.flatten()方法将二维的坐标网格和v数组展平即可。

完整可运行代码 #

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

# 直接指定采样点数生成数组
y = np.linspace(0, 5.3, 130)
z = np.linspace(0, 0.3, 48)

# 生成完整二维网格
xx, yy = np.meshgrid(y, z)
# 计算v，维度与xx、yy完全一致
v = np.cos(2.0 * math.pi * xx) * np.sin(math.pi * yy)

# 绘制等高线图
h = plt.contourf(y, z, v)
plt.colorbar()
plt.show()

# 绘制3D散点图，所有参数展平为一维
fig = plt.figure(figsize=(6, 6))
ax = fig.add_subplot(111, projection='3d')
ax.scatter(xx.flatten(), yy.flatten(), v.flatten(),
           linewidths=1, alpha=.7,
           edgecolor='k',
           s=20, # 总点数较多，调小尺寸避免重叠
           c=v.flatten())
plt.show()

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