Fast nonzero indices per row/column for (sparse) 2D numpy array

Question

I am looking for the fastest way to obtain a list of the nonzero indices of a 2D array per row and per column. The following is a working piece of code:

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

Example input:

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

Example output

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])]

(The lists are called preds and descs because they refer to the predecessors and descendants in a DAG when the matrix is interpreted as an adjacency matrix but this is not essential to the question.)

I was wondering whether this might be doable with some sort of sparse matrix (CSR, CSC, COO etc.) from scipy.sparse but I am unfamiliar with them and have not got that working. I don't necessarily need to use these types if a faster option exists.

Timing example: For timing purposes, the following matrix is a good representative:

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

Thank you.

Background: In my code, these two lines take 75% of the time for a 4000x4000 matrix whereas the ensuing topological sort and DP algorithm take only the rest of the quarter. If someone know how to do this much more efficiently it would be greatly appreciated. Roughly 5% of the matrix has nonzero values.

(on suggestion I moved the question here from: https://stackoverflow.com/questions/62065793/fast-nonzero-indices-per-row-column-for-sparse-2d-numpy-array Contains several useful answers)

Answer 1

Ideally you would have the matrix already in a sparse matrix data structure. But for this example we can do the conversion via

In [31]: M = scipy.sparse.coo_matrix(np.array([[0,0,0,0],[1,0,0,0],[1,1,0,0],[1,1,1,0]]))

Then you can do

In [32]: Mcsr = M.tocsr() 

In [33]: np.split(Mcsr.indices, Mcsr.indptr)         
Out[33]: 
[array([], dtype=int32),
 array([], dtype=int32),
 array([0], dtype=int32),
 array([0, 1], dtype=int32),
 array([0, 1, 2], dtype=int32),
 array([], dtype=int32)]

to get descs. Similarly,

In [34]: Mcsc = M.tocsc() 

In [35]: np.split(Mcsc.indices, Mcsc.indptr)        
Out[35]: 
[array([], dtype=int32),
 array([1, 2, 3], dtype=int32),
 array([2, 3], dtype=int32),
 array([3], dtype=int32),
 array([], dtype=int32),
 array([], dtype=int32)]

gives preds. I have no idea if this is more effective in practice though.