Question

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

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

在我的问题中，我试图通过对每个方形（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没有接受非方形矩阵（因此迫使我的基于循环的解决方案）。也许einsum也可以帮到这里吗？

Answer 1

一种方法是重塑为2D，使用对角线值设置ncols+1步的列。重塑创建视图，因此允许我们直接访问这些对角线位置。因此，实施将是 -

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

Answer 2

如果您执行np.source(np.fill_diagonal)，您会在第二种情况下看到它使用了一个跨步的＆＃39;方法

    if a.ndim == 2:
        step = a.shape[1] + 1
        end = a.shape[1] * a.shape[1]
    a.flat[:end:step] = val

@Divakar's解决方案通过＆＃39;展平＆＃39;在2个方面。

您可以使用M.sum(axis=1)对列进行求和。虽然我模糊地回忆起一些时间，发现einsum实际上要快一点。 sum更为传统。

有人要求能够在einsum扩展维度，但我认为不会发生这种情况。

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

2 个答案: