Python中内存高效的大量numpy数组

Question

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

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

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

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

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 分钟）。然而，我最近开始在FUN_1 中使用更大、更精细的[X, F(X)] 值的分辨率表作为上述插值步骤，返回b_arr，现在SORT 真的变慢了，尽管其他一切都保持不变。

有趣的是，我什至没有在排序滞后的那一步对插值进行排序。以下是不同插值文件的一些 sn-ps - 在每种情况下，较小的文件大约小 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

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

Answer 1

目前对np.argsort 的每次调用都会生成一个(868940742, 1) int64 索引数组，它本身将占用大约7 GB。此外，当您使用这些索引对full_arr 的列进行排序时，您将生成另一个(868940742, 1) 浮点数组，因为fancy indexing always returns a copy rather than a view。

一个相当明显的改进是使用.sort() method 对full_arr 进行适当的排序。不幸的是，.sort() 不允许您直接指定要排序的行或列。但是，您可以为结构化数组指定一个作为排序依据的字段。因此，您可以通过将view 作为具有三个浮点字段的结构化数组添加到您的数组中，然后按以下字段之一进行排序，从而强制对三列之一进行就地排序：

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

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

此方法的功劳应归功于 Joe Kington 的回答 here。

---

虽然它需要更少的内存，但不幸的是，与使用np.argsort 生成索引数组相比，按字段对结构化数组进行排序要慢得多，正如您在下面的 cmets 中提到的（请参阅this previous question）。如果您使用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() 使用 Quicksort。该算法的qsort 变体通过递归选择数组中的“枢轴”元素，然后对数组重新排序，使得所有小于枢轴值的元素都放在它之前，所有大于枢轴值的元素被放置在它之后。等于枢轴的值已经排序。具有更少的唯一值意味着，平均而言，更多的值将等于任何给定扫描的主值，因此需要更少的扫描来对数组进行完全排序。

例如：

%%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

在这个例子中，两个数组的 dtypes 是相同的。如果较小的数组与较大的数组相比具有较小的项大小，那么由于花哨的索引而复制它的成本也会更小。

Answer 2

编辑：如果任何编程新手和numpy 遇到此帖子，我想指出考虑您正在使用的np.dtype 的重要性。就我而言，我实际上能够摆脱使用半精度浮点，即np.float16，它将内存中的 20GB 对象减少到 5GB，并使排序更易于管理。 numpy 使用的默认值是np.float64，这是您可能不需要的很多精度。在这里查看doc，它描述了不同数据类型的容量。感谢 @ali_m 在 cmets 中指出这一点。

我在解释这个问题时做得不好，但我发现了一些有用的解决方法，我认为这些解决方法对于需要对真正庞大的 numpy 数组进行排序的任何人来说都是有用的。

我正在从 22 个包含元素 [position, value] 的人类基因组数据“子阵列”构建一个非常大的 numpy 阵列。最终，最终数组必须根据特定列中的值“就地”进行数字排序，而不是对行内的值进行混排。

子数组维度如下：

arr1.shape = (N1, 2)
...
arr22.shape = (N22, 2)

sum([N1..N2]) = 868940742 即有接近 1BN 个位置需要排序。

首先，我使用函数process_sub_arrs 处理这22 个子数组，该函数返回与输入长度相同的一维数组的三元组。我将一维数组堆叠到一个新的 (N, 3) 数组中，并将它们插入到为完整数据集初始化的 np.zeros 数组中：

    full_arr = np.zeros([868940742, 3])
    i, j = 0, 0

    for arr in list(arr1..arr22):  
        # indices (i, j) incremented at each loop based on sub-array size
        j += len(arr)
        full_arr[i:j, :] = np.column_stack( process_sub_arrs(arr) )
        i = j

    return full_arr

编辑：由于我意识到我的数据集可以用半精度浮点数表示，我现在初始化 full_arr 如下：full_arr = np.zeros([868940742, 3], dtype=np.float16)，它的大小只有 1/4，而且更容易排序。

结果是一个巨大的 20GB 数组：

full_arr.nbytes = 20854577808

正如@ali_m 在他的详细帖子中指出的那样，我之前的例程效率低下：

sort_idx = np.argsort(full_arr[:,idx])
full_arr = full_arr[sort_idx]

数组sort_idx，其大小是full_arr 的33%，在对full_arr 排序后挂起并浪费内存。由于“花哨”的索引，这种排序可能会生成full_arr 的副本，这可能会将内存使用量推高到已经用于保存大量数组的内存的 233%！这是一个缓慢的步骤，持续大约十分钟，并且严重依赖虚拟内存。

我不确定“花式”排序是否会生成持久副本。观察我机器上的内存使用情况，full_arr = full_arr[sort_idx] 似乎删除了对未排序原始的引用，因为大约 1 秒后剩下的就是排序数组和索引使用的内存，即使存在临时副本.

argsort() 用于节省内存的更简洁用法是：

    full_arr = full_arr[full_arr[:,idx].argsort()]

这仍然会导致分配时出现峰值，其中同时创建了临时索引数组和临时副本，但内存几乎立即再次释放。

@ali_m 指出了一个很好的技巧（归功于 Joe Kington），它可以在full_arr 上生成带有view 的事实上的结构化数组。好处是这些可以“就地”排序，保持稳定的行顺序：

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

视图非常适合执行数学数组运算，但对于排序来说，即使是我数据集中的单个子数组，它的效率也太低了。一般来说，结构化数组似乎并不能很好地扩展，即使它们具有非常有用的属性。如果有人知道为什么会这样，我很想知道。

使用非常大的数组最大限度地减少内存消耗并提高性能的一个不错的选择是构建一个由小而简单的函数组成的管道。函数完成后会清除局部变量，因此如果中间数据结构正在构建并占用内存，这可能是一个很好的解决方案。

这是我用来加速大规模数组排序的管道草图：

def process_sub_arrs(arr):
    """process a sub-array and return a 3-tuple of 1D values arrays"""

    return values1, values2, values3

def build_arr():
    """build the initial array by joining processed sub-arrays"""

    full_arr = np.zeros([868940742, 3])
    i, j = 0, 0

    for arr in list(arr1..arr22):  
        # indices (i, j) incremented at each loop based on sub-array size
        j += len(arr)
        full_arr[i:j, :] = np.column_stack( process_sub_arrs(arr) )
        i = j

    return full_arr

def sort_arr():
    """return full_arr and sort_idx"""

    full_arr = build_arr()
    sort_idx = np.argsort(full_arr[:, index])

    return full_arr[sort_idx]

def get_sorted_arr():
    """call through nested functions to return the sorted array"""

    sorted_arr = sort_arr()
    <process sorted_arr>

    return statistics

调用堆栈：get_sorted_arr --> sort_arr --> build_arr --> process_sub_arrs

一旦每个内部函数完成get_sorted_arr()，最后只保存排序后的数组，然后返回一个小的统计数据数组。

编辑：这里还值得指出的是，即使您能够使用更紧凑的dtype 来表示您的庞大数组，您也需要使用更高的精度进行汇总计算。例如，由于full_arr.dtype = np.float16，命令np.mean(full_arr[:,idx]) 尝试以半精度浮点数计算平均值，但是当对一个庞大的数组求和时，这很快就会溢出。使用np.mean(full_arr[:,idx], dtype=np.float64) 将防止溢出。

我最初发布这个问题是因为我对相同大小的数据集突然开始阻塞我的系统内存这一事实感到困惑，尽管在新的“慢”集合中唯一值的比例存在很大差异。 @ali_m 指出，确实，具有更少唯一值的更统一的数据更容易排序：

快速排序的 qsort 变体通过递归选择一个 'pivot' 数组中的元素，然后重新排序数组，使所有 小于枢轴值的元素放在它之前，并且所有 大于枢轴值的元素放在它之后。 等于枢轴的值已经排序，所以直观地说， 数组中的唯一值越少，数字越小 需要进行的交换。

关于这一点，我最终尝试解决此问题的最后一个更改是提前对较新的数据集进行舍入，因为插值步骤剩余的小数精度水平过高。这最终比其他节省内存的步骤产生了更大的影响，表明排序算法本身是这种情况下的限制因素。

期待其他 cmet 或任何人可能对这个主题提出的建议，我几乎肯定会误会一些技术问题，所以我很高兴收到回复 :-)