如何对二维粒子物理模拟的双重for循环进行numpy向量化优化？

两个单位向量的夹角余弦值等于二者的点积，不需要额外计算夹角θ，省去angle_between的调用开销，再通过numpy广播机制完全替代双重循环即可实现大幅提速。

完整向量化代码 #

import numpy as np

# 提前计算公共常量
cos_alpha = np.cos(alpha)

# 向量化计算所有cosθ_i：输出shape为(N,N)
cos_theta_i = np.sum(orientation[:, np.newaxis, :] * r, axis=-1)
# 向量化计算所有cosθ_j：输出shape为(N,N)
cos_theta_j = np.sum(orientation[np.newaxis, :, :] * r, axis=-1)

# 向量化计算sigmoid乘积
sigmoid_i = 1 / (1 + np.exp(-w * (cos_theta_i - cos_alpha)))
sigmoid_j = 1 / (1 + np.exp(-w * (cos_theta_j - cos_alpha)))
res = sigmoid_i * sigmoid_j

# 仅保留i<j的上三角部分，其余位置置0，和原逻辑完全对齐
output = np.zeros_like(res)
upper_mask = np.triu_indices_from(output, k=1)
output[upper_mask] = res[upper_mask]

# 若你原本确实需要和r同shape的(N,N,2)输出，可添加下方代码
# output = np.repeat(output[..., np.newaxis], 2, axis=-1)

优化说明 #

• 广播机制：通过给orientation增加空维度，自动匹配r的(N,N,2) shape，一次性完成所有粒子对的点积计算，所有运算都在numpy底层C实现，比Python级别的双重循环快2~3个数量级，粒子数越大提升越明显。
• 冗余计算裁剪：直接用点积得到cosθ，省去了角度计算、三角函数调用的额外开销。
• 结果对齐：通过triu_indices_from仅保留i<j的上三角赋值，其余位置为0，和原循环逻辑完全一致。

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