我有一个3x 3的矩阵,我可以计算如下
e_matrix = np.array([m[i]*m[j]-n[i]*n[j] for i in range(3) for j in range(3)])
字符串
其中m
和n
都是长度为3的向量。
现在假设m
和n
是形状为(K,3)
的矩阵,我想做一个类似于上面的计算,得到一个形状为(3,3,K)
的e_matrix
。
我意识到我可以做一个天真的方法,例如。
e_matrix = np.zeros((3,3,K))
for k in range(K):
e_matrix[:,:,k] = np.array([m[k,i]*m[k,j]-n[k,i]*n[k,j] for i in range(3) for j in range(3)])
型
是否存在矢量化方法?理想情况下,它可以与Numba/JIT编译一起工作(所以没有,例如np.tensordot
)。
3条答案
按热度按时间a0zr77ik1#
通常在raw numpy中通过broadcasting:
字符串
我想这和伦巴是兼容的,但我不太熟悉...
ryoqjall2#
这是使用numba的“矢量化”变体。
字符串
当K非常大(例如,100,000)时,该函数很快,但不幸的是,当K很小时,由于开销很大,该函数非常慢。
如果性能是你的优先考虑,我更喜欢一个简单的jit函数。
型
基准:
型
结果(K = 100_000):
型
结果(K = 100):
型
p4rjhz4m3#
您可以通过更改
[:, None]
的.来使Chrysophylaxs' answer与numba(<0. 57)njit兼容,以使用ndarray.reshape()
代替,因为这是supported by numba:字符串
输出量:
型