用模对 numpy 数组求和

Question

考虑以下设置：

import numpy as np
import itertools as it
A = np.random.rand(3,3,3,16,3,3,3,16)  # sum elements of A to arrive at...
B = np.zeros((4,4))  # a 4x4 array (output)

我有一个大数组 'A' 我想求和，但要以一种非常具体的方式求和。 'A' 的形状为 (x,x,x,16,x,x,x,16)，其中 'x' 是某个整数。期望的结果是一个 4x4 矩阵 'B'，我可以像这样通过 for 循环计算它：

%%timeit
for x1,y1,z1,s1 in it.product(range(3), range(3), range(3), range(16)):
    for x2,y2,z2,s2 in it.product(range(3), range(3), range(3), range(16)):
        B[s1%4, s2%4] += A[x1,y1,z1,s1,x2,y2,z2,s2]

>> 134 ms ± 1.27 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

其中B的元素是'A'中那个维度16个元素的两个轴的"modulo-4"，这里索引为s1和s2.

如何通过广播或其他方式实现同样的效果？显然，'x' 越大（'A' 中的维度），for 循环的计算时间将呈指数增长，这并不理想。

编辑：

C = np.zeros((4,4))
for i,j in it.product(range(4), range(4)):
    C[i,j] = A[:,:,:,i::4,:,:,:,j::4].sum()

这似乎也有效。但仍然涉及1个for循环。有没有办法让它更快？

Answer 1

这里有一个更干净、更快速的解决方案。不幸的是，它们并不相同......

def clean(A):
    return A.reshape(4*n*n*n, 4, 4*n*n*n, 4).sum(axis=(0, 2))

def fast(A):
    return np.bincount(np.tile(np.arange(16).reshape(4, 4), (4, 4)).ravel(), A.sum((0,1,2,4,5,6)).ravel(), minlength=16).reshape(4, 4)

在 n==6 fast 大约快三倍。

用模对 numpy 数组求和

Summing over numpy array with modulo

python

arrays

numpy

python-3.x

array-broadcasting