python - 用于在 NumPy 中创建 block 矩阵序列的向量化语法

Question

我有两个 3D 数组 A和B有形状 (k, n, n)和(k, m, m)分别。我想创建一个矩阵 C形状(k, n+m, n+m)这样对于每个 0 <= i < k ，二维矩阵 C[i,:,:]是通过输入 A[i, :, :] 得到的 block 对角矩阵在左上角n x n部分和B[i, :, :]右下m x m部分。

目前我正在使用以下方法来实现 NumPy:

C = np.empty((k, n+m, n+m))
for i in range(k):
    C[i, ...] = np.block([[A[i,...], np.zeros((n,m))],
                          [np.zeros((m,n)), B[i,...]]])

我想知道是否有办法在没有 for 的情况下做到这一点环形。我想如果k太大了，我的解决方案效率不是很高。

Answer 1

IIUC 您可以简单地切片和分配 -

C = np.zeros((k, n+m, n+m),dtype=np.result_type(A,B))
C[:,:n,:n] = A
C[:,n:,n:] = B