Add norm_sq/slice_norms/to_scipy_sparse to normalized views - #47
Merged
Conversation
parafac2 touches its input matrix through matmul/rmatmul (already supported via __matmul__/__rmatmul__) plus a squared-Frobenius-norm reduction it currently only knows how to compute for a plain np.ndarray or scipy.sparse array. Add that as norm_sq()/slice_norms() on NormalizedViewBase, computed with new numba kernels that reuse the same O(nnz) traversal as the existing matmul/toarray kernels, so a normalized view can satisfy parafac2's duck-typed backend contract without materializing anything. Also add to_scipy_sparse() (the uncentered, scaled sparse term with the same sparsity pattern as the raw array) and a means property (the per-gene correction to subtract from it), for code that only knows how to move a plain NumPy/SciPy array onto a device -- e.g. so a normalized view can be materialized into a real (and, for typical single-cell data, tiny) sparse array before running through parafac2's existing CuPy/MLX GPU path, rather than needing new GPU-native kernels of its own. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
parafac2touches its input matrixXthrough exactly three things: matmul/rmatmul (
X @ rhs/lhs @ X),.shape/.dtype, and a squared-Frobenius-normreduction (
sum((X - means) ** 2), overall and per-condition-group). Thenormalized VCSR/VCSC views already implement the first two --
__matmul__/__rmatmul__/.shape/.dtype-- but not the norm reduction, so passing oneof these views as
parafac2'sXcurrently crashes immediately incalc_norm_sq(unsupported operand type(s) for ** ...).This adds:
norm_sq(): squared Frobenius norm of the full normalized matrix, inO(nnz + n_cols). Uses the samesum(Delta^2) - 2*sum(Delta*offset) + n_rows*sum(offset^2)expansionparafac2.utils.calc_norm_sq's own sparsepath already uses for a scipy CSR array plus an external
meansvector --just carried out against this view's own uncentered
Deltaterm instead(see
_vcs_matmul.py), sinceoffset = col_post_scale * col_meanisexactly that
means.slice_norms(condition_idxs, n_cond): the per-condition-group versionof the same, matching
parafac2.utils.calc_slice_norms.meansproperty:col_post_scale * col_mean, the per-gene correctiona caller following that same sparse-plus-external-means convention should
subtract externally.
to_scipy_sparse(): materializes the uncentered, scaledDeltatermas a real
scipy.sparsearray (same sparsity pattern as the underlying rawarray -- no denser than the input). Combined with
means, this lets codethat only knows how to move a plain NumPy/SciPy array onto a device (e.g.
parafac2's CuPy/MLX GPU backend) run on a normalized view without anyGPU-native kernels of its own -- for typical single-cell data the
materialized array is a small fraction of the dense matrix's size.
All four are implemented with new
numbakernels mirroring the existingtoarray/matmul kernels' per-format (VCSR/VCSC) traversal, so they stayO(nnz)and never materialize a dense array.This is the
vsparsehalf of a two-repo change; the companion PR onparafac2generalizescalc_norm_sq/calc_slice_norms/GPUMatrixto usethese (and any future type's) duck-typed methods. Verified end-to-end with
both repos' local checkouts: a
vsparse-backedAnnDatanow fits correctlyvia
parafac2_ndonbackend='cpu', and materializing viato_scipy_sparse()+meanslets the same dataset fit on a real CuPy GPU.Test plan
uv run pytest(1241 passed, 49 skipped) -- includes new propertytests in
tests/test_property_norm_stats.py(norm_sq/slice_normsagainst a dense
toarray()reference, for both VCSC/VCSR and everyrecipe;
to_scipy_sparse()+meansreconstructingtoarray(); samesparsity pattern as the raw array) and example tests added to
tests/test_vcs_norm.py(all-zero-matrix edge cases, return type/dtype)uv run ruff check ./uv run ruff format --check ./uv run codespelluv run ty checkvsparse+parafac2checkouts, including a real CuPy GPU fit after materializing via
to_scipy_sparse()/means.