使用 CuPy 或 NumPy 高效计算分区总和

Question

我有一个很长的数组*，长度为 L（我们称之为 values），我想对其求和，以及一个相同长度的排序一维数组 L，其中包含 N 整数对原始数组进行分区——我们称这个数组为labels。

我目前正在做的是这个（module 是 cupy 或 numpy）：

result = module.empty(N)
for i in range(N):
    result[i] = values[labels == i].sum()

但这不是最有效的方法（应该可以摆脱for 循环，但如何？）。由于labels 已排序，我可以轻松确定断点并将这些索引用作开始/停止点，但我看不出这如何解决for 循环问题。

请注意，如果可能的话，我想避免创建大小为 NxL 的数组，因为 L 非常大。

我在 cupy 工作，但也欢迎任何 numpy 解决方案，并且可能会被移植。在 cupy 中，ReductionKernel 似乎就是这种情况，但我不太明白该怎么做。

---

* 在我的例子中，values 是一维的，但我认为解决方案不会依赖于此

Answer 1

您正在描述 groupby sum 聚合。您可以为此编写一个 CuPy RawKernel，但使用在 GPU 数据帧库 cuDF 中实现的现有 groupby 聚合会非常更容易。它们可以互操作，而无需您复制数据。如果您在生成的 cuDF 系列上调用 .values，它将为您提供一个 CuPy 数组。

如果你回到 CPU，你可以用 pandas 做同样的事情。

import cupy as cp
import pandas as pd

N = 100
values = cp.random.randint(0, N, 1000)
labels = cp.sort(cp.random.randint(0, N, 1000))

L = len(values)
result = cp.empty(L)
for i in range(N):
    result[i] = values[labels == i].sum()
    
result[:5]
array([547., 454., 402., 601., 668.])

import cudf

df = cudf.DataFrame({"values": values, "labels": labels})
df.groupby(["labels"])["values"].sum().values[:5]
array([547, 454, 402, 601, 668])

Answer 2

这是一个解决方案，它使用N x <max partition size in labels> 数组而不是N x L 数组（如果不同分区之间的差异不太大，则该数组不应该很大）：

1. 将数组调整为二维数组，每行都有分区。由于行的长度等于最大分区的大小，因此用零填充不可用的值（因为它不会影响任何总和）。这使用了@Divakar 给出的here 的解决方案。

def jagged_to_regular(a, parts):
    lens = np.ediff1d(parts,to_begin=parts[0])
    mask = lens[:,None]>np.arange(lens.max())
    out = np.zeros(mask.shape, dtype=a.dtype)
    out[mask] = a
    return out

parts_stack = jagged_to_regular(values, labels)

1. 沿轴 1 求和：

result = np.sum(parts_stack, axis = 1)

如果您想要一个 CuPy 实现，jagged_to_regular 中没有直接的 CuPy alternative 到 numpy.ediff1d。在这种情况下，您可以将语句替换为 numpy.diff，如下所示：

lens = np.insert(np.diff(parts), 0, parts[0])

然后继续使用 CuPy 作为 numpy 的替代品。