✨ Extend _l2norm for sparse input

[mbuttner]

Is your feature request related to a problem? Please describe.
I tried to compute a WNN integration on a MuData object, i.e. I executed the following steps:

mu.pp.l2norm(mdata)
mu.pp.neighbors(mdata)
mu.tl.umap(mdata)

However, my Python session keeps crashing at the l2norm step. Omitting this step, computing neighbors and a UMAP works fine. I looked into the mu.pp.l2norm function and isolated the function np.linalg.norm() in the _l2norm() function as the culprit. It crashes reliably for sparse matrices. Converting the data matrix into dense format fixes the issue. This PR extends the _l2norm function to sparse matrices.

[ilia-kats]

Thanks, looks good to me. I'll merge as soon as you make the CI pass (I think black is complaining about the trailing whitespace in line 159)

[mbuttner]

Thanks! I did some more testing today and spotted a different behavior in my testing environment, where the norm matrix is still sparse after dividing with the l2norm vector, and add handling of infinite values for sparse matrices.

[ilia-kats]

I think for csr_matrix, csc_matrix, and coo_matrix it would probably be more efficient to do something like norm.data[~np.isfinite(norm.data)] = 0. Sparse matrices in AnnData are usually CSR or CSC (in fact, IO only has support for CSR and CSC implemented), so it would probably make sense to also add that special case.

[mbuttner]

Thats a very good point! I simplified it in commit e7b7b12

[ilia-kats]

I don't think that will do what you're expecting it to do. The line below (X[:] = norm) overwrites the matrix in the AnnData object with the normalized values. Now you're replacing the local X variable with something else, so the normalized values will be lost. The question is: Is there ever a situation where X is sparse and norm is not? At least if X is a csr_matrix or a csc_matrix norm should always be a COO matrix.

Then you could also get rid of the if clause in line 162.

[mbuttner]

Thanks for the feedback. I removed the lines in commit e7b7b12

[ilia-kats]

I'm not sure at the moment, but can you test if this does what we expect it to do for sparse matrices? I.e. it doesn't change the sparsity structure? Otherwise it would perhaps make sense to do something like

if sparse_X and (isspmatrix_csc(X) or isspmatrix_csr(X) or issspmatrix_coo(X)):
    X.data[:] = norm.data[:]
else:
    X[:] = norm

(I don't know enough about the other sparse matrix types off the top of my head, but this should cover most usecases)

[mbuttner]

I had a weird case under Python 3.8 and scipy==1.10.0 when norm became dense while X was still sparse, and then X[:] = norm crashed my kernel. That's why I wanted to catch that error.

[mbuttner]

I'll test your suggestion.

[mbuttner]

I tested your suggestion and included it in the code with minimal adaptations (commit cd3f942).

[mbuttner]

Following up on the part where norm became dense and X was still sparse. I added a clause to convert norm into a csr_matrix in that case.

[mbuttner]

Please merge if you think it fits to muon. I tested my version with sparse matrices and did not spot any issues.

[ilia-kats]

Thank you for the contribution.