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

Question

I have a very (very, very) large two dimensional array - on the order of a thousand columns, but a couple of million rows (enough so that it doesn't fit in to memory on my 32GB machine). I want to compute the variance of each of the thousand columns. One key fact which helps: my data is 8-bit unsigned ints.

Here's how I'm planning on approaching this. I will first construct a new two dimensional array called counts with shape (1000, 256), with the idea that counts[i,:] == np.bincount(bigarray[:,i]) . Once I have this array, it will be trivial to compute the variance.

Trouble is, I'm not sure how to compute it efficiently (this computation must be run in real-time, and I'd like bandwidth to be limited by how fast my SSD can return the data). Here's something which works, but is god-awful slow:

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

Is there any way to write this to run faster? Something like this:

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

Answer 1

I think this is what you want:

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

---

But if your array has millions of rows, you are still going to have to iterate over millions of those. The following trick gives close to a 5x speed-up on my system:

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

And now:

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