使用 numpy 数组有效地索引 numpy 数组

Posted 2023-04-18

【中文标题】使用 numpy 数组有效地索引 numpy 数组

【英文标题】：Efficiently indexing numpy array with a numpy array

【发布时间】：2013-03-27 20:40:43

【问题描述】：

我有一个非常（非常、非常）大的二维数组 - 大约有一千列，但有几百万行（足以让它不适合我的 32GB 机器的内存）。我想计算每一千列的方差。一个有帮助的关键事实：我的数据是 8 位无符号整数。

这就是我打算如何处理这个问题。我将首先构造一个名为 counts 的新二维数组，其形状为 (1000, 256)，其想法是 counts[i,:] == np.bincount(bigarray[:,i])。一旦我有了这个数组，计算方差就很简单了。

问题是，我不确定如何有效地计算它（这个计算必须实时运行，我希望带宽受限于我的 SSD 可以多快返回数据）。下面是一些可行的方法，但速度慢得可怕：

counts = np.array((1000,256))
for row in iterator_over_bigaray_rows():
    for i,val in enumerate(row):
        counts[i,val] += 1

有什么办法可以写得更快吗？像这样的：

counts = np.array((1000,256))
for row in iterator_over_bigaray_rows():
    counts[i,:] = // magic np one-liner to do what I want

【问题讨论】：

+1 处理庞大数组的聪明方法！但是你仍然需要遍历很多行......

是的，但无法避免。我实际上并没有那样迭代它（我以块的形式从磁盘上加载它，然后异步地迭代块）。

【参考方案1】：

我想这就是你想要的：

counts[np.arange(1000), row] += 1

---

但是，如果您的数组有数百万行，您仍然需要遍历数百万行。以下技巧使我的系统速度提高了近 5 倍：

chunk = np.random.randint(256, size=(1000, 1000))

def count_chunk(chunk):
    rows, cols = chunk.shape
    col_idx = np.arange(cols) * 256
    counts = np.bincount((col_idx[None, :] + chunk).ravel(),
                         minlength=256*cols)
    return counts.reshape(-1, 256)

def count_chunk_by_rows(chunk):
    counts = np.zeros(chunk.shape[1:]+(256,), dtype=np.int)
    indices = np.arange(chunk.shape[-1])
    for row in chunk:
        counts[indices, row] += 1
    return counts

现在：

In [2]: c = count_chunk_by_rows(chunk)

In [3]: d = count_chunk(chunk)

In [4]: np.all(c == d)
Out[4]: True

In [5]: %timeit count_chunk_by_rows(chunk)
10 loops, best of 3: 80.5 ms per loop

In [6]: %timeit count_chunk(chunk)
100 loops, best of 3: 13.8 ms per loop