Векторизовать разреженную сумму без scipy.sparse

Question

Я пытаюсь сделать пространственные производные и почти сумел получить все циклы из моего кода, но когда я пытаюсь суммировать все в конце, у меня проблема.

У меня есть набор N~=250k узлов. Я нашел индексы i,j пар узлов с i.size=j.size=~7.5M, которые находятся на определенном расстоянии поиска, первоначально исходящем от np.triu_indices(n,1), и прошли через последовательность булевых масок, чтобы смыть узлы, не влияющие друг на друга. Теперь я хочу подытожить влияние на каждый узел из других узлов.

настоящее время у меня это:

def sparseSum(a,i,j,n): 
    return np.array([np.sum(a[np.logical_or(i==k,j==k)],axis=0) for k in range(n)]) 

Это очень медленно. Я бы хотел что-то векторизовать. Если бы я был SciPy я мог сделать

def sparseSum(a,i,j,n): 
    sp=scipy.sparse.csr_matrix((a,(i,j)),shape=(n,n))+ scipy.sparse.csr_matrix((a,(j,i)),shape=(n,n)) 
    return np.sum(sp, axis=0) 

Но я делаю все это в рамках реализации Abaqus, который не включает в себя SciPy. Есть ли способ сделать это только для numpy?

Answer 1

Подход № 1: Вот подход использования matrix-multiplication и broadcasting -

K = np.arange(n)[:,None] 
mask = (i == K) | (j == K) 
out = np.dot(mask,a) 

Подход № 2: Для случаев с небольшим числом столбцов, мы можем использовать np.bincount для такого bin- на основе суммирования по каждому колонку, как так -

def sparseSum(a,i,j,n): 
    if len(a.shape)==1: 
     out=np.bincount(i,a,minlength=n)+np.bincount(j,a) 
    else: 
     ncols = a.shape[1] 
     out = np.empty((n,ncols)) 
     for k in range(ncols): 
      out[:,k] = np.bincount(i,a[:,k],minlength=n) + np.bincount(j,a[:,k]) 
    return out 

Answer 2

Вот не ключ решения, но один, который добавляет столбцы разреженная матрица. Он по существу вычисляет и использует представление csc

def sparse_col_sums(i, j, a, N): 
    order = np.lexsort(j, i) 
    io, jo, ao = i[order], j[order], a[order] 
    col_bnds = io.searchsorted(np.arange(N)) 
    return np.add.reduceat(ao, col_bnds)