NumPy：如何将向量值函数应用于网格点对？

原代码错误原因 #

np.meshgrid(X1,X2)返回的是两个形状为(3,3)的数组组成的列表。当直接将这个列表传给np.vectorize(f)时，np.vectorize会将这两个数组视为多个输入参数，并且对每个位置的标量元素单独处理——即把X[0][i,j]作为唯一参数传给f，此时f中的x是一个标量，自然无法执行x[0]索引操作，从而抛出IndexError。

方案1：修改函数参数为独立变量 #

将函数f改为接受两个独立的参数x1和x2，再将meshgrid生成的两个数组分别传入矢量化后的函数：

import numpy as np

def f(x1, x2):
    return x1 + x2

X1 = np.array([0, 1, 2])
X2 = np.array([0, 1, 2])
X = np.meshgrid(X1, X2)

result = np.vectorize(f)(X[0], X[1])
print(result)
# 输出：
# [[0 1 2]
#  [1 2 3]
#  [2 3 4]]

原理：np.vectorize会将传入的多个数组对应位置的元素作为参数传递给f，每个位置的(x1[i,j], x2[i,j])会被正确传入函数计算，完全适配带条件语句等复杂逻辑的f。

方案2：保留原函数参数形式，使用signature指定输入形状 #

如果不想修改f的参数形式，可以将meshgrid的结果合并为每个位置包含(x1,x2)对的数组，同时通过signature参数告诉np.vectorize函数f接受的是一个长度为2的一维数组，输出是标量：

import numpy as np

def f(x):
    return x[0] + x[1]

X1 = np.array([0, 1, 2])
X2 = np.array([0, 1, 2])
X = np.meshgrid(X1, X2)

# 将两个网格数组合并为形状(3,3,2)的数组，每个元素是(x1,x2)对
X_combined = np.stack(X, axis=-1)
# 使用signature指定输入为(n,)的数组，输出为标量
result = np.vectorize(f, signature='(n)->()')(X_combined)
print(result)
# 输出：
# [[0 1 2]
#  [1 2 3]
#  [2 3 4]]

原理：signature参数强制np.vectorize按照指定的输入输出形状处理数据，不再将数组拆分为单个标量，而是把每个(2,)的子数组作为整体传入f，适合需要保留原函数参数结构的场景。

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