python - 从 numpy 数组中提取对角 block

Question

我正在寻找一种巧妙的方法来提取位于 (2N)x(2N) numpy 数组主对角线上的大小为 2x2 的对角线 block (也就是说，将有 N 个这样的 block )。这概括了 numpy.diag，它返回沿主对角线的元素，人们可能认为这些元素是 1x1 block (当然 numpy 不以这种方式表示它们)。

更广泛地说，人们可能希望从 (MN)x(MN) 数组中提取 N MxM block 。这似乎是 scipy.linalg.block_diag 的补充，在 How can I transform blocks into a blockdiagonal matrix (NumPy) 中进行了巧妙的讨论。 ，将 block 从 block_diag 放置它们的地方拉出来。

从the solution to that question借用代码，这是如何设置的:

>>> a1 = np.array([[1,1,1],[1,1,1],[1,1,1]])
>>> a2 = np.array([[2,2,2],[2,2,2],[2,2,2]])
>>> a3 = np.array([[3,3,3],[3,3,3],[3,3,3]])
>>> import scipy.linalg
>>> scipy.linalg.block_diag(a1, a2, a3)
array([[1, 1, 1, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 0, 0, 0, 0, 0, 0],
       [1, 1, 1, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 2, 2, 2, 0, 0, 0],
       [0, 0, 0, 2, 2, 2, 0, 0, 0],
       [0, 0, 0, 2, 2, 2, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 3, 3, 3],
       [0, 0, 0, 0, 0, 0, 3, 3, 3],
       [0, 0, 0, 0, 0, 0, 3, 3, 3]]) 

那么，我们希望有一个像这样的函数

>>> A = scipy.linalg.block_diag(a1, a2, a3)
>>> extract_block_diag(A, M=3)
array([[[1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]],
       [[2, 2, 2],
        [2, 2, 2],
        [2, 2, 2]],
       [[3, 3, 3],
        [3, 3, 3],
        [3, 3, 3]]]) 

继续与 numpy.diag 进行类比，人们可能还希望提取非对角线 block :第 k 个 block 对角线上的 N - k 个 block 。 (顺便说一下，block_diag 的扩展允许将 block 放置在主对角线之外肯定会有用，但这不是这个问题的范围。)在上面的数组的情况下，这可能会产生:

>>> extract_block_diag(A, M=3, k=1)
array([[[0, 0, 0],
        [0, 0, 0],
        [0, 0, 0]],
       [[0, 0, 0],
        [0, 0, 0],
        [0, 0, 0]]])

我看到 this question 中涵盖了 stride_tricks 的使用旨在产生类似于此功能的东西，但我理解跨步在字节级别运行，这听起来不太合适。

出于动机，这源于我希望提取协方差矩阵(即方差)的对角线元素的情况，其中元素本身不是标量而是 2x2 矩阵。

编辑:基于the suggestion from Chris ，我做了以下尝试:

def extract_block_diag(A,M,k=0):
    """Extracts blocks of size M from the kth diagonal
    of square matrix A, whose size must be a multiple of M."""

    # Check that the matrix can be block divided
    if A.shape[0] != A.shape[1] or A.shape[0] % M != 0:
        raise StandardError('Matrix must be square and a multiple of block size')

    # Assign indices for offset from main diagonal
    if abs(k) > M - 1:
        raise StandardError('kth diagonal does not exist in matrix')
    elif k > 0:
        ro = 0
        co = abs(k)*M 
    elif k < 0:
        ro = abs(k)*M
        co = 0
    else:
        ro = 0
        co = 0

    blocks = np.array([A[i+ro:i+ro+M,i+co:i+co+M] 
                       for i in range(0,len(A)-abs(k)*M,M)])
    return blocks

上面的数据会返回以下结果:

D = extract_block_diag(A,3)
[[[1 1 1]
  [1 1 1]
  [1 1 1]]

 [[2 2 2]
  [2 2 2]
  [2 2 2]]

 [[3 3 3]
  [3 3 3]
  [3 3 3]]]

D = extract_block_diag(A,3,-1)
[[[0 0 0]
  [0 0 0]
  [0 0 0]]

 [[0 0 0]
  [0 0 0]
  [0 0 0]]]

Answer 1

作为起点，您可以使用类似的东西

>>> a
array([[1, 1, 1, 0, 0, 0, 0, 0, 0],
   [1, 1, 1, 0, 0, 0, 0, 0, 0],
   [1, 1, 1, 0, 0, 0, 0, 0, 0],
   [0, 0, 0, 2, 2, 2, 0, 0, 0],
   [0, 0, 0, 2, 2, 2, 0, 0, 0],
   [0, 0, 0, 2, 2, 2, 0, 0, 0],
   [0, 0, 0, 0, 0, 0, 3, 3, 3],
   [0, 0, 0, 0, 0, 0, 3, 3, 3],
   [0, 0, 0, 0, 0, 0, 3, 3, 3]]) 

>>> M = 3
>>> [a[i*M:(i+1)*M,i*M:(i+1)*M] for i in range(a.shape[0]/M)]
[array([[1, 1, 1],
   [1, 1, 1],
   [1, 1, 1]]), array([[2, 2, 2],
   [2, 2, 2],
   [2, 2, 2]]), array([[3, 3, 3],
   [3, 3, 3],
   [3, 3, 3]])]