稀疏 ndarray 求和

Question

我有一个 3 维数组 (np.ndarray)，其中大部分为 0。现在我想在第一个维度上对它们求和，但这相当慢。我研究了 csr_matrix，但 csr 不支持 3 维数组。有没有一种更快的方法来对几乎稀疏的 nd 数组求和？以下是我当前代码的摘录。

相关问题： sparse 3d matrix/array in Python?（创建一个自制的稀疏 ndarray 类，矫枉过正？）

r = np.array([  [[1, 0, 0, 0],
                 [1, 0, 0, 0],
                 [0, 0, 1, 0]],

                [[0, 1, 0, 0],
                 [0, 0, 0, 1],
                 [0, 0, 2, 0]],

                [[0, 1, 0, 0],
                 [0, 0, 0, 0],
                 [0, 0, 0, 0]],

                [[0, 0, 0, 1],
                 [0, 0, 0, 0],
                 [0, 0, 0, 0]]], dtype=int)
np.sum(r,axis=0)
Out[35]: 
array([[1, 2, 0, 1],
       [1, 0, 0, 1],
       [0, 0, 3, 0]])

编辑

在下面 hpaulj 的回答之后，我做了更多的计时测试，见下文。似乎重塑对总和没有多大好处，而将它们转换为 csr_matrix 并返回 numpy 会杀死性能。我仍在考虑直接使用索引（下面称为rand_persons、rand_articles 和rand_days，因为在我最初的问题中，我也使用这些索引制作了大 ndarray。

from timeit import timeit
from scipy.sparse import csr_matrix
import numpy as np

def create_test_data():
    '''
    dtype = int64
    1% nonzero, 1000x1000x100: 1.3 s, 
    1% nonzero, 10000x1000x100: 13.3 s
    0.1% nonzero, 10000x1000x100: 2.7 s
    1ppm nonzero, 10000x1000x100: 0.007 s
    '''
    global purchases
    N_persons = 10000
    N_articles = 1000
    N_days = 100
    purchases = np.zeros(shape=(N_days, N_persons, N_articles), dtype=int)
    N_elements = N_persons * N_articles * N_days
    rand_persons = np.random.choice(a=range(N_persons), size=N_elements / 1e6, replace=True)
    rand_articles = np.random.choice(a=range(N_articles), size=N_elements / 1e6, replace=True)
    rand_days = np.random.choice(a=range(N_days), size=N_elements / 1e6, replace=True)
    for (i, j, k) in zip(rand_persons, rand_articles, rand_days):
        purchases[k, i, j] += 1

def sum_over_first_dim_A():
    '''
    0.1% nonzero, 10000x1000x99: 1.57s (average over 10)
    1ppm nonzero, 10000x1000x99: 1.70s (average over 10)
    '''
    global purchases
    d = purchases[:99, :, :]
    np.sum(d, axis=0)
def sum_over_first_dim_B():
    '''
    0.1% nonzero, 10000x1000x99: 1.55s (average over 10)
    1ppm nonzero, 10000x1000x99: 1.37s (average over 10)
    '''
    global purchases
    d = purchases[:99, :, :]
    (N_days, N_persons, N_articles) = d.shape 
    d.reshape(N_days, -1).sum(0).reshape(N_persons, N_articles) 
def sum_over_first_dim_C():
    '''
    0.1% nonzero, 10000x1000x99: 7.54s (average over 10)
    1ppm nonzero, 10000x1000x99: 7.44s (average over 10)
    '''
    global purchases
    d = purchases[:99, :, :]
    (N_days, N_persons, N_articles) = d.shape 
    r = csr_matrix(d.reshape(N_days, -1))
    t = r.sum(axis=0)
    np.reshape(t, newshape=(N_persons, N_articles))

if __name__ == '__main__':
    print (timeit(create_test_data, number=10))
    print (timeit(sum_over_first_dim_A, number=10))
    print (timeit(sum_over_first_dim_B, number=10))
    print (timeit(sum_over_first_dim_C, number=10))

编辑 2

我现在找到了一种更快的求和方法：我用稀疏矩阵创建了一个 numpy 数组。但是，这些矩阵的初始创建还有一段时间。我现在用一个循环来做这个。有没有办法加快速度？

def create_test_data():
    [ ... ]
    '''
    0.1% nonzero, 10000x1000x100: 2.1 s
    1ppm nonzero, 10000x1000x100: 0.45 s
    '''
    global sp_purchases
    sp_purchases = np.empty(N_days, dtype=lil_matrix)
    for i in range(N_days):
        sp_purchases[i] = lil_matrix((N_persons, N_articles))
    for (i, j, k) in zip(rand_persons, rand_articles, rand_days):
        sp_purchases[k][i, j] += 1

def sum_over_first_dim_D():
    '''
    0.1% nonzero, 10000x1000x99: 0.47s (average over 10)
    1ppm nonzero, 10000x1000x99: 0.41s (average over 10)
    '''
    global sp_purchases
    d = sp_purchases[:99]
    np.sum(d)

Answer 1

您可以将数组重新整形为 2d，求和，然后重新整形

r.reshape(4,-1).sum(0).reshape(3,4)   # == r.sum(0)

这种重塑不会增加太多的处理时间。您可以将该 2d 转换为稀疏的，看看是否可以节省任何时间。我的猜测是，您的数组必须非常大且非常稀疏，才能击败直接的 numpy 总和。如果您有其他理由使用稀疏格式，它可能是值得的，但只需要做这个总和，不。但是你自己测试一下。

Answer 2

由于您的数据已经是稀疏格式（索引和值），您可以自己计算总和。只需创建一个与最终求和数组大小相同的数组，然后遍历索引，将相应的值求和到正确的插槽中。下面的sum2d 函数显示了在对第一个维度求和的情况下您将如何做到这一点：

import timeit
import numpy as np

n = 1000
s = 1000
inds = np.random.randint(0, n, size=(s, 3))
vals = np.random.normal(size=s)


def sum3d():
    a = np.zeros((n, n, n))
    for [i, j, k], v in zip(inds, vals):
        a[i, j, k] = v

    return a.sum(axis=0)


def sum2d():
    b = np.zeros((n, n))
    for [i, j, k], v in zip(inds, vals):
        b[j, k] += v

    return b


kwargs = dict(repeat=3, number=1)
print(min(timeit.repeat('sum3d()', 'from __main__ import sum3d', **kwargs)))
print(min(timeit.repeat('sum2d()', 'from __main__ import sum2d', **kwargs)))
assert np.allclose(sum3d(), sum2d())