从 numpy 数组中提取对角线块
Posted
技术标签:
【中文标题】从 numpy 数组中提取对角线块【英文标题】：Extracting diagonal blocks from a numpy array 【发布时间】：2012-06-05 13:57:15 【问题描述】：我正在寻找一种巧妙的方法来提取位于 (2N)x(2N) numpy 数组的主对角线上的大小为 2x2 的对角块（也就是说，将有 N 个这样的块）。这概括了 numpy.diag，它返回沿主对角线的元素，人们可能会将其视为 1x1 块（当然 numpy 不会以这种方式表示它们）。
为了更广泛地表达这一点，人们可能希望从 (MN)x(MN) 数组中提取 N MxM 块。这似乎是 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 的扩展允许将块放置在主对角线之外肯定会很有用，但这不是这个问题的范围。）对于上面的数组，这可能会产生：
>>> 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]]]
【问题讨论】：
【参考方案1】：您不想使用直接方法的任何特定原因？
>>> 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]
>>>
【讨论】：
【参考方案2】：作为起点，您可以使用类似的东西
>>> 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]])]
【讨论】：
这对我来说非常好。我认为使用numpy.array() 围绕列表理解的调用，输出将是一个 (a.shape[0]/M,M,M) 数组。并且索引算法可以推广到非对角块。太好了！
很高兴我能帮上忙。但是请注意，我还没有真正考虑过当M 没有整齐地划分为矩阵形状时会发生什么。如果您将其包装在 extract_block_diag 函数中，您很可能需要进行一些输入检查。【参考方案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
【讨论】：
【参考方案4】：根据 unutbu 在 https://***.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]]]
请注意，上面的函数返回一个视图，如果您更改返回的数组数组中的任何内容，原始的也会受到影响。如有必要，请制作一份副本。
【讨论】：
【参考方案5】：将 numpy 导入为 np
def extract_blocks(array):
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, 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]])
对于extract_blocks(d)中的h：
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]]
【讨论】：
使用 extract_blocks(array) 吐出任何大小的所有对角线块。以上是关于从 numpy 数组中提取对角线块的主要内容，如果未能解决你的问题，请参考以下文章
python使用numpy中的diagonal函数获取2D numpy数组的对角线元素使用numpy中的diagonal函数和sum函数获取2D numpy数组的迹（matrix trace）