Задайте строки матрицы scipy.sparse, которые удовлетворяют определенным условиям для нулей

Question

Интересно, какой лучший способ заменить строки, которые не удовлетворяют определенному условию нулями для разреженных матриц. Например (я использую простые массивы для иллюстрации):

Я хочу заменить каждую строку, сумма которых больше, чем 10 с рядом нулей

a = np.array([[0,0,0,1,1], 
       [1,2,0,0,0], 
       [6,7,4,1,0], # sum > 10 
       [0,1,1,0,1], 
       [7,3,2,2,8], # sum > 10 
       [0,1,0,1,2]]) 

Я хочу, чтобы заменить а [2] и [ 4] с нулями, так что мой вывод должен выглядеть следующим образом:

array([[0, 0, 0, 1, 1], 
     [1, 2, 0, 0, 0], 
     [0, 0, 0, 0, 0], 
     [0, 1, 1, 0, 1], 
     [0, 0, 0, 0, 0], 
     [0, 1, 0, 1, 2]]) 

Это довольно прямо вперед для плотных матриц:

row_sum = a.sum(axis=1) 
to_keep = row_sum >= 10 
a[to_keep] = np.zeros(a.shape[1]) 

Однако, когда я пытаюсь:

s = sparse.csr_matrix(a) 
s[to_keep, :] = np.zeros(a.shape[1]) 

Я получаю эту ошибку:

raise NotImplementedError("Fancy indexing in assignment not " 
NotImplementedError: Fancy indexing in assignment not supported for csr matrices. 

Следовательно, мне нужно другое решение для разреженных матриц. Я пришел с этим:

def zero_out_unfit_rows(s_mat, limit_row_sum): 
    row_sum = s_mat.sum(axis=1).T.A[0] 
    to_keep = row_sum <= limit_row_sum 
    to_keep = to_keep.astype('int8') 
    temp_diag = get_sparse_diag_mat(to_keep) 
    return temp_diag * s_mat 

def get_sparse_diag_mat(my_diag): 
    N = len(my_diag) 
    my_diags = my_diag[np.newaxis, :] 
    return sparse.dia_matrix((my_diags, [0]), shape=(N,N)) 

Это основывается на том факте, что, если мы устанавливаем 2 и 4 элементы диагонали единичной матрицы к нулю, то строки из предварительно умноженного матрицы устанавливаются в нуль.

Однако, я чувствую, что есть лучшее, более scipynic-решение. Есть ли лучшее решение?

Answer 1

Не уверен, что это очень scithonic, но многие операции с разреженными матрицами лучше делать, обратившись непосредственно к кишкам. Для вашего случая я лично сделал бы:

a = np.array([[0,0,0,1,1], 
       [1,2,0,0,0], 
       [6,7,4,1,0], # sum > 10 
       [0,1,1,0,1], 
       [7,3,2,2,8], # sum > 10 
       [0,1,0,1,2]]) 
sps_a = sps.csr_matrix(a) 

# get sum of each row: 
row_sum = np.add.reduceat(sps_a.data, sps_a.indptr[:-1]) 

# set values to zero 
row_mask = row_sum > 10 
nnz_per_row = np.diff(sps_a.indptr) 
sps_a.data[np.repeat(row_mask, nnz_per_row)] = 0 
# ask scipy.sparse to remove the zeroed entries 
sps_a.eliminate_zeros() 

>>> sps_a.toarray() 
array([[0, 0, 0, 1, 1], 
     [1, 2, 0, 0, 0], 
     [0, 0, 0, 0, 0], 
     [0, 1, 1, 0, 1], 
     [0, 0, 0, 0, 0], 
     [0, 1, 0, 1, 2]]) 
>>> sps_a.nnz # it does remove the entries, not simply set them to zero 
10