Generalize normalized views into declarative normalization recipes - #41
Merged
Conversation
Adapts VCSCArrayNormalized/VCSRArrayNormalized to fishidaho's proposed y = (g(x*a*b) - c) * s form (#40), with five built-in recipes (raw, cp10k_log1p, parafac2, scanpy, pearson). .normalized(view, recalculate=True) lets a caller switch recipes without recomputing statistics once each has been built; VCSCAnnData.normalized() additionally records a/b/c/s into obs/varm/uns (named vsparse_a/b/c/s) and tracks staleness across indexing, so recalculate=False can opt into carried-over (population-stale) stats instead of paying for a fresh O(nnz) pass. Matmul/matmat kernels are generalized to the same g/s parameterization so they keep working as a view for every recipe. Adds recipe-correctness tests against an independent numpy reference and benchmarks comparing the view path to materialize-then-multiply for every recipe (both timing and peak allocation), gated in baselines.json. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Replaces the per-case recorded-with-margin ceilings with a flat cutoff across every normalized-view benchmark, rather than one keyed to each case's specific measured ratio. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
For VCSC, a column's gsum (b) depends only on that column's own nonzeros, so it's already final by the time the g-transform pass over the same nonzeros needs it -- no need to wait for other columns. The refactor to per-recipe statistics had split this into two separate numba dispatches (_column_stats_major_is_col + _column_gstats_major_is_col), doubling parallel-region launch overhead for parafac2/pearson without doing any extra elementwise work. Fuses them back into one kernel, parameterized by which parts a given recipe actually needs (need_b/need_gstats), matching the pre-refactor VCSC fusion. VCSR keeps two separate passes, since gene_scale there isn't final until every row has scattered into it. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
from_stats() set self._dual_arr directly, but that attribute is only declared on the _VCSNormalizedBase subclass (in _vcs_norm.py), not on the NormalizedViewBase base class where from_stats() lives -- ty can't verify it exists on Self there. Replaces the direct attribute set with an _init_extra() hook (no-op on the base class, overridden by _VCSNormalizedBase to initialize _dual_arr), called from both __init__ and from_stats(). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Member
Author
|
@fishidaho interested in whether you see any gaps with this interface. |
`VCSCAnnData.normalized()` resolved its `view` argument to a `Recipe`, then handed `recipe.name` back down to `_VCSBase.normalized()`, which re-resolved that name through `RECIPES`. Any `Recipe` the caller built themselves is not in `RECIPES`, so the round trip raised `ValueError: unknown normalization view`, even though the same object works at the array level. Pass the resolved `Recipe` straight down instead, and widen both signatures to `str | Recipe` to say so. The `_norm_cache`/`_vcs_norm_cache` keys move from `recipe.name` to the resolved `Recipe` (a frozen, slotted dataclass, so hashable). Two distinct caller-built recipes are now free to share a `name` without silently handing back each other's cached view. Also validate `g_code` in `resolve_recipe()`: `_g()` falls through to the identity branch for anything it doesn't recognize, so a typo'd code would have silently produced an untransformed matrix rather than an error. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019ZwuJMcTRUacxqARwjMngE
`_norm_cache`/`_vcs_norm_cache` were plain dicts holding every computed view strongly, with no eviction. Two problems followed. A view can carry a `_dual_arr` -- a whole opposite-format copy of the array, cached by the matmul kernels whenever regrouping fits one chunk's budget. Because the cache held views strongly, dropping every reference to a view did not release its dual: the array kept it alive for its own lifetime. Touching all five built-in recipes on an 8k x 1.5k array and discarding each view left 5 duals resident, 8.4 MB against 2.3 MB of index+value bytes for the array itself. And nothing capped the number of entries, so a caller cycling through recipes accumulated one entry per distinct recipe forever. Replace both with `_NormCache`: an LRU bounded at `NORM_CACHE_MAXSIZE` (4) that holds each view *weakly* and its statistics strongly. The statistics are `O(n_rows + n_cols)` -- 0.55 MB per entry on a 60k x 3k array, measured equal to the predicted `8 * (n_rows + 3 * n_cols)` -- so retaining those is cheap, while nothing `O(nnz)` survives the caller dropping a view. `recalculate=False` keeps its contract: a hit on a live view returns that exact object, so identity is stable while the caller holds it; a hit on a collected one rebuilds from the retained statistics, still skipping the `O(nnz)` passes (1334x faster than recomputing, measured). That rebuild goes through a new `_from_internal_stats()` rather than `from_stats()`, because `a`/`b` are exposed as reciprocals of the internal `row_scale`/`gene_scale` and `1 / (1 / x)` is not exactly involutive in float64. Restored views are now bit-identical to a fresh recompute, which is asserted rather than approximated in the tests. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_019ZwuJMcTRUacxqARwjMngE
Every normalized-view benchmark compared `nv @ B` against `nv.toarray() @ B`
-- a dense materialization of the implicit-zero-filled matrix, which nobody
would actually perform. It is where the "20-100x faster, 150-200x less peak
memory" figures come from, and it flatters the view for a reason that has
nothing to do with the kernels: `time_ratio_view_over_materialize` clears its
0.3 ceiling no matter what the per-nonzero cost is.
Add the comparison the prototype benchmarks used instead. `_sparse_delta()`
builds `Delta` as a scipy CSR -- exactly what a `to_csr()` on the view would
return -- and the baseline multiplies that and adds the same rank-1
correction the view adds. Both sides then do identical arithmetic, asserted
to rtol=1e-9 inside the case, so the ratio isolates the kernel.
CPU is the metric that matters here, not wall: our kernels are
`parallel=True` against a single-threaded scipy, so wall time hides a
per-nonzero cost a shared workstation still pays. Measured with numba pinned
to one thread, on 60k x 2k @ 5% (6M nonzeros):
raw (identity) 1.53x scipy
pearson (sqrt) 1.89x
scanpy (log1p) 5.12x
cp10k_log1p (log1p) 5.10x
parafac2 (log10(1+1000x)) 5.87x
So the transform is essentially the whole cost, consistent with the 2.13x the
prototype measured at 336M nonzeros. `normalized_cpu_vs_sparse_1t` gates
these; the parallel wall/CPU ratios are recorded for context only, since they
depend on the runner's core count.
Three things the measurement needed to be reproducible:
- `benchmarks/run.py` gives any case whose name ends in `_1t` a subprocess
with `NUMBA_NUM_THREADS=1`. It has to be set before numba is imported --
capping the pool at runtime leaves the idle workers spinning and
`process_time()` counts them, which swung the ratio 2x between runs.
- 6M nonzeros rather than 2M, and `repeat=15`. The identity-transform kernels
land within ~1.5x of scipy, so shorter runs left jitter a large share of
the ratio.
- `gc.collect()` between recipes. Each builds its own ~150 MB delta, and
without it the recipes measured last ran against a progressively more
fragmented heap -- scanpy alone swung between 2.8x and 6.1x.
Ceilings are set from the observed maximum over five runs x 1.5 rather than
from one `--record` sample: the ratios have a low tail, and recording off one
of those puts the ceiling under the steady-state value. Verified clean over
four consecutive runs.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_019ZwuJMcTRUacxqARwjMngE
Contributor
|
This implementation matches the initial streaming prototype that I had been benchmarking against previously. The added three commits are fixes I may have found along the way and not anydisagreements with it. What changed1. It resolved mine = Recipe("mine", 1e4, False, G_LOG1P, True, True)
v.normalized(mine) # fine
va.normalized(mine) # ValueError: unknown normalization view 'mine'Now the resolved 2. The cache pinned
Replaced both with an LRU bounded at 3. Added a CPU gate against a sparse baseline ( The benchmarks compared Timing results
Parallel wall time, same comparison, 48 threads:
For matrix-vector products the view is slower than scipy on wall clock for all three log recipes. There isn't enough work per nonzero at k=1 to cover the transform, and the parallelism doesn't rescue it. The matmat story is good while the matvec story isn't, and randomized SVD does plenty of both (but its hard to say if they cancel). |
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
2 participants
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
Adapts
VCSCArrayNormalized/VCSRArrayNormalizedper fishidaho's proposal in #40: normalization as a declarative recipey[i,j] = (g(x[i,j] * a[i] * b[j]) - c[j]) * s[j], whereais a per-cell scale,b/c/sper-gene scale/center/post-scale, andga monotone transform withg(0) = 0.vsparse.RECIPES:raw,cp10k_log1p,parafac2(the existing default behavior),scanpy,pearson..normalized(view="parafac2", recalculate=True)—recalculate=Falsereuses a previously-built view for that recipe instead of recomputing itsO(nnz)statistics, so switching between recipes already computed this session is free.VCSCAnnData.normalized(...)additionally records the active recipe's statistics intoobs["vsparse_a"],varm["vsparse_b"/"c"/"s"], anduns["vsparse"](recipe name + astaleflag), so they persist throughwrite_h5ad/write_zarr. Indexing an AnnData carries these forward but marks themstale=True(population-level parts like the depth median or gene mean no longer reflect the subset);recalculate=Falsereuses them anyway,recalculate=True(default) recomputes and clears staleness.g/sparameterization, so@/toarray()/select()/__getitem__all keep working as views for every recipe, not just the original log10 transform.Test plan
test_vcs_norm_recipes.py: every recipe'stoarray()/matmul/rmatmul checked against an independent numpy reference;.normalized(view, recalculate=...)caching semantics;VCSCAnnDataobs/varm/uns storage and staleness across indexing.test_vcs_norm.py/test_norm_selection.py(defaultparafac2recipe) still pass unmodified.benchmarks/cases.py): for every recipe, matmat/matvec/rmatmat/rmatvec via the view vs. materialize-then-multiply — view is ~20-100x faster and ~150-200x less peak memory in every case, gated inbaselines.json.pytest -q→ 1585 passed, 97 skipped.ruff check/ruff format --checkclean.🤖 Generated with Claude Code