快速设置(M x N x N)矩阵对角线的方法? Einsum/n维fill_diagonal?
问题描述:
我正在尝试基于矩阵编写快速优化的代码,并且最近发现了einsum作为实现显着加速的工具。快速设置(M x N x N)矩阵对角线的方法? Einsum/n维fill_diagonal?
是否可以使用它来有效地设置多维数组的对角线,还是只能返回数据?
在我的问题中,我试图通过求和每个正方形(N×N)矩阵中的列来设置矩阵矩阵(形状:M x N x N)的对角线。
我现在的(基于循环慢,)的解决方案是:
# Build dummy array
dimx = 2 # Dimension x (likely to be < 100)
dimy = 3 # Dimension y (likely to be between 2 and 10)
M = np.random.randint(low=1, high=9, size=[dimx, dimy, dimy])
# Blank the diagonals so we can see the intended effect
np.fill_diagonal(M[0], 0)
np.fill_diagonal(M[1], 0)
# Compute diagonals based on summing columns
diags = np.einsum('ijk->ik', M)
# Set the diagonal for each matrix
# THIS IS LOW. CAN IT BE IMPROVED?
for i in range(len(M)):
np.fill_diagonal(M[i], diags[i])
# Print result
M
这能在所有的改进吗?看来np.fill_diagonal不接受非方形矩阵(因此强制我的基于循环的解决方案)。也许爱因斯也可以在这里帮助?
答
一种方法是重塑2D
,将对角线值设置为步骤ncols+1
。重新塑造创造了一个视图,因此我们可以直接访问这些对角线位置。因此,实现起来 -
s0,s1,s2 = M.shape
M.reshape(s0,-1)[:,::s2+1] = diags
答
如果你这样做np.source(np.fill_diagonal)
你会看到,在2D情况下,它通过使用“跨入”的做法
if a.ndim == 2:
step = a.shape[1] + 1
end = a.shape[1] * a.shape[1]
a.flat[:end:step] = val
@Divakar's
解决方案适用给你的3D案例在2个维度上“展平”。
您可以用M.sum(axis=1)
对列进行求和。虽然我隐约记得一些时间发现einsum
实际上有点快。 sum
稍微传统一些。
有人要求能够扩大尺寸einsum
,但我不认为会发生。
谢谢。这很好。有没有其他方法可以加快速度,避免进行多次改造? – PhysLQ
@PhysLQ那么,重塑视图是对数组进行查看访问并重构几乎是免费的最佳方式。我没有看到任何其他用于重建所提议的方法的用法,还是我们需要在其他地方进行重塑? – Divakar