Python 中海量 numpy 数组的内存高效排序

我需要使用 numpy 对非常大的基因组数据集进行排序。我有一个 26 亿个浮点数的数组，维度 =(868940742, 3)加载后，它会占用我机器上大约 20GB 的内存。我有一台 2015 年初的 13 英寸 MacBook Pro，配备 16GB RAM、500GB 固态硬盘和 3.1 GHz intel i7 处理器。只是加载数组会溢出到虚拟内存，但不会到我的机器受到影响的程度，或者我必须停止我正在做的所有其他事情。

我从 22 个较小的数组一步步构建这个非常大的数组(N, 2)子数组。

功能FUN_1生成 2 个新的(N, 1)使用我称之为的 22 个子数组中的每一个的数组sub_arr.

第一个输出为FUN_1是通过插入值生成的sub_arr[:,0]在阵列上b = array([X, F(X)])第二个输出是通过放置生成的sub_arr[:, 0]使用数组放入 bin 中r = array([X, BIN(X)])。我称这些输出为b_arr and rate_arr， 分别。该函数返回一个 3 元组(N, 1) arrays:

import numpy as np

def FUN_1(sub_arr):
    """interpolate b values and rates based on position in sub_arr"""

    b = np.load(bfile)
    r = np.load(rfile)

    b_arr = np.interp(sub_arr[:,0], b[:,0], b[:,1])
    rate_arr = np.searchsorted(r[:,0], sub_arr[:,0])  # HUGE efficiency gain over np.digitize...

    return r[rate_r, 1], b_arr, sub_arr[:,1] 

我在 for 循环中调用该函数 22 次，并填充预先分配的零数组full_arr = numpy.zeros([868940742, 3])与价值观：

full_arr[:,0], full_arr[:,1], full_arr[:,2] = FUN_1

就这一步节省内存而言，我认为这是我能做的最好的事情，但我愿意接受建议。不管怎样，到目前为止我都没有遇到问题，而且只需要大约 2 分钟。

这是排序例程（有两个连续排序）

for idx in range(2):
    sort_idx = numpy.argsort(full_arr[:,idx])
    full_arr = full_arr[sort_idx]
    # ...
    # <additional processing, return small (1000, 3) array of stats>

现在这种方法已经开始工作了，尽管速度很慢（大约需要 10 分钟）。然而，我最近开始使用更大、更精细的分辨率表[X, F(X)]上述插值步骤的值FUN_1返回b_arr现在排序确实变慢了，尽管其他一切都保持不变。

有趣的是，我什至没有在排序现在滞后的步骤中对插值进行排序。以下是不同插值文件的一些片段 - 较小的插值文件在每种情况下大约小 30%，并且就第二列中的值而言更加统一；较慢的具有更高的分辨率和更多的独特值，因此插值的结果可能更独特，但我不确定这是否会产生任何影响......？

更大、更慢的文件：

17399307    99.4
17493652    98.8
17570460    98.2
17575180    97.6
17577127    97
17578255    96.4
17580576    95.8
17583028    95.2
17583699    94.6
17584172    94

更小、更统一的常规文件：

1       24  
1001    24  
2001    24  
3001    24  
4001    24  
5001    24
6001    24
7001    24

我不确定是什么导致了这个问题，我对任何建议或关于在这种类型的内存限制情况下进行排序的一般输入感兴趣！

目前每次致电np.argsort正在生成一个(868940742, 1)int64 索引数组，它本身就占用约 7 GB。此外，当您使用这些索引对列进行排序时full_arr你正在生成另一个(868940742, 1)浮点数组，因为花式索引总是返回副本而不是视图 http://docs.scipy.org/doc/numpy/user/basics.indexing.html#index-arrays.

一个相当明显的改进是排序full_arr使用其到位.sort() method http://docs.scipy.org/doc/numpy/reference/generated/numpy.sort.html。很遗憾，.sort()不允许您直接指定要排序的行或列。然而，你can指定结构化数组的排序依据字段。因此，您可以通过获取一个值来强制对三列之一进行就地排序view http://docs.scipy.org/doc/numpy/reference/generated/numpy.ndarray.view.html作为具有三个浮点字段的结构化数组添加到数组上，然后按以下字段之一进行排序：

full_arr.view('f8, f8, f8').sort(order=['f0'], axis=0)

在这种情况下我正在排序full_arr位于第 0 个字段，对应于第一列。请注意，我假设有三个 float64 列（'f8') - 如果您的数据类型不同，您应该相应地更改它。这还要求您的数组是连续的并且采用行优先格式，即full_arr.flags.C_CONTIGUOUS == True.

这种方法的功劳应该归功于 Joe Kington 的回答here https://stackoverflow.com/a/2828371/1461210.

虽然它需要更少的内存，但不幸的是，与使用相比，按字段对结构化数组进行排序要慢得多np.argsort生成索引数组，正如您在下面的评论中提到的那样（请参阅上一个问题 https://stackoverflow.com/q/19682521/1461210）。如果你使用np.argsort要获取一组索引进行排序，您可能会通过使用获得适度的性能增益np.take而不是直接索引来获取排序数组：

 %%timeit -n 1 -r 100 x = np.random.randn(10000, 2); idx = x[:, 0].argsort()
x[idx]
# 1 loops, best of 100: 148 µs per loop

 %%timeit -n 1 -r 100 x = np.random.randn(10000, 2); idx = x[:, 0].argsort()
np.take(x, idx, axis=0)
# 1 loops, best of 100: 42.9 µs per loop

但是，我不希望在内存使用方面看到任何差异，因为两种方法都会生成副本。

关于为什么对第二个数组进行排序更快的问题 - 是的，当数组中唯一值较少时，您应该期望任何合理的排序算法都会更快，因为平均而言，它要做的工作较少。假设我有一个 1 到 10 之间的随机数字序列：

5  1  4  8  10  2  6  9  7  3

有10个！ = 3628800 种可能的排列这些数字的方式，但只有一种是按升序排列的。现在假设只有 5 个唯一数字：

4  4  3  2  3  1  2  5  1  5

现在有 2⁵ = 32 种方法可以按升序排列这些数字，因为我可以交换排序向量中的任何一对相同的数字而不会破坏顺序。

默认情况下，np.ndarray.sort() uses 快速排序 https://en.wikipedia.org/wiki/Quicksort. The qsort https://en.wikipedia.org/wiki/Quicksort#Repeated_elements该算法的变体通过递归地选择数组中的“枢轴”元素，然后对数组重新排序，使得所有小于枢轴值的元素都放置在它之前，所有大于枢轴值的元素都放置在它之后。等于主元的值已经排序。具有较少的唯一值意味着，平均而言，在任何给定的扫描中，更多的值将等于主元值，因此需要更少的扫描来对数组进行完全排序。

例如：

%%timeit -n 1 -r 100 x = np.random.random_integers(0, 10, 100000)
x.sort()
# 1 loops, best of 100: 2.3 ms per loop

%%timeit -n 1 -r 100 x = np.random.random_integers(0, 1000, 100000)
x.sort()
# 1 loops, best of 100: 4.62 ms per loop

在此示例中，两个数组的数据类型相同。如果较小的数组与较大的数组相比具有较小的项目大小，那么由于花哨的索引而复制它的成本也会较小。