假设我有一个形状为 (3, 10, 3) 的数组 b 和另一个数组 v = [8, 9, 4] 形状 (3,)，见下文。对于 b 中形状为 (10, 3) 的 3 个数组中的每一个，我需要对由 v 确定的行数求和，即对于 i = 0, 1, 2 我需要得到 np.sum(b[i, 0:v[i]], axis=0)。我的解决方案(如下所示)使用 for 循环，我猜这是低效的。我想知道是否有一种有效的(矢量化的)方法来完成我上面描述的事情。

注意:我的实际数组有更多维度，这些数组用于说明。

v = np.array([8,9,4]) 
b = np.array([[[0., 1., 0.], 
               [0., 0., 1.], 
               [0., 0., 1.], 
               [0., 0., 1.], 
               [1., 0., 0.], 
               [1., 0., 0.], 
               [0., 0., 1.], 
               [1., 0., 0.], 
               [0., 1., 0.], 
               [1., 0., 0.]], 
 
              [[0., 0., 1.], 
               [0., 1., 0.], 
               [1., 0., 0.], 
               [0., 0., 1.], 
               [1., 0., 0.], 
               [1., 0., 0.], 
               [1., 0., 0.], 
               [0., 1., 0.], 
               [0., 0., 1.], 
               [0., 1., 0.]], 
 
              [[1., 0., 0.], 
               [1., 0., 0.], 
               [1., 0., 0.], 
               [0., 0., 1.], 
               [0., 1., 0.], 
               [0., 1., 0.], 
               [1., 0., 0.], 
               [1., 0., 0.], 
               [0., 0., 1.], 
               [1., 0., 0.]]]) 
 
n=v.shape[0] 
vv=np.zeros([n, n]) 
for i in range(n): 
   vv[i]=np.sum( b[i,0:v[i]],axis=0)

输出:

vv 
array([[3., 1., 4.], 
       [4., 2., 3.], 
       [3., 0., 1.]])

编辑: 下面是数组 v 和 b 的更多实际示例。

v= np.random.randint(0,300, size=(32, 98,3)) 
b = np.zeros([98, 3, 300, 3]) 
for i in range(3):  
   for j in range(98): 
       b[j,i] = np.random.multinomial(1,[1./3, 1./3, 1./3], 300) 
 
v.shape 
Out[292]: (32, 98, 3) 
 
b.shape 
Out[293]: (98, 3, 300, 3)

我需要做和以前一样的事情，所以最终结果是一个形状为(32,98,3,3)的数组。请注意，我必须在每次迭代时执行上述操作，这就是我寻找高效实现的原因。

请您参考如下方法:

这是答案中提供的不同方法的性能比较:

sliced_reduce
sliced_sum
sliced_sum_numba
reduce_cumulative (原创 here )
baseline - “经典”Python for 循环(见下文)。

性能说明

sliced_reduce 将索引对的顺序从升序反转为降序，将多余元素的计算变为空操作；然而，通过这种方式，数组不会按内存布局顺序遍历，并且似乎会使该方法减慢 ~30%。
reduce_cumulative 执行许多不必要的添加操作，这取决于开始和停止索引的分布。对于开始索引全为零且停止索引均匀分布的 OP 示例，平均而言，这将是严格必要的操作次数的两倍。对于其他分布(例如非零起始索引)，与其他方法相比，这个分数可能会发生很大变化，从而降低性能。请检查您自己的情况。
[免责声明] 与所有性能评估一样，这些是提供广泛概述的粗略指南，但它们并不能避免您在特定机器上针对特定用例自行运行性能测试绝对确定选择最佳选项。

使用 OP 中的示例维度:

In [15]: np.random.seed(0) 
 
In [16]: b = np.random.randint(0, 1000, size=(98, 3, 300, 3)) 
 
In [17]: v = np.random.randint(-299, 300, size=(32, 98, 3)) 
 
In [18]: %timeit sliced_reduce(b, np.zeros_like(v), v, np.add, axis=2) 
11.3 ms ± 110 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 
 
In [19]: %timeit sliced_sum(b, np.zeros_like(v), v, axis=2) 
54.9 ms ± 153 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) 
 
In [20]: %timeit sliced_sum_numba(b, np.zeros_like(v), v, 2) 
16.3 ms ± 609 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 
 
In [21]: %timeit reduce_cumulative(b, np.zeros_like(v), v, np.add, axis=2) 
2.05 ms ± 30.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each) 
 
In [22]: %timeit baseline(b, np.zeros_like(v), v, axis=2) 
79 ms ± 625 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

基线实现:

def baseline(a, i, j, axis=None): 
    if axis is None: 
        axis = len(i.shape) 
    i = (a.shape[axis] + i) % a.shape[axis] 
    j = (a.shape[axis] + j) % a.shape[axis] 
    m = len(i.shape) - axis 
    result = np.empty(i.shape + a.shape[axis+1:], dtype=a.dtype) 
    for k in np.ndindex(i.shape): 
        result[k] = np.sum(a[k[m:] + (slice(i[k], j[k]),)], axis=0) 
    return result

性能图

除了 OP 的特定示例案例的时间安排之外，检查算法如何随数据和索引数组的大小进行扩展也很有启发性。在这里，我们可以将形状分成三个不同的部分:

索引数组的前导维度(那些不存在于数据数组中的维度)。在 OP 示例中，这是 (32,)。
索引和数据数组的公共(public)维度(从前导维度到缩减轴的维度)。在 OP 示例中，这是 (98, 3)。
要缩小的轴的大小。在 OP 示例中，这是 300。
(数据数组的尾部维度由所有算法类似地处理，因此不需要特定的缩放比例。)

因此，我们可以为三种不同的情况创建性能图:改变主要维度的大小、公共(public)维度和要减少的轴的大小。边界从 1 到 N 中选择，其中 N 是 2 的最大幂，因此所涉及的数组都没有更多超过 5,000,000 个元素(输入、索引、输出；中间数组可能更大(例如 sliced_reduce))。

代码见下。

主要维度

常用维度

降维

代码

from string import ascii_lowercase as symbols 
import numba 
import numpy as np 
import perfplot 
 
np.random.seed(0) 
 
def sliced_reduce(a, i, j, ufunc=np.add, axis=2): 
    indices = np.tile( 
        np.repeat( 
            np.arange(np.prod(a.shape[:axis])) * a.shape[axis], 
            2 
        ), 
        np.prod(i.shape[:len(i.shape) - axis]) 
    ) 
    indices[::2] += (a.shape[axis] + i.ravel()) % a.shape[axis] 
    indices[1::2] += (a.shape[axis] + j.ravel()) % a.shape[axis] 
    indices = indices.reshape(-1, 2)[::-1].ravel()  # This seems to be counter-effective, please check for your own case. 
    result = ufunc.reduceat(a.reshape(-1, *a.shape[axis+1:]), indices)[::2]  # Select only even to odd. 
    result[indices[::2] == indices[1::2]] = ufunc.reduce([]) 
    return result[::-1].reshape(*(i.shape + a.shape[axis+1:])) 
 
def sliced_sum(a, i, j, axis=2): 
    l = len(i.shape) - axis 
    m = len(i.shape) - l 
    n = len(a.shape) - axis - 1 
    leading = symbols[:l] 
    common = symbols[l:l+m] 
    summation = symbols[l+m] 
    trailing = symbols[l+m+1:l+m+1+n] 
    i = (a.shape[axis] + i) % a.shape[axis] 
    j = (a.shape[axis] + j) % a.shape[axis] 
    indices, i, j = np.broadcast_arrays(np.arange(a.shape[axis]), 
                                        np.expand_dims(i, -1), np.expand_dims(j, -1)) 
    active_elements = (i <= indices) & (indices < j) 
    return np.einsum(f'{leading + common + summation},{common + summation + trailing}->{leading + common + trailing}', 
                     active_elements, a) 
 
def sliced_sum_numba(a, i, j, axis=2): 
    i = (a.shape[axis] + i) % a.shape[axis] 
    j = (a.shape[axis] + j) % a.shape[axis] 
    m = np.prod(i.shape[:len(i.shape) - axis], dtype=int) 
    n = np.prod(i.shape[len(i.shape) - axis:], dtype=int) 
    a_flat = a.reshape(-1, *a.shape[axis:]) 
    i_flat = i.ravel() 
    j_flat = j.ravel() 
    result = np.empty((m*n,) + a.shape[axis+1:], dtype=a.dtype) 
    numba_sum(a_flat, i_flat, j_flat, m, n, result) 
    return result.reshape(*(i.shape + a.shape[axis+1:])) 
 
@numba.jit(parallel=True, nopython=True) 
def numba_sum(a, i, j, m, n, out): 
    for index in numba.prange(m*n): 
        out[index] = np.sum(a[index % n, i[index]:j[index]], axis=0) 
 
def reduce_cumulative(a, i, j, ufunc=np.add, axis=2): 
    i = (a.shape[axis] + i) % a.shape[axis] 
    j = (a.shape[axis] + j) % a.shape[axis] 
    a = np.insert(a, 0, 0, axis) 
    c = ufunc.accumulate(a, axis=axis) 
    pre = np.ix_(*(range(x) for x in i.shape)) 
    l = len(i.shape) - axis 
    return c[pre[l:] + (j,)] - c[pre[l:] + (i,)] 
 
def baseline(a, i, j, axis=2): 
    i = (a.shape[axis] + i) % a.shape[axis] 
    j = (a.shape[axis] + j) % a.shape[axis] 
    m = len(i.shape) - axis 
    result = np.empty(i.shape + a.shape[axis+1:], dtype=a.dtype) 
    for k in np.ndindex(i.shape): 
        result[k] = np.sum(a[k[m:] + (slice(i[k], j[k]),)], axis=0) 
    return result 
 
a = np.random.randint(0, 1000, size=(98, 3, 300, 3)) 
j = np.random.randint(-299, 300, size=(32, 98, 3)) 
i = np.zeros_like(j) 
check = [f(a, i, j) for f in [sliced_reduce, sliced_sum, sliced_sum_numba, reduce_cumulative, baseline]] 
assert all(np.array_equal(check[0], x) for x in check[1:]) 
 
perfplot.show( 
    # Leading dimensions: 
    # setup = lambda n: (np.random.randint(0, 1000, size=(98, 3, 300, 3)), 
    #                    np.zeros((n, 98, 3), dtype=int), 
    #                    np.random.randint(-299, 300, size=(n, 98, 3))), 
    # Common dimensions: 
    # setup = lambda n: (np.random.randint(0, 1000, size=(n, 3, 300, 3)), 
    #                    np.zeros((32, n, 3), dtype=int), 
    #                    np.random.randint(-299, 300, size=(32, n, 3))), 
    # Reduced dimension: 
    setup = lambda n: (np.random.randint(0, 1000, size=(98, 3, n, 3)), 
                       np.zeros((32, 98, 3), dtype=int), 
                       np.random.randint(-n+1, n, size=(32, 98, 3))), 
    kernels=[ 
        lambda a: sliced_reduce(*a), 
        lambda a: sliced_sum(*a), 
        lambda a: sliced_sum_numba(*a), 
        lambda a: reduce_cumulative(*a), 
        lambda a: baseline(*a), 
    ], 
    labels=['sliced_reduce', 'sliced_sum', 'sliced_sum_numba', 'reduce_cumulative', 'baseline'], 
    # n_range=[2 ** k for k in range(13)],  # Leading dimensions. 
    # n_range=[2 ** k for k in range(11)],  # Common dimensions. 
    n_range=[2 ** k for k in range(2, 13)],  # Reduced dimension. 
    # xlabel='Size of leading dimension', 
    # xlabel='Size of first common dimension (second is 3)', 
    xlabel='Size of reduced dimension', 
)

python之每次使用不同的切片范围沿不同的暗淡对数组求和

性能说明

性能图

主要维度

常用维度

降维

代码