numpy - 设置(M x N x N)矩阵对角线的快速方法？ Einsum/n维fill_diagonal？

Question

我正在尝试编写基于矩阵的快速优化代码，并且最近发现 einsum 作为实现显着加速的工具。

是否可以使用它来有效地设置多维数组的对角线，还是只能返回数据？

在我的问题中，我试图通过对每个方阵 (N x 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 不接受非方矩阵(因此强制我基于循环的解决方案)。也许 einsum 也可以提供帮助？

Answer 1

一种方法是 reshape 为 2D，在 ncols+1 步长处使用对角线值设置列。 reshape 创建一个 View ，因此允许我们直接访问那些对角线位置。因此，实现将是 -

s0,s1,s2 = M.shape
M.reshape(s0,-1)[:,::s2+1] = diags