从矩阵的 3D 堆栈在 numpy/scipy 中构造块对角矩阵的 3D 堆栈的有效方法

Question

我正在尝试从给定的矩阵堆栈 (nXmXm) 中以 nXMXM 的形式在 numpy/scipy 中构造一个块对角矩阵堆栈，其中 M=k*m，k 是矩阵堆栈的数量。目前，我在 for 循环中使用 scipy.linalg.block_diag 函数来执行此任务：

import numpy as np
import scipy.linalg as linalg

a = np.ones((5,2,2))
b = np.ones((5,2,2))
c = np.ones((5,2,2))

result = np.zeros((5,6,6))

for k in range(0,5):
    result[k,:,:] = linalg.block_diag(a[k,:,:],b[k,:,:],c[k,:,:])

但是，由于在我的情况下 n 变得相当大，我正在寻找一种比 for 循环更有效的方法。我找到了3D numpy array into block diagonal matrix，但这并不能真正解决我的问题。我能想象的任何事情都是将每个矩阵堆栈转换为块对角线

import numpy as np
import scipy.linalg as linalg

a = np.ones((5,2,2))
b = np.ones((5,2,2))
c = np.ones((5,2,2))

a = linalg.block_diag(*a)
b = linalg.block_diag(*b)
c = linalg.block_diag(*c)

并通过重塑来构造结果矩阵

result = linalg.block_diag(a,b,c)

result = result.reshape((5,6,6))

不会重塑。我什至不知道，如果这种方法会更有效，所以我问我是否走在正确的轨道上，或者是否有人知道构建这个块对角 3D 矩阵的更好方法，或者我是否必须坚持for 循环解决方案。

编辑： 由于我是这个平台的新手，我不知道该把它留在哪里（编辑或回答？），但我想分享我的最终解决方案：panadestein 的 highlightet 解决方案非常好用且简单，但我现在使用高维数组，我的矩阵位于最后两个维度。此外，我的矩阵不再具有相同的维度（主要是 1x1、2x2、3x3 的混合），因此我采用了 V. Ayrat 的解决方案，并进行了细微的更改：

def nd_block_diag(arrs):
    shapes = np.array([i.shape for i in arrs])

    out = np.zeros(np.append(np.amax(shapes[:,:-2],axis=0), [shapes[:,-2].sum(), shapes[:,-1].sum()]))
    r, c = 0, 0
    for i, (rr, cc) in enumerate(shapes[:,-2:]):
        out[..., r:r + rr, c:c + cc] = arrs[i]
        r += rr
        c += cc

    return out

如果输入数组的形状正确（即要广播的维度不会自动添加），它也适用于数组广播。感谢 pandestein 和 V. Ayrat 的友好和快速的帮助，我学到了很多关于列表推导和数组索引/切片的可能性！

Answer 1

我不认为您可以逃脱所有可能的循环来解决您的问题。我发现方便，也许比你的for循环更效率的方式是使用列表理解：

import numpy as np
from scipy.linalg import block_diag

# Define input matrices

a = np.ones((5, 2, 2))
b = np.ones((5, 2, 2))
c = np.ones((5, 2, 2))

# Generate block diagonal matrices

mats = np.array([a, b, c]).reshape(5, 3, 2, 2)
result = [block_diag(*bmats) for bmats in mats]

也许这可以给你一些想法来改善你的实现。

Answer 2

block_diag 也只是遍历形状。几乎所有时间都花在复制数据上，这样您就可以随心所欲地进行操作，例如只需很少更改 block_diag 的源代码

arrs = a, b, c
shapes = np.array([i.shape for i in arrs])
out = np.zeros([shapes[0, 0], shapes[:, 1].sum(), shapes[:, 2].sum()])
r, c = 0, 0

for i, (_, rr, cc) in enumerate(shapes):
    out[:, r:r + rr, c:c + cc] = arrs[i]
    r += rr
    c += cc

print(np.allclose(result, out))
# True