Эффективно индексируя массив numpy с массивом numpy

Question

У меня есть очень (очень, очень) большой двумерный массив - порядка тысячи столбцов, но несколько миллионов строк (достаточно, чтобы он не вписывался в память на моей 32-гигабайтной машине). Я хочу вычислить дисперсию каждой из тысяч столбцов. Один из ключевых фактов, который помогает: мои данные - 8-битные беззнаковые int.

Вот как я планирую приблизиться к этому. Сначала я построю новый двухмерный массив, называемый , числом с формой (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 

Answer 1

Я думаю, что это то, что вы хотите:

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

---

Но если массив имеет миллионы строк, вы все равно придется перебирать миллионы тех, , Следующий трюк дает близко к 5x ускорения на моей системе:

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