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

时间:2013-03-27 20:40:44

标签: python numpy bigdata

我有一个非常(非常非常)大的二维数组 - 大约有一千列,但有几百万行(足以让它不适合我的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 个答案:

答案 0 :(得分: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