Быстрые ненулевые индексы на строку / столбец для (разреженного) массива 2D numpy

Question

Я ищу самый быстрый способ получить список ненулевых индексов 2D-массива для каждой строки и каждого столбца. Ниже приведен рабочий фрагмент кода:

preds = [matrix[:,v].nonzero()[0] for v in range(matrix.shape[1])]
descs = [matrix[v].nonzero()[0] for v in range(matrix.shape[0])]

Пример ввода:

matrix = np.array([[0,0,0,0],[1,0,0,0],[1,1,0,0],[1,1,1,0]])

Пример вывода

preds = [array([1, 2, 3]), array([2, 3]), array([3]), array([], dtype=int64)]
descs = [array([], dtype=int64), array([0]), array([0, 1]), array([0, 1, 2])]

(Списки называются preds и descs, потому что они относятся к предшественникам и потомкам в DAG, когда матрица интерпретируется как матрица смежности, но это не является существенным для вопроса.)

Пример синхронизации: Для целей синхронизации следующая матрица является хорошим представителем:

test_matrix = np.zeros(shape=(4096,4096),dtype=np.float32)
for k in range(16):
    test_matrix[256*(k+1):256*(k+2),256*k:256*(k+1)]=1

Фон: В моем коде эти две строки занимают 75% времени для Матрица 4000x4000, тогда как последующая топологическая сортировка и алгоритм DP занимают только оставшуюся часть квартала. Примерно 5% значений в матрице отличны от нуля, поэтому может быть применимо решение с разреженной матрицей.

Спасибо.

(По предложению, размещенному здесь также: https://scicomp.stackexchange.com/questions/35242/fast-nonzero-indices-per-row-column-for-sparse-2d-numpy-array Там также есть ответы, время на которые я укажу в комментариях. Эта ссылка содержит принятый ответ, который работает вдвое быстрее. )

Answer 1

Если у вас достаточно мотивации, Нумба может творить удивительные вещи. Вот быстрая реализация необходимого вам logi c. Вкратце, он вычисляет эквивалент np.nonzero(), но попутно включает информацию для последующего преобразования индексов в нужный вам формат. Информация основана на sparse.csr.indptr и sparse.csc.indptr.

import numpy as np
import numba as nb


@nb.jit
def cumsum(arr):
    result = np.empty_like(arr)
    cumsum = result[0] = arr[0]
    for i in range(1, len(arr)):
        cumsum += arr[i]
        result[i] = cumsum
    return result


@nb.jit
def count_nonzero(arr):
    arr = arr.ravel()
    n = 0
    for x in arr:
        if x != 0:
            n += 1
    return n


@nb.jit
def row_col_nonzero_nb(arr):
    n, m = arr.shape
    max_k = count_nonzero(arr)
    indices = np.empty((2, max_k), dtype=np.uint32)
    i_offset = np.zeros(n + 1, dtype=np.uint32)
    j_offset = np.zeros(m + 1, dtype=np.uint32)
    n, m = arr.shape
    k = 0
    for i in range(n):
        for j in range(m):
            if arr[i, j] != 0:
                indices[:, k] = i, j
                i_offset[i + 1] += 1
                j_offset[j + 1] += 1
                k += 1
    return indices, cumsum(i_offset), cumsum(j_offset)


def row_col_idx_nonzero_nb(arr):
    (ii, jj), jj_split, ii_split = row_col_nonzero_nb(arr)
    ii_ = np.argsort(jj)
    ii = ii[ii_]
    return np.split(ii, ii_split[1:-1]), np.split(jj, jj_split[1:-1])

По сравнению с вашим подходом (row_col_idx_sep() ниже) и множеством других, согласно @ hpaulj answer (row_col_idx_sparse_lil()) и @ knl ответ от scicomp.stackexchange.com (row_col_idx_sparse_coo()):

def row_col_idx_sep(arr):
    return (
        [arr[:, j].nonzero()[0] for j in range(arr.shape[1])],
        [arr[i, :].nonzero()[0] for i in range(arr.shape[0])],)

def row_col_idx_zip(arr):
    n, m = arr.shape
    ii = [[] for _ in range(n)]
    jj = [[] for _ in range(m)]
    x, y = np.nonzero(arr)
    for i, j in zip(x, y):
        ii[i].append(j)
        jj[j].append(i)
    return jj, ii

import scipy as sp
import scipy.sparse


def row_col_idx_sparse_coo(arr):
    coo_mat = sp.sparse.coo_matrix(arr)
    csr_mat = coo_mat.tocsr()
    csc_mat = coo_mat.tocsc()
    return (
        np.split(csc_mat.indices, csc_mat.indptr)[1:-1],
        np.split(csr_mat.indices, csr_mat.indptr)[1:-1],)

def row_col_idx_sparse_lil(arr):
    lil_mat = sp.sparse.lil_matrix(arr)
    return lil_mat.T.rows, lil_mat.rows

Для входных данных, созданных с использованием:

def gen_input(n, density=0.1, dtype=np.float32):
    arr = np.zeros(shape=(n, n), dtype=dtype)
    indices = tuple(np.random.randint(0, n, (2, int(n * n * density))).tolist())
    arr[indices] = 1.0
    return arr

Можно получить (ваш test_matrix имел приблизительно 0,06 ненулевой плотности):

m = gen_input(4096, density=0.06)
%timeit row_col_idx_sep(m)
# 1 loop, best of 3: 767 ms per loop
%timeit row_col_idx_zip(m)
# 1 loop, best of 3: 660 ms per loop
%timeit row_col_idx_sparse_coo(m)
# 1 loop, best of 3: 205 ms per loop
%timeit row_col_idx_sparse_lil(m)
# 1 loop, best of 3: 498 ms per loop
%timeit row_col_idx_nonzero_nb(m)
# 10 loops, best of 3: 130 ms per loop

Это означает, что это почти вдвое быстрее самого быстрого scipy.sparse -основанный подход.

Answer 2

In [182]: arr = np.array([[0,0,0,0],[1,0,0,0],[1,1,0,0],[1,1,1,0]])                      

Данные представлены во всем массиве nonzero, но не разбиты на массивы по строкам / столбцам:

In [183]: np.nonzero(arr)                                                                
Out[183]: (array([1, 2, 2, 3, 3, 3]), array([0, 0, 1, 0, 1, 2]))
In [184]: np.argwhere(arr)                                                               
Out[184]: 
array([[1, 0],
       [2, 0],
       [2, 1],
       [3, 0],
       [3, 1],
       [3, 2]])

Возможно, можно разбить array([1, 2, 2, 3, 3, 3]) на подсписки, [1,2,3],[2,3],[3],[] на основе другого массива. Но для разработки logi c для этого может потребоваться некоторое время, и нет никакой гарантии, что это будет быстрее, чем ваши итерации строки / столбца.

Логические операции могут уменьшить логический массив до столбца или строки , давая строки или столбцы, в которых встречается ненулевое значение, но опять же не рваное:

In [185]: arr!=0                                                                         
Out[185]: 
array([[False, False, False, False],
       [ True, False, False, False],
       [ True,  True, False, False],
       [ True,  True,  True, False]])
In [186]: (arr!=0).any(axis=0)                                                           
Out[186]: array([ True,  True,  True, False])
In [187]: np.nonzero((arr!=0).any(axis=0))                                               
Out[187]: (array([0, 1, 2]),)
In [188]: np.nonzero((arr!=0).any(axis=1))                                               
Out[188]: (array([1, 2, 3]),)
In [189]: arr                                                                            
Out[189]: 
array([[0, 0, 0, 0],
       [1, 0, 0, 0],
       [1, 1, 0, 0],
       [1, 1, 1, 0]])

Формат scipy.sparse lil генерирует нужные вам данные:

In [190]: sparse                                                                         
Out[190]: <module 'scipy.sparse' from '/usr/local/lib/python3.6/dist-packages/scipy/sparse/__init__.py'>
In [191]: M = sparse.lil_matrix(arr)                                                     
In [192]: M                                                                              
Out[192]: 
<4x4 sparse matrix of type '<class 'numpy.longlong'>'
    with 6 stored elements in List of Lists format>
In [193]: M.rows                                                                         
Out[193]: array([list([]), list([0]), list([0, 1]), list([0, 1, 2])], dtype=object)
In [194]: M.T                                                                            
Out[194]: 
<4x4 sparse matrix of type '<class 'numpy.longlong'>'
    with 6 stored elements in List of Lists format>
In [195]: M.T.rows                                                                       
Out[195]: array([list([1, 2, 3]), list([2, 3]), list([3]), list([])], dtype=object)

Но время вероятно, ничем не лучше, чем итерация строки или столбца.