python - 如何对数组中的范围值求和，其中每行都有不同的范围

Question

我想计算数组中某个范围的总和(简单) - 但我不想这样做，但要计算 n 次，并且应该求和的范围来自第二个数组。

我有一个包含 0 和 1 的二维数组:

count = np.array(\
[[0,1,0,0,1,0,1],
 [0,0,1,1,1,0,0]])

我有一个结构化的二维数组，它有一个字段，其中的范围是对计数数组求和的范围。

dtype=[..., ('ranges', 'u1', (2, 2)) , ...]

table['ranges'] 如下所示:

 [
 [[1, 3], [0, 4]],
 [[0, 0], [3, 4]],
 [[0, 0], [2 4]],
 [[0, 0], [3 4]],
 [[3, 7], [1 5]]]

(通常在 20 到几百行之间)。

这个例子的结果应该是

[2, # = (1 +0) + (0 + 0 +1)
 1, # = ( ) + (1)
 2,  # = ( ) + (1 + 1)
 1, # = ( ) + (1)
 5] # = (0 + 1 +0 +1 ) + (0 + 1 + 1 + 1)

首先我开始:

        result = np.zeros(table.size, dtype=np.int)

        for index, r in enumerate(table):
            for index, range in enumerate(r['ranges']):
                result[index] += np.sum(counts[index][range[0]:range[1]])

给出了正确的结果，但不是效率的例子。

我还尝试消除第二个循环并对其进行更多的 numpyfy:

        result = np.zeros(table.size, dtype=np.int)

        for index, (from1, to1, from2, to2) in \
                enumerate(np.nditer(table['ranges'], flags=['external_loop'])):
            counts[index] += np.sum(counts[0][from1:to1]) +\
                np.sum(counts[1][from2:to2])

但是这些代码行仍然是应用程序花费大部分时间的地方。该应用程序比这要大得多，但根据分析器，这些行花费的时间大约减半。

所以基本上我正在寻找一种方法来摆脱循环并在 numpy 中完成这一切。我一直在寻找类似的东西

counts=np.sum(counts[1][table['ranges'][0][0]:table['ranges'][0][1])+np.sum(counts[2][table['ranges'][1][0]:table['ranges'][1][1])

但到目前为止还没有真正找到一个好的方法来做到这一点。

更新进行了一些时间比较:

import numpy as np
import timeit as ti

table = np.empty(5,
                 dtype=[('s1', np.int8),
                        ('ranges', 'u1', (2, 2)),
                        ('s2', np.int16)])

table["ranges"] = [((1, 3), (0, 4)),
                   ((0, 0), (3, 4)),
                   ((0, 0), (2, 4)),
                   ((0, 0), (3, 4)),
                   ((3, 7), (1, 5))]

results = np.zeros(table.size)

counts = np.array([[0, 1, 0, 0, 1, 0, 1],
                   [0, 0, 1, 1, 1, 0, 0]])


# version one
def rv1(table, counts, results):
    for row_index, r in enumerate(table):
        for index, crange in enumerate(r['ranges']):
            results[row_index] += np.sum(counts[index][crange[0]:crange[1]])


# version two
def rv2(table, counts, results):
    for rowindex, (f1, t1, f2, t2) in \
            enumerate(np.nditer(table['ranges'], flags=['external_loop'])):
        results[rowindex] += np.sum(counts[0][f1:t1]) +\
            np.sum(counts[1][f2:t2])


# version 3 (TomNash)
def rvTN(table, counts, results):
    ranges=table["ranges"]
    result=[
        sum(counts[0][slice(*ranges[i][0])]) + sum(counts[1][slice(*ranges[i][1])])
            for i in range(len(ranges))]
    results+=result


results = np.zeros(table.size)
rv1(table, counts, results)
print ("rv1 result" , results)


results = np.zeros(table.size)
rv2(table, counts, results)
print ("rv2 result", results)

results = np.zeros(table.size)
rvTN(table, counts, results)
print ("slice*(TN) result", results)



print ("double loop time " , ti.timeit(lambda : rv1(table, counts, results)))
print ("nditer time " ,  ti.timeit(lambda : rv2(table, counts, results)))
print ("slice* time " ,  ti.timeit(lambda : rv3(table, counts, results)))

我明白了

double loop result [3. 1. 2. 1. 5.]
nditer result [3. 1. 2. 1. 5.]
slice* result [3. 1. 2. 1. 5.]
double loop time  42.41987561201677
nditer time  36.45269059110433
slice* time  24.102186055853963

所以 TomNashs 版本快了大约 30%。不幸的是，这仍然有点慢。

Answer 1

遇到同样的问题后，我找到了两种编写完全矢量化版本的方法。

一种解决方案涉及使用 ufunc.reduceat numpy.add 上的方法，但由于 reduceat 行为的奇怪规范，这最终变得相当困惑和低效。关于 numpy issue #834 有一些讨论关于添加一个新的 ufunc 方法，可以使该任务成为一个高效的单行任务，但尚未实现任何内容。

我提出的第二个解决方案似乎更高效，也更容易理解，工作原理如下:

• 首先，使用numpy.cumsum获取沿 counts 长轴的累积总和 cumulative_counts。

• 接下来，获取要求和的范围的所有起始索引 starts 和结束索引 ends，这可以通过对数组进行切片轻松获得在tables['ranges']中。

• 现在，您可以通过使用 starts 和 ends 索引到 cumulative_counts 来获取每个范围的开始和结束处的累积总和的数组.

• 最后，要获得每个范围内的总和，只需从最终总和中减去所有起始总和即可。

有一个小问题，那就是为了使索引正常工作，cumulative_counts 需要在其第一个索引处为零，然后才是cumsum > 结果。

首先将这个概念放在一维中，因为它稍微清晰一点:

for i in range(len(counts)):
    cumulative_counts = np.empty(len(counts[i]) + 1, dtype=counts.dtype)
    cumulative_counts[0] = 0
    cumulative_counts[1:] = np.cumsum(counts[i])
    starts, ends = table['ranges'][:,i,:].T
    results += cumulative_counts[ends] - cumulative_counts[starts]

这已经更快了 - 循环仅在短轴上(计数的 2 行)，其余部分被矢量化。

但也可以消除外循环并对 2D 输入应用相同的方法，为您的示例提供完全矢量化的解决方案:

cumulative_counts = np.empty((counts.shape[0], counts.shape[1] + 1), dtype=counts.dtype)
cumulative_counts[:,0] = 0
cumulative_counts[:,1:] = np.cumsum(counts, axis=1)
starts, ends = table['ranges'].T.swapaxes(1, 2)
rows = np.arange(len(counts)).reshape(1,-1)
row_results = cumulative_counts[rows, ends] - cumulative_counts[rows, starts]
results += row_results.sum(axis=1)

将此版本添加到您的计时比较中，它在我的系统上比 TomNash 的切片方法快 30%。我预计对于更大的阵列，加速会得到放大。

rv1 result [3. 1. 2. 1. 5.]
rv2 result [3. 1. 2. 1. 5.]
slice*(TN) result [3. 1. 2. 1. 5.]
cumsum(ML) result [3. 1. 2. 1. 5.]
double loop time  118.32055400998797
nditer time  101.83165721700061
slice* time  50.249228183995
cumsum time  38.694901300012134