Skip to content

Add norm_sq/slice_norms/to_scipy_sparse to normalized views - #47

Merged
aarmey merged 1 commit into
mainfrom
parafac2-duck-typed-backend
Sep 13, 2026
Merged

Add norm_sq/slice_norms/to_scipy_sparse to normalized views#47
aarmey merged 1 commit into
mainfrom
parafac2-duck-typed-backend

Conversation

@aarmey

@aarmey aarmey commented Sep 13, 2026

Copy link
Copy Markdown
Member

Summary

parafac2 touches its input matrix X through exactly three things: matmul/
rmatmul (X @ rhs/lhs @ X), .shape/.dtype, and a squared-Frobenius-norm
reduction (sum((X - means) ** 2), overall and per-condition-group). The
normalized VCSR/VCSC views already implement the first two -- __matmul__/
__rmatmul__/.shape/.dtype -- but not the norm reduction, so passing one
of these views as parafac2's X currently crashes immediately in
calc_norm_sq (unsupported operand type(s) for ** ...).

This adds:

  • norm_sq(): squared Frobenius norm of the full normalized matrix, in
    O(nnz + n_cols). Uses the same sum(Delta^2) - 2*sum(Delta*offset) + n_rows*sum(offset^2) expansion parafac2.utils.calc_norm_sq's own sparse
    path already uses for a scipy CSR array plus an external means vector --
    just carried out against this view's own uncentered Delta term instead
    (see _vcs_matmul.py), since offset = col_post_scale * col_mean is
    exactly that means.
  • slice_norms(condition_idxs, n_cond): the per-condition-group version
    of the same, matching parafac2.utils.calc_slice_norms.
  • means property: col_post_scale * col_mean, the per-gene correction
    a caller following that same sparse-plus-external-means convention should
    subtract externally.
  • to_scipy_sparse(): materializes the uncentered, scaled Delta term
    as a real scipy.sparse array (same sparsity pattern as the underlying raw
    array -- no denser than the input). Combined with means, this lets code
    that 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 any
    GPU-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 numba kernels mirroring the existing
toarray/matmul kernels' per-format (VCSR/VCSC) traversal, so they stay
O(nnz) and never materialize a dense array.

This is the vsparse half of a two-repo change; the companion PR on
parafac2 generalizes calc_norm_sq/calc_slice_norms/GPUMatrix to use
these (and any future type's) duck-typed methods. Verified end-to-end with
both repos' local checkouts: a vsparse-backed AnnData now fits correctly
via parafac2_nd on backend='cpu', and materializing via
to_scipy_sparse() + means lets the same dataset fit on a real CuPy GPU.

Test plan

  • uv run pytest (1241 passed, 49 skipped) -- includes new property
    tests in tests/test_property_norm_stats.py (norm_sq/slice_norms
    against a dense toarray() reference, for both VCSC/VCSR and every
    recipe; to_scipy_sparse() + means reconstructing toarray(); same
    sparsity 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 codespell
  • uv run ty check
  • Manual end-to-end smoke test with local vsparse + parafac2
    checkouts, including a real CuPy GPU fit after materializing via
    to_scipy_sparse()/means.

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>
@aarmey
aarmey merged commit a748874 into main Sep 13, 2026
5 of 6 checks passed
@aarmey
aarmey deleted the parafac2-duck-typed-backend branch September 13, 2026 02:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant