设置（M x N x N）矩阵的对角线的快速方法？ Einsum / n维填充对角线？答案

【问题标题】：Fast way to set diagonals of an (M x N x N) matrix? Einsum / n-dimensional fill_diagonal?设置（M x N x N）矩阵的对角线的快速方法？ Einsum / n维填充对角线？
【发布时间】：2017-10-14 20:42:50
【问题描述】：

我正在尝试基于矩阵编写快速、优化的代码，并且最近发现了 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 也可以在这里提供帮助？

【问题讨论】：

标签： numpy matrix optimization diagonal numpy-einsum

【解决方案1】：

一种方法是重塑为2D，将列设置在ncols+1 的步长处，并使用对角线值。重塑创建了一个视图，因此我们可以直接访问这些对角线位置。因此，实现将是 -

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

【讨论】：

谢谢。这很好用。有没有进一步的方法来加快速度，以避免进行多次重塑？
@PhysLQ 好吧，重塑以创建视图是对数组进行查看访问的最佳方式，并且重塑几乎是免费的。我没有看到建议的方法对 reshape 有任何其他用法，还是我们需要在其他地方进行 reshape ？

【解决方案2】：

如果您使用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 解决方案通过在 2 维上“展平”将其应用于您的 3d 案例。

您可以将列与M.sum(axis=1) 相加。虽然我隐约记得有一些时间发现einsum 实际上要快一点。 sum 更传统一点。

有人要求在einsum 中扩展维度，但我认为这不会发生。

【讨论】：

那么，复制和粘贴np.fill_diagonal 的源代码如何帮助改进它以加快对角线填充？
我的回答更像是你的一个脚注。