如何用Numpy优雅生成3D绘图所需的二维Z值数组？

嘿，这个需求太常见了！要快速生成对应二维Z数组，完全不用写繁琐的嵌套循环，用Numpy的向量化特性就能搞定最优雅的写法，分两种情况给你说：

情况1：你的映射函数是Numpy向量化的 #

如果你的compute_z函数本身就是用Numpy操作编写的（比如用np.sin、np.sqrt这些原生Numpy函数），那步骤超简单：

1. 先把一维的X、Y转成网格矩阵（如果还没转的话）
2. 直接把网格矩阵传入函数，Numpy会自动广播计算每个(x,y)对应的Z值

示例代码：

import numpy as np

# 假设你的一维X、Y数组
X = np.linspace(-5, 5, 100)
Y = np.linspace(-5, 5, 100)

# 向量化的映射函数
def compute_z(x, y):
    return np.sin(np.sqrt(x**2 + y**2))

# 生成网格并计算Z
X_grid, Y_grid = np.meshgrid(X, Y)
Z = compute_z(X_grid, Y_grid)

这种写法完全利用Numpy的底层优化，不仅代码简洁，运行效率也比手动循环高得多。

情况2：你的映射函数是纯Python函数（非向量化） #

如果你的compute_z是用纯Python写的（比如用标准库math模块的函数），可以用np.vectorize把它包装成向量化函数，然后同样直接传入网格矩阵：

示例代码：

import numpy as np
import math

# 纯Python的映射函数
def compute_z(x, y):
    return math.sin(math.sqrt(x**2 + y**2))

# 转成向量化函数
vec_compute_z = np.vectorize(compute_z)

# 生成网格并计算Z
X_grid, Y_grid = np.meshgrid(X, Y)
Z = vec_compute_z(X_grid, Y_grid)

注意：np.vectorize本质是语法糖，内部还是循环，但胜在写法简洁，适合快速实现；如果追求极致效率，建议把纯Python函数改成Numpy向量化的形式。

更进阶的紧凑写法 #

如果不想单独写meshgrid的行，还可以把生成网格和计算Z合并成一行（可读性稍降，看个人偏好）：

Z = compute_z(*np.meshgrid(X, Y))

这样一行就搞定了网格生成和Z计算，非常紧凑！

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