大型网格的 Numpy 索引

Posted 2023-03-11

【中文标题】大型网格的 Numpy 索引

【英文标题】：Numpy indexing of large meshgrid

【发布时间】：2022-01-17 14:57:46

【问题描述】：

我想要一种方法来索引 d 维 numpymeshgrid：

实际上没有将完整的密集网格存储在内存中 支持全密集网格网格支持的所有索引类型

举个例子：

x = np.random.randn(Nx)
y = np.random.randn(Ny)
z = np.random.randn(Nz)

x_g, y_g, z_g = np.meshgrid(x, y, z, indexing='ij')
X = np.concatenate([x_g[..., None], y_g[..., None], z_g[..., None]], axis=-1)

一旦我有了X，我就可以使用numpy的所有索引方法，例如：

X[10:20,1:9:3,:]

X[(
    [0, 1, 3, 5],
    [1, 1, 3, 3],
    [2, 7, 8, 9]
    )]

在某些情况下，产品 Nx * Ny * Nz 太大而无法放入内存（例如 Nx = Ny = Nz = 4000），但我希望能够使用索引取出我的多维数据集。

有没有一种方法可以在不重新编码 numpy 的索引逻辑的情况下实现这一点？ 3d 中的答案如下所示：

def meshgrid_index(x, y, z, index):
    # index is the argument normally passed to X.__getitem__
    # should support all types of indexing
    ...

谢谢！

【参考方案1】：

不幸的是，Numpy 主要使用 数组视图。实际上，视图只是对 原始内存缓冲区、数据类型、维度（d）、步幅（d 相对整数）的引用和一个形状（d 个正整数）。步幅非常适合生成重复元素的视图，颠倒项目的顺序，但它们不足以生成您想要的网格网格（我认为这可以生成一个子集但不能生成完整网格，因为步幅需要以某种方式在某些时候改变）。

因此，您需要创建网格阵列（可能只是一部分），或者在没有 Numpy 网格功能的情况下自行创建。带有循环的 Numba 代码通常非常有效且可读性强。你想要的可能是一个（通用的）惰性 ND 视图，它很少被语言和库支持（因为它实现起来很复杂和/或通常效率不高）。

---

初步答案：

当内存不足时，解决此问题的通常方法是逐块。这个想法是使用基本的 Python 循环遍历所有可能的块，并对 x、y 和 z 数组进行切片以生成一块网格网格（例如 np.meshgrid(x[xChunkId:xChunkId+xChunkSize], y[yChunkId:yChunkId+yChunkSize], z[zChunkId:zChunkId+zChunkSize], indexing='ij')）。由于 CPU 缓存（假设块足够小以适合缓存），这种方法通常会加快计算可以放入 RAM 中的巨大数组。只要块足够大，循环就不会引入显着的开销。

请注意，生成如此巨大的（临时）数组甚至可能不是强制性的。事实上，像 Numba 或 Cython 这样的工具可以帮助您在网格上进行迭代，而无需像这样生成临时数组。这需要更少的内存并且在实践中更快（由于 CPU 缓存），更不用说这些工具可用于使您的代码并行运行（假设您的算法可以并行化）。请注意，这些工具还可以大大降低小块的循环成本（并且使用小块通常对性能有好处，同样是因为 CPU 缓存）。

如果由于需要一次计算整个数组而无法做到这一点，那么唯一的解决方案是将巨大的数组存储/加载到存储设备中。一种快速便捷的方法是使用memory-mapped arrays。请注意，这会慢得多（尤其是使用 HDD）。

请注意，将此特定数组存储在存储设备上会非常低效，因为它包含大量重复值。

为了性能，减小数组的大小是关键。实际上，默认情况下，您的数组可能会占用 1.4 TiB 的内存，而 Numpy 往往会创建许多临时数组。因此，最终所需的存储容量可以迅速超过 10 TiB。对于这种用途，HDD 非常慢，而高性能 SSD/持久内存非常昂贵。处理较小的类型，如 32 位浮点数或整数（使用 dtype 参数和 astype 方法）会有很大帮助。这不是免费的，尽管这样做是以降低准确性和可能的​​溢出为代价的。

【讨论】：

感谢您的回复，很多有趣的点！但是，在我看来，它们中没有一个适用于我需要做的事情。我真的没有大数据问题——我想通过索引从大立方体中提取的位本身不一定很大。我只是在寻找支持所有 numpy 索引方法的机器，而无需实际计算/存储（在内存或磁盘中）多维数据集。

哈好吧。对不起。我误解了确切的问题。我编辑了答案。简而言之：这在 Numpy 中是不可能直接实现的（至少没有视图）。