numpy-groupies

本包包含一组优化的工具库，用于执行可大致认为是“分组索引操作”的任务。其中最突出的工具是 aggregate，其详细描述在页面下方。

安装

如果您有 pip，则只需

pip install numpy_groupies

请注意，numpy_groupies 没有任何强制依赖项（甚至 numpy 也是可选的），因此即使没有包管理器，您也应该能够相当容易地安装它。如果您只想安装 aggregate 的一个特定实现（例如 aggregate_numpy.py），您可以下载该文件，并将 utils.py 文件的内容复制粘贴到该文件的顶部（替换 from .utils import (...) 行）。

aggregate

import numpy as np
import numpy_groupies as npg
group_idx = np.array([   3,   0,   0,   1,   0,   3,   5,   5,   0,   4])
a =         np.array([13.2, 3.5, 3.5,-8.2, 3.0,13.4,99.2,-7.1, 0.0,53.7])
npg.aggregate(group_idx, a, func='sum', fill_value=0)
# >>>          array([10.0, -8.2, 0.0, 26.6, 53.7, 92.1])

aggregate 接收一个值数组，以及一个数组，其中包含每个值的分组编号。然后它返回每个组中值的总和（或平均值、标准差、任何...等）。您可能之前已经遇到过这个想法——参见 Matlab 的 accumarray 函数，或 pandas 分组概念，或 MapReduce 模式，或简单地是 基本直方图。

一些实现的功能并不减少数据，而是在迭代数据或排列它们时累积值。输出大小与输入大小匹配。

group_idx = np.array([4, 3, 3, 4, 4, 1, 1, 1, 7, 8, 7, 4, 3, 3, 1, 1])
a =         np.array([3, 4, 1, 3, 9, 9, 6, 7, 7, 0, 8, 2, 1, 8, 9, 8])
npg.aggregate(group_idx, a, func='cumsum')
# >>>          array([3, 4, 5, 6,15, 9,15,22, 7, 0,15,17, 6,14,31,39])

输入

该函数接受各种不同的输入组合，产生各种不同的输出形状。我们简要描述了输入的一般含义，然后更详细地介绍不同的组合。

• group_idx - 非负整数数组，用作对 a 中值进行分组的“标签”。
• a - 要聚合的值数组。
• func='sum' - 用于聚合的函数。有关更多详细信息，请参阅下面的部分。
• size=None - 输出数组的形状。如果 None，则 group_idx 中的最大值将设置输出的大小。
• fill_value=0 - 用于在 group_idx 输入数组中任何地方都没有出现的输出组中的值。
• order='C' - 对于多维输出，这控制内存中的布局，可以是 'F' 以 Fortran 风格。
• dtype=None - 输出的 dtype。None 表示为给定的 a、func 和 fill_value 选择一个合理的类型。
• axis=None - 解释如下。
• ddof=0 - 传递给方差和标准差的计算（请参阅函数部分）。

• 形式 1 是最简单的，它接收长度匹配的 1D group_idx 和 a，并产生一个 1D 输出。
• 形式 2 与形式 1 类似，但接收标量 a，该标量被广播到 group_idx 的长度。请注意，这通常并不那么有用。
• 形式 3 更复杂。 group_idx 的长度与 a.shape[axis] 相同。组被广播到 a 的其他轴/轴上，因此输出形状为 n_groups x a.shape[0] x ... x a.shape[axis-1] x a.shape[axis+1] x ... a.shape[-1]，即输出有两个或更多维度。
• 形式 4 也产生具有两个或更多维度的输出，但与形式 3 的原因非常不同。在这里，a 是 1D，group_idx 是 2D，而在形式 3 中，a 是 ND，group_idx 是 1D，我们提供了 axis 的值。 a 的长度必须匹配 group_idx.shape[1]，group_idx.shape[0] 的值确定输出中的维度数，即 group_idx[:,99] 给出 a[99] 的 (x,y,z) 组索引。
• 表单5与表单4相同，但包含标量a。与表单2一样，这很少有帮助。

性能注意事项。 输出的顺序不太可能影响聚合的性能（尽管它可能影响您对该输出的后续使用），然而多维a或group_idx的顺序可能会影响性能：在表单4中，如果group_idx中的列在内存中是连续的，则最佳，即group_idx[:, 99]对应于连续的内存块；在表单3中，如果group_idx[i]的a中的所有数据都是连续的，则最佳，例如如果axis=1，则我们希望a[:, 55]是连续的。

可用函数

默认情况下，aggregate假设您想要对每个组内的值求和，但是您可以使用func关键字参数指定另一个函数。这个func可以是任何自定义的可调用对象，但是您可能想要以下优化函数之一。请注意，并非所有实现都提供所有函数。

• 'sum' - 每个组内项的和（见上面的示例）。
• 'prod' - 每个组内项的乘积
• 'mean' - 每个组内项的平均值
• 'var' - 每个组内项的方差。使用ddof关键字参数表示自由度。计算中使用的除数是N - ddof，其中N表示元素的数量。默认情况下，ddof为零。
• 'std' - 每个组内项的标准差。使用ddof关键字参数表示自由度（见上面的var）。
• 'min' - 每个组内项的最小值。
• 'max' - 每个组内项的最大值。
• 'first' - 每个组内a中的第一个项。
• 'last' - 每个组内a中的最后一个项。
• 'argmax' - 每个组内a中最大值的索引。
• 'argmin' - 每个组内a中最小值的索引。

上述函数也有一个nan形式，它们跳过计算结果中的nan值，而不是传播到计算结果中

• 'nansum'、'nanprod'、'nanmean'、'nanvar'、'nanstd'、'nanmin'、'nanmax'、'nanfirst'、'nanlast'、'nanargmax'、'nanargmin'

以下函数与上述函数略有不同，因为它们总是返回布尔值。它们对nan的处理也与上述不同

• 'all' - 如果组内的所有项都是真值，则返回True。请注意，np.all(nan)是True，即nan实际上是真值。
• 'any' - 如果组内的任何项是真值，则返回True。
• 'allnan' - 如果组内的所有项都是nan，则返回True。
• 'anynan' - 如果组内的任何项是nan，则返回True。

以下函数不会减少数据，而是生成与输入大小匹配的输出

• 'cumsum' - 每个组内项的累积和。
• 'cumprod' - 每个组内项的累积乘积。（仅numba）
• 'cummin' - 每个组内项的累积最小值。（仅numba）
• 'cummax' - 每个组内项的累积最大值。（仅numba）
• 'sort' - 以升序对每个组内的项进行排序，使用reverse=True来反转顺序。

最后，有三个函数不会将每个组减少为单个值，而是返回组内的完整项集

• 'array' - 简单地返回分组项，使用与a中出现的相同顺序。（仅numpy）

示例

计算连续整数的和，然后计算这些连续整数的乘积。

group_idx = np.arange(5).repeat(3)
# group_idx: array([0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4])
a = np.arange(group_idx.size)
# a: array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14])
x = npg.aggregate(group_idx, a) # sum is default
# x: array([ 3, 12, 21, 30, 39])
x = npg.aggregate(group_idx, a, 'prod')
# x: array([ 0, 60, 336, 990, 2184])

忽略NaN值计算方差，将所有NaN值组设置为nan。

x = npg.aggregate(group_idx, a, func='nanvar', fill_value=nan)

计算每个组中元素的数量。请注意，这相当于执行np.bincount(group_idx)，实际上这就是NumPy实现的方式。

x = npg.aggregate(group_idx, 1)

将1000个值累加成一个15x15x15的三维立方体。注意，在这个例子中，三个维度的大小相同，但这不必是这种情况。

group_idx = np.random.randint(0, 15, size=(3, 1000))
a = np.random.random(group_idx.shape[1])
x = npg.aggregate(group_idx, a, func="sum", size=(15,15,15), order="F")
# x.shape: (15, 15, 15)
# np.isfortran(x): True

使用自定义函数生成一些字符串。

group_idx = np.array([1, 0,  1,  4,  1])
a = np.array([12.0, 3.2, -15, 88, 12.9])
x = npg.aggregate(group_idx, a,
              func=lambda g: ' or maybe '.join(str(gg) for gg in g), fill_value='')
# x: ['3.2', '12.0 or maybe -15.0 or maybe 12.9', '', '', '88.0']

使用axis参数同时进行三行的求和聚合。

a = np.array([[99, 2,  11, 14,  20],
	   	   [33, 76, 12, 100, 71],
		   [67, 10, -8, 1,   9]])
group_idx = np.array([[3, 3, 7, 0, 0]])
x = npg.aggregate(group_idx, a, axis=1)
# x : [[ 34, 0, 0, 101, 0, 0, 0, 11],
#      [171, 0, 0, 109, 0, 0, 0, 12],
#      [ 10, 0, 0,  77, 0, 0, 0, -8]]

多种实现

提供了多种aggregate的实现。如果您使用from numpy_groupies import aggregate，则将自动选择最佳可用实现。否则，您可以直接选择特定版本，例如from numpy_groupies import aggregate_nb as aggregate或从实现模块导入aggregate from numpy_groupies.aggregate_weave import aggregate。

目前有以下实现

• numpy - 这是默认实现。它使用纯numpy，主要依赖于np.bincount和基本的索引魔术。它不包含除numpy以外的其他依赖项，并且对于偶尔的使用表现合理。
• numba - 这是性能最好的实现，基于numba和LLVM提供的jit编译。
• 纯Python - 这种实现没有依赖项，仅使用标准库。它非常慢，只有在没有numpy的情况下才应该使用。
• numpy ufunc - 仅用于基准测试。这种实现使用NumPy的ufunc的.at方法（例如add.at），这似乎是为执行与aggregate相同的计算而设计的，然而NumPy实现相当不完整。
• pandas - 仅用于参考。 Pandas的groupby概念与aggregate执行的任务相同。但是，Pandas实际上并不比默认的NumPy实现快。另外，请注意，在这里使用Pandas可能有改进的空间。最值得注意的是，当计算相同数据的多个聚合（例如'min'和'max'）时，Pandas可能更有效地使用。

所有实现都具有相同的调用语法和产生相同的输出，但存在一些浮点误差。但是，某些实现仅支持有效输入的子集，有时会抛出NotImplementedError。

基准测试

此存储库中包含测试和基准测试脚本。要进行基准测试，请从此存储库的根目录运行python -m numpy_groupies.benchmarks.generic。

下面我们使用从[0, 1000)中均匀选择的500,000个索引。a的值从区间[0,1)中均匀选择，小于0.2的值设置为0（以便在布尔运算中充当假值）。对于nan-运算，另外20%的值设置为nan，剩余的值位于区间[0.2,0.8)。

基准测试结果以ms为单位给出，在2.40GHz的i7-7560U上运行

Linux(x86_64), Python 3.10.12, Numpy 1.25.2, Numba 0.58.0, Pandas 2.0.2

开发

该项目由@ml31415启动，其中numba和weave的实现也是由他完成的。纯Python和numpy的实现由@d1manson编写。

作者希望numpy的ufunc.at方法或其他numpy或scipy中的aggregate实现最终足够快，使这个包变得不再必要。实际上，Numpy 1.25确实包含了关于ufunc速度的重大改进，大大缩小了numpy和numba实现之间的速度差距。