如何借助Meshgrid工具优化Numpy中目标数组C的生成代码？

已通过以下代码生成网格数组omega、phiy和phix：

import numpy as np

phi = np.linspace(-0.75, 0.75, num=4)
omega = np.linspace(-100, 100, num=5)
omega, phiy, phix = np.meshgrid(omega, phi, phi, sparse=True, indexing='ij')

rand = np.random.rand(5, 4)

需要生成数组C，当前通过三重循环实现：

def f(j,k):
   return np.argmin(np.abs(phi[j]**2 + np.sin(phi[k]) - phi))

C = np.empty((5,4,4))
for i in range(5):
  for j in range(4):
    for k in range(4):
       C[i,j,k] = rand[i, f(j,k)]

请问能否借助numpy的网格/矢量化工具，实现更紧凑、高效的写法？

当然可以，核心是矢量化计算所有(j,k)组合对应的索引，再利用numpy的广播索引直接生成C，彻底去掉Python循环：

1. 矢量化生成所有(j,k)对应的索引 #

先一次性计算所有j和k组合下，phi[j]² + sin(phi[k])在phi中最接近值的索引：

# 生成j、k的二维网格，覆盖所有(4,4)组合
j_grid, k_grid = np.meshgrid(np.arange(4), np.arange(4), indexing='ij')
# 计算所有组合的目标值：phi[j]^2 + sin(phi[k])
target_vals = phi[j_grid] ** 2 + np.sin(phi[k_grid])
# 广播计算每个目标值与phi所有元素的绝对差，取最小差的索引
idx = np.argmin(np.abs(target_vals[..., np.newaxis] - phi), axis=-1)

这里通过[..., np.newaxis]给target_vals增加维度，让它能和一维的phi做广播运算，最后沿新增维度取argmin，得到形状为(4,4)的索引数组idx，对应原函数f(j,k)的所有结果。

2. 广播索引生成C #

有了idx后，直接利用numpy的广播机制生成C：

C = rand[:, idx]

rand形状为(5,4)，idx形状为(4,4)，numpy会自动把rand广播为(5,4,4)，再按idx的索引取值，最终得到形状为(5,4,4)的C，和原循环结果完全一致。

完整代码 #

import numpy as np

phi = np.linspace(-0.75, 0.75, num=4)
omega = np.linspace(-100, 100, num=5)
omega, phiy, phix = np.meshgrid(omega, phi, phi, sparse=True, indexing='ij')

rand = np.random.rand(5, 4)

# 矢量化计算所有(j,k)对应的索引
j_grid, k_grid = np.meshgrid(np.arange(4), np.arange(4), indexing='ij')
target_vals = phi[j_grid] ** 2 + np.sin(phi[k_grid])
idx = np.argmin(np.abs(target_vals[..., np.newaxis] - phi), axis=-1)

# 生成最终数组C
C = rand[:, idx]

效率优势 #

这种矢量化写法完全依托numpy的C底层运算，避免了Python循环的性能损耗。当phi的长度增大时（比如从4到100），速度会比循环版本快几十甚至上百倍。

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