如何计算numpy/pandas的“submatrix”条目之和?

对NumPy的一种方法是:

import numpy as np

def as_submatrices(x, rows, cols=None, writeable=False):
    from numpy.lib.stride_tricks import as_strided
    if cols is None: cols = rows
    x = np.asarray(x)
    x_rows, x_cols = x.shape
    s1, s2 = x.strides
    if x_rows % rows != 0 or x_cols % cols != 0:
        raise ValueError('Invalid dimensions.')
    out_shape = (x_rows // rows, x_cols // cols, rows, cols)
    out_strides = (s1 * rows, s2 * cols, s1, s2)
    return as_strided(x, out_shape, out_strides, writeable=writeable)

def sum_submatrices(x, rows, cols=None):
    if cols is None: cols = rows
    x = np.asarray(x)
    x_sub = as_submatrices(x, rows, cols)
    x_sum = np.sum(x_sub, axis=(2, 3))
    x_rows, x_cols = x.shape
    return np.repeat(np.repeat(x_sum, rows, axis=0), cols, axis=1)

x = np.arange(64).reshape((8, 8))

print(sum_submatrices(x, 4))
# [[216 216 216 216 280 280 280 280]
#  [216 216 216 216 280 280 280 280]
#  [216 216 216 216 280 280 280 280]
#  [216 216 216 216 280 280 280 280]
#  [728 728 728 728 792 792 792 792]
#  [728 728 728 728 792 792 792 792]
#  [728 728 728 728 792 792 792 792]
#  [728 728 728 728 792 792 792 792]]

print(sum_submatrices(x, 2))
# [[ 18  18  26  26  34  34  42  42]
#  [ 18  18  26  26  34  34  42  42]
#  [ 82  82  90  90  98  98 106 106]
#  [ 82  82  90  90  98  98 106 106]
#  [146 146 154 154 162 162 170 170]
#  [146 146 154 154 162 162 170 170]
#  [210 210 218 218 226 226 234 234]
#  [210 210 218 218 226 226 234 234]]

print(sum_submatrices(x, 2, 8))
# [[120 120 120 120 120 120 120 120]
#  [120 120 120 120 120 120 120 120]
#  [376 376 376 376 376 376 376 376]
#  [376 376 376 376 376 376 376 376]
#  [632 632 632 632 632 632 632 632]
#  [632 632 632 632 632 632 632 632]
#  [888 888 888 888 888 888 888 888]
#  [888 888 888 888 888 888 888 888]]

As pointed out by Divakar , np.broadcast_to 快一点 np.repeat

def sum_submatrices(x, rows, cols=None):
    if cols is None: cols = rows
    x = np.asarray(x)
    x_sub = as_submatrices(x, rows, cols)
    x_sum = np.sum(x_sub, axis=(2, 3), keepdims=True)
    x_sum = np.broadcast_to(x_sum, x_sub.shape)
    return x_sum.transpose((0, 2, 1, 3)).reshape(x.shape)

这和Divakar的答案基本上是一样的,只有一个更好,因为它不使用跨步技巧和转置。

einsum 和 np.repeat (唯一的条件是 n

def sum_chunks(n, x):
    """
    Tiles an array into NxN chunks, based on the sum of the chunk
    :param n: dimension of sub-matrices
    :param x: input array
    :return: Tiled array
    """
    h, w = x.shape
    out = x.reshape(h//n, n, -1, n).swapaxes(1,2).reshape(-1, n, n)
    s = np.einsum('ijk->i', out)
    return np.repeat(np.repeat(s.reshape(h//n, w//n), n, axis=0), n, axis=1)

您可以使用此解决方案将阵列拆分为任意大小的子阵列,求和,然后重复原始大小:

>>> sum_chunks(4, np.arange(64).reshape(8,8))
array([[216, 216, 216, 216, 280, 280, 280, 280],
       [216, 216, 216, 216, 280, 280, 280, 280],
       [216, 216, 216, 216, 280, 280, 280, 280],
       [216, 216, 216, 216, 280, 280, 280, 280],
       [728, 728, 728, 728, 792, 792, 792, 792],
       [728, 728, 728, 728, 792, 792, 792, 792],
       [728, 728, 728, 728, 792, 792, 792, 792],
       [728, 728, 728, 728, 792, 792, 792, 792]])

sum_chunks(2, np.arange(64).reshape(8,8))
array([[ 18,  18,  26,  26,  34,  34,  42,  42],
       [ 18,  18,  26,  26,  34,  34,  42,  42],
       [ 82,  82,  90,  90,  98,  98, 106, 106],
       [ 82,  82,  90,  90,  98,  98, 106, 106],
       [146, 146, 154, 154, 162, 162, 170, 170],
       [146, 146, 154, 154, 162, 162, 170, 170],
       [210, 210, 218, 218, 226, 226, 234, 234],
       [210, 210, 218, 218, 226, 226, 234, 234]])

这里有一些考虑到性能的东西可以利用 np.broadcast_to 总结整形后做复制部分-

def sum_chunks_broadcasted(x, M, N): # M,N : no. of blocks along height and width
    m,n = x.shape
    s = x.reshape(M,m//M,N,n//N).sum((1,3),keepdims=1)
    return np.broadcast_to(s,(M,m//M,N,n//N)).reshape(m,n)

样本运行-

In [143]: x = np.arange(48).reshape(8,6)

In [144]: x
Out[144]: 
array([[ 0,  1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17],
       [18, 19, 20, 21, 22, 23],
       [24, 25, 26, 27, 28, 29],
       [30, 31, 32, 33, 34, 35],
       [36, 37, 38, 39, 40, 41],
       [42, 43, 44, 45, 46, 47]])

In [145]: sum_chunks_broadcasted(x, M=2, N=3) # 2x3 total windows
Out[145]: 
array([[ 76,  76,  92,  92, 108, 108],
       [ 76,  76,  92,  92, 108, 108],
       [ 76,  76,  92,  92, 108, 108],
       [ 76,  76,  92,  92, 108, 108],
       [268, 268, 284, 284, 300, 300],
       [268, 268, 284, 284, 300, 300],
       [268, 268, 284, 284, 300, 300],
       [268, 268, 284, 284, 300, 300]])

与其他一般矢量化方法的比较 @user3483203's sum_chunks 和 @jdehesa's sum_submatrices 关于各种窗口形状和数字的大数组-

In [83]: x = np.random.rand(8000, 8000)

2) 4x4个窗口:

In [152]: %timeit sum_submatrices(x, 8000//4, cols=8000//4)
1 loop, best of 3: 271 ms per loop

In [153]: %timeit sum_chunks(8000//4, x)
1 loop, best of 3: 372 ms per loop

In [154]: %timeit sum_chunks_broadcasted(x, M=4, N=4)
10 loops, best of 3: 81 ms per loop

In [155]: %timeit sum_submatrices(x, 8000//40, cols=8000//40)
1 loop, best of 3: 271 ms per loop

In [156]: %timeit sum_chunks(8000//40, x)
1 loop, best of 3: 385 ms per loop

In [157]: %timeit sum_chunks_broadcasted(x, M=40, N=40)
10 loops, best of 3: 84 ms per loop

4) 总共400x400个窗口:

In [158]: %timeit sum_submatrices(x, 8000//400, cols=8000//400)
1 loop, best of 3: 318 ms per loop

In [159]: %timeit sum_chunks(8000//400, x)
1 loop, best of 3: 396 ms per loop

In [160]: %timeit sum_chunks_broadcasted(x, M=400, N=400)
10 loops, best of 3: 123 ms per loop

dx = 8
dy = 8
x_subs = 2
y_subs = 2

arr = np.arange(dx * dy).reshape(dy, dx)
sums = [
    [second_split.sum() for second_split in np.split(first_split, y_subs, axis=1)]
    for first_split in np.split(arr, x_subs, axis=0)
]
sums_filled = np.repeat(np.repeat(sums, dx, axis=0), dy, axis=1)

jdehesa的。