这是我运行您的代码的示例:
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. ]]