從一個numpy數組中提取對角塊

Question

我正在尋找一種簡潔的方式來提取大小爲2x2的對角塊，它位於（2N）x（2N）numpy數組的主對角線上（也就是說，將會有N這樣的塊）。這概括了numpy.diag，它返回沿着主對角線的元素，人們可能認爲它是1x1塊（儘管當然numpy並不代表它們）。爲了更廣泛地描述這一點，人們可能希望從（MN）x（MN）陣列中提取N個M×M個塊。這似乎是scipy.linalg.block_diag的補充，在How can I transform blocks into a blockdiagonal matrix (NumPy)中進行了巧妙的討論，將塊放在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，人們也可能希望提取非對角塊：在第k個塊對角線上它們的N-k個。 （順便說一下，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旨在產生類似於此功能的東西，但我知道striding在字節級操作，這聽起來不太合適。

通過動機的方式，這是由於我希望提取協方差矩陣（即方差）的對角元素的情況，其中元素本身不是標量而是2×2矩陣。

編輯：基於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]])] 

Answer 2

任何特別的原因，你不想使用簡單的方法？

>>> A=np.array([[1,1,1,1],[2,2,2,2],[3,3,3,3],[4,4,4,4]]) 
>>> M1s,M2s=0,2 # start from A[M1s,M2s] 
>>> M=2 # extract an MxM block 
>>> for a in A[M1s:M1s+M]: 
... print a[M2s:M2s+M] 
... 
[1 1] 
[2 2] 
>>> 

Answer 3

您還可以在欣賞做到這一點。這可能比索引方法更快。

import numpy as np 
import scipy.linalg 

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]]) 

b = scipy.linalg.block_diag(a1, a2, a3) 
b[1,4] = 4 
b[1,7] = 5 
b[4,1] = 6 
b[4,7] = 7 
b[7,1] = 8 
b[7,4] = 9 
print b 

def extract_block_diag(a, n, k=0): 
    a = np.asarray(a) 
    if a.ndim != 2: 
     raise ValueError("Only 2-D arrays handled") 
    if not (n > 0): 
     raise ValueError("Must have n >= 0") 

    if k > 0: 
     a = a[:,n*k:] 
    else: 
     a = a[-n*k:] 

    n_blocks = min(a.shape[0]//n, a.shape[1]//n) 

    new_shape = (n_blocks, n, n) 
    new_strides = (n*a.strides[0] + n*a.strides[1], 
        a.strides[0], a.strides[1]) 

    return np.lib.stride_tricks.as_strided(a, new_shape, new_strides) 

print "-- Diagonal blocks:" 
c = extract_block_diag(b, 3) 
print c 

print "-- They're views!" 
c[0,1,2] = 9 
print b[1,2] 

print "-- 1st superdiagonal" 
c = extract_block_diag(b, 3, 1) 
print c 

print "-- 2nd superdiagonal" 
c = extract_block_diag(b, 3, 2) 
print c 

print "-- 3rd superdiagonal (empty)" 
c = extract_block_diag(b, 3, 3) 
print c 

print "-- 1st subdiagonal" 
c = extract_block_diag(b, 3, -1) 
print c 

print "-- 2nd subdiagonal" 
c = extract_block_diag(b, 3, -2) 
print c 

Answer 4

，我得到以下基於unutbu對https://stackoverflow.com/a/8070716/219229答案：（順便說一句什麼錯在字節級操作？）

from numpy.lib.stride_tricks import as_strided 

... 

def extract_block_diag(A, M=3, k=0): 
    ny, nx = A.shape 
    ndiags = min(map(lambda x: x//M, A.shape)) 
    offsets = (nx*M+M, nx, 1) 
    strides = map(lambda x:x*A.itemsize, offsets) 
    if k > 0: 
     B = A[:,k*M] 
     ndiags = min(nx//M-k, ny//M) 
    else: 
     k = -k 
     B = A[k*M] 
     ndiags = min(nx//M, ny//M-k) 
    return as_strided(B, shape=(ndiags,M,M), 
        strides=((nx*M+M)*A.itemsize, nx*A.itemsize, A.itemsize)) 

---

實例：

# a1, a2, a3 from your example 
>>> bigA = scipy.linalg.block_diag(a1, a2, a3) 
>>> extract_block_diag (bigA, 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]]]) 
>>> extract_block_diag (bigA, 3)[2] 
array([[3, 3, 3], 
     [3, 3, 3], 
     [3, 3, 3]]) 
>>> extract_block_diag(np.arange(1,9*9+1).reshape(9,9),3,1) 
[[[ 4 5 6] 
[13 14 15] 
[22 23 24]] 

[[34 35 36] 
[43 44 45] 
[52 53 54]]] 

---

注意上面的函數返回一個視圖，其中，如果你改變陣列的返回的數組內的話，就是原始也受到影響。如有必要請複印一份。

Answer 5

進口numpy的作爲NP

DEF extract_blocks（數組）：

prev = -1 
for i in xrange(len(array)-1): 
    if array[i+1][i] == 0 and array[i][i+1] == 0: 
     yield array[prev + 1: i + 1, prev + 1: i + 1] 
     print prev + 1, i 
     prev = i 
yield array[prev + 1: len(array), prev + 1: len(array)] 

d = np.array（[[4,4，0，0，0，0，0，0，0 ]，[4,4,0,0,0,0,0,0,0]，[0,0,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,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]]）

for h in extract_blocks（d）：

print h 

[[4 4] [4 4]]，[[1]]，[[2 2 2] [2 2 2] [2 2 2]]，[[3 3 3] [3 3 3] [3 3 3]