如何对Numpy函数进行“双重向量化”以高效计算跨数组结果？

假设有一个函数f，其前3个位置参数支持向量化；该函数完全由支持广播的Numpy向量化函数（算术、三角函数等）实现。
函数f的前两个参数x和y为等长一维数组，x[i]与y[i]对应数据集中的第i个数据点；第三个参数q是另一集合中的一维数组，长度与x、y不同，q[k]对应其中第k个数据点。

需要计算所有(i,k)对对应的f(x[i], y[i], q[k])值，并将结果整理为矩阵。目标是实现如下嵌套循环计算的完全向量化版本：

result = np.empty((len(x), len(q)))
for k in range(len(q)):
    for i in range(len(x)):
        result[i, k] = f(x[i], y[i], q[k])

目前使用的是仅对i索引向量化的版本，仍需循环q的元素：

result = np.empty((len(x), len(q)))
for k in range(len(q)):
    result[:, k] = f(x, y, q[k])

是否存在高效方法实现对两个索引的完全向量化？比如利用Numpy的广播技巧？

示例函数：余弦定理实现 #

以余弦定理函数为例，它完全由Numpy广播兼容的操作构成：

def law_of_cosines(a, b, ϑ):
    return np.sqrt(
        np.square(a) +
        np.square(b) +
        2.0 * a * b * np.cos(ϑ)
    )

解决方案：利用Numpy广播实现完全向量化 #

核心思路是通过维度扩展让x、y和q的形状满足广播规则，从而一次性计算所有(i,k)对的结果：

1. 将x和y从一维数组（形状(N,)）扩展为二维数组（形状(N, 1)），这样它们可以和形状为(M,)的q进行广播，最终得到形状为(N, M)的结果。
2. 直接调用f函数，传入扩展后的x、y和原始q，Numpy会自动完成广播计算。

针对示例函数的实现代码：

import numpy as np

def law_of_cosines(a, b, ϑ):
    return np.sqrt(
        np.square(a) +
        np.square(b) +
        2.0 * a * b * np.cos(ϑ)
    )

# 示例数据
x = np.array([1.0, 2.0, 3.0])
y = np.array([4.0, 5.0, 6.0])
q = np.array([np.pi/3, np.pi/2, np.pi/4, np.pi/6])

# 维度扩展：x -> (3,1), y -> (3,1)
result = law_of_cosines(x[:, np.newaxis], y[:, np.newaxis], q)
# 或者用np.expand_dims：np.expand_dims(x, axis=1)
print(result.shape)  # 输出 (3, 4)，符合预期的(N,M)矩阵

原理说明 #

当x被扩展为(N,1)，y扩展为(N,1)，q保持(M,)时，Numpy的广播机制会自动将q扩展为(1,M)，然后和(N,1)的x、y进行逐元素运算，最终得到(N,M)的结果矩阵。这种方式完全避免了Python层面的循环，所有计算都在底层C实现的Numpy操作中完成，效率远高于循环版本。

通用适配 #

对于任意满足条件的函数f（仅使用Numpy广播兼容的操作），只需对前两个参数进行最后一维扩展，第三个参数保持原形状即可实现完全向量化计算。例如：

# 通用写法
result = f(x[..., np.newaxis], y[..., np.newaxis], q)

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