從 3D 矩陣堆棧在 numpy/scipy 中構造 3D 塊對角矩陣堆棧的有效方法

Question

我正在嘗試從給定的矩陣堆棧（nXmXm）中以 nXMXM 的形式在 numpy/scipy 中構造一個塊對角矩陣堆棧，其中 M=k*m，k 是矩陣堆棧的數量。 目前，我在 for 循環中使用 scipy.linalg.block_diag function 來執行此任務：

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 數組成塊對角矩陣，但這並不能真正解決我的問題。 我能想象的任何事情都是將每個矩陣堆棧轉換為塊對角線

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 解決方案非常好用且簡單，但我m 現在使用更高維度的 arrays，我的矩陣位於最后兩個維度。 此外，我的矩陣不再具有相同的維度（主要是 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

如果輸入 arrays 的形狀正確（即要廣播的尺寸不會自動添加），它也適用於數組廣播。 感謝 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