矢量化物理模拟？

Question

我正在尝试模拟一些二维粒子。每个粒子都是一个有方向的圆。方向由二维单位向量指定。

在我的模拟的一部分中，我想计算粒子对之间的角度和每个粒子的方向的函数。这应该对每个粒子对进行。在视觉上，我想计算 $\theta_i$ 和 $\theta_j$ 的函数（参见图片链接）。

我已经计算了每个粒子对的成对位移单位向量。这是一个名为 r 的形状为 (N, N, 2) 的 numpy 数组，其中 N 是粒子的总数。我还计算了笛卡尔坐标中每个粒子的方向。这是一个名为orientation of shape (N, 2)的numpy数组。

我已经能够编写我需要的代码作为双 for 循环。

output = np.zeros(r.shape)
for i in range(len(r)):
      for j in range(len(r[i])):
        if (i < j):
          theta_i = angle_between(orientation[i], r[i,j])
          theta_j = angle_between(orientation[j], r[i,j])
          output[i][j] = (1 / (1 + np.exp(-w*(np.cos(theta_i) - np.cos(alpha))))) * (1 / (1 + np.exp(-w*(np.cos(theta_j) - np.cos(alpha)))))

不过，运行此代码需要很长时间，尤其是对于大型系统。有没有办法对双 for 循环进行矢量化，使其运行得更快？

Answer 1

这是我运行您的代码的示例：

import numpy as np
np.random.seed(0)

N = 3
r_hat = np.random.rand(N, N, 2)
moment_cartesian = np.random.rand(N, 2)
output = np.zeros(r_hat.shape)
output = np.zeros((N, N))
w = 1
alpha = 1
theta_i_ij = np.zeros((N, N))
theta_j_ij = np.zeros((N, N))
for i in range(len(r_hat)):
      for j in range(len(r_hat[i])):
        if (i < j):
          theta_i = angle_between(moment_cartesian[i], r_hat[i,j])
          theta_j = angle_between(moment_cartesian[j], r_hat[i,j])
          theta_i_ij[i, j] = theta_i
          theta_j_ij[i, j] = theta_j
          output[i][j] = (1 / (1 + np.exp(-w*(np.cos(theta_i) - np.cos(alpha))))) * (1 / (1 + np.exp(-w*(np.cos(theta_j) - np.cos(alpha)))))

print(f'theta_i_ij = \n{theta_i_ij}')
print(f'theta_j_ij = \n{theta_j_ij}')
print(f'output = \n{output}')

输出：

theta_i_ij = 
[[0.         0.99438035 0.98889055]
 [0.         0.         0.9954101 ]
 [0.         0.         0.        ]]
theta_j_ij = 
[[0.         0.99873953 0.99891568]
 [0.         0.         0.90135909]
 [0.         0.         0.        ]]
output = 
[[0.         0.25072287 0.25127888]
 [0.         0.         0.26052633]
 [0.         0.         0.        ]]

以下是代码矢量化的示例：

innerp    = np.einsum('jk,ijk->ij', moment_cartesian, r_hat)
norm      = np.linalg.norm(moment_cartesian, axis = -1)[None, :]*np.linalg.norm(r_hat, axis = -1)
theta_j_v = innerp/norm

innerp    = np.einsum('ik,ijk->ij', moment_cartesian, r_hat)
norm      = np.linalg.norm(moment_cartesian, axis = -1)[:, None]*np.linalg.norm(r_hat, axis = -1)
theta_i_v = innerp/norm

output_v = (1 / (1 + np.exp(-w*(np.cos(theta_i_v) - np.cos(alpha))))) * (1 / (1 + np.exp(-w*(np.cos(theta_j_v) - np.cos(alpha)))))


print(f'theta_i_v = \n{theta_i_v}')
print(f'theta_j_v = \n{theta_j_v}')
print(f'output_v = \n{output_v}')

输出：

theta_i_v = 
[[0.99717384 0.99438035 0.98889055]
 [0.9090371  0.95351666 0.9954101 ]
 [0.99986414 0.98876437 0.87290329]]
theta_j_v = 
[[0.99717384 0.99873953 0.99891568]
 [0.96281813 0.95351666 0.90135909]
 [0.98397106 0.9796661  0.87290329]]
output_v = 
[[0.25059435 0.25072287 0.25127888]
 [0.26327747 0.25972068 0.26052633]
 [0.2516916  0.25331216 0.27620631]]

但我必须说，我没有想到一种方法来矢量化您的代码，而无需计算矩阵的下三角形。 您实际上根本不需要矢量化您的代码。 如果您使用 Numba 编译包含您的代码的函数，它执行相同操作的速度与矢量化解决方案一样快，甚至更快。

这是我的 Numba 示例：

from numba import jit

@jit(nopython=True)
def computations(r_hat, moment_cartesian, alpha, w):
    output = np.zeros((N, N))
    for i in range(len(r_hat)):
        for j in range(len(r_hat[i])):
            if (i < j):
                v0 = r_hat[i,j]
                v1 = moment_cartesian[i]
                v2 = moment_cartesian[j]
                theta_i = np.sum(v0*v1)/(np.sum(v0**2)*np.sum(v1**2))**0.5
                theta_j = np.sum(v0*v2)/(np.sum(v0**2)*np.sum(v2**2))**0.5
                output[i][j] = (1 / (1 + np.exp(-w*(np.cos(theta_i) - np.cos(alpha))))) * (1 / (1 + np.exp(-w*(np.cos(theta_j) - np.cos(alpha)))))
    return output

print(f'output = \n{computations(r_hat, moment_cartesian, alpha, w)}')

输出：

output = 
[[0.         0.25072287 0.25127888]
 [0.         0.         0.26052633]
 [0.         0.         0.        ]]