Add per-axis sum/mean/max/min and elementwise arithmetic by aarmey · Pull Request #22 · meyer-lab/vsparse

aarmey · 2026-09-02T18:19:53Z

Summary

Add mean(axis), max(axis), min(axis) to _VCSBase (shared by VCSCArray/VCSRArray), wired through the existing major/minor-axis reduction split used by sum(axis). max/min correctly account for implicit (structural) zeros, matching scipy's sparse-array semantics.
Add elementwise arithmetic: __add__/__radd__/__sub__/__rsub__ (scalar 0 no-ops, nonzero scalar raises NotImplementedError like scipy does) and extend __mul__/.multiply()/__truediv__ beyond scalars to dense arrays and other VCS/scipy sparse arrays. All fall back to scipy internally to compute the result, but re-wrap a sparse result back into the same VCSCArray/VCSRArray class rather than returning a raw scipy array.
Fixed a latent bug found while extending __mul__/__truediv__: passing an unsupported operand type (e.g. a dict) previously produced np.asarray(NotImplemented) instead of letting Python raise TypeError — _elementwise now propagates NotImplemented correctly.
Updated test_unsupported_scalar_operands_raise in tests/test_ops.py, whose old expectations (mismatched-shape operands silently returning NotImplemented) no longer apply now that non-scalar operands are handled elementwise (with broadcasting, matching numpy/scipy).

Test plan

uv run pytest -q — 907 passed, 37 skipped
uv run ruff check .
uv run ty check

VCSCArray/VCSRArray only supported scalar-only mul/truediv and a global sum(); everything else (per-axis reductions, add/sub, and array-array/array-dense multiply) fell back to converting to scipy and losing the VCS type. This wires up mean/max/min per-axis (matching scipy's implicit-zero-aware semantics) and add/sub/multiply against scalars, dense arrays, and other VCS/scipy sparse arrays, re-wrapping sparse results back into the same VCS class. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

#22/#23 added astype, native minor-axis selection, and elementwise add/sub/multiply (which rebuild through scipy). Every one of them constructs its result with `type(self)(...)`, so all of them inherit the narrowing in `_VCSBase.__init__` for free -- no per-method change was needed, which is the payoff for putting the rule at the construction choke point rather than in `from_scipy`. Verified rather than assumed: an int64-indexed input stays int32 through astype, `v[:, cols]`, both-axes selection, scalar mul/div/neg, add/sub against another VCS array, copy, log1p, `_transpose_major` and `T`. The scipy round-trip in the elementwise path is the one that could plausibly have handed back int64, so it gets its own case. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

#22 added per-axis max/min and getnnz. Both take the same shape as the `_minor_sums` this branch was already fixing, and both were written the same way -- so the defect this PR exists to remove arrived in three new places at once: _minor_reduce np.repeat(values, group_sizes) then ufunc.at 16.07 B/nnz _minor_nnz np.bincount(int32 indices) promotes to intp 8.00 B/nnz Leaving those while fixing `sum` would make the PR incoherent, so they get the same treatment. `_ops.minor_extrema` computes the extremum over stored values and the per-index count in one parallel pass over the value groups; `_ops.minor_counts` does the count alone. Measured on 2e6 nonzeros, against the implementations they replace: max(axis=0) 32.01 MB -> 0.39 MB (81x) 15.3x faster min(axis=0) 32.01 MB -> 0.39 MB (81x) getnnz(axis=0) 16.00 MB -> 0.20 MB (81x) 8.7x faster Output is identical in every case, and matches a dense reference. #22's own tests for these methods pass unchanged -- the implicit-zero correction stays in `_minor_reduce` rather than moving into the kernel, so the semantics those tests pin are untouched. The accumulator budget now takes a per-element size, since the extrema kernel keeps an extremum *and* a count per slot: 64 MiB still bounds the block, it just buys fewer threads for a wider accumulator. Renames tests/test_reductions.py to test_reduction_memory.py -- #22 added test_reductions_and_arith.py for the semantics, and two files a letter apart covering different concerns is a trap. The docstring now says which is which. Not fixed here: `_select_minor` (new in #23) has the same nnz-sized temporary, at 15.07 B/nonzero. It's an indexing path rather than a reduction, so it's filed separately as ISSUE-30 instead of widening this PR. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Clean merge. #23 touched _anndata_class only in _subset_2d's comment; this branch touches _write_group/write_h5ad/write_zarr, so the two don't overlap. Nothing in #22/#23 addresses obs/var metadata encoding, and the categorical conversion is unaffected by the indexing changes.

…d main Three new fast cases: minor-axis max/min, per-minor-index getnnz, and minor-axis selection. All three are new surface from #22/#23, and all three shipped carrying an nnz-sized temporary -- which is the best argument this gate could have for existing. The pattern this suite was built to catch reappeared in new code within the same week, unnoticed, because nothing measured it: minor_extrema_peak_mb 66.27 MB (max/min: np.repeat + ufunc.at) minor_getnnz_peak_mb 32.02 MB (np.bincount promotes int32 -> intp) minor_selection_peak_mb 62.27 MB (_select_minor, ISSUE-30) for results of length n_minor, on a 4M-nonzero array. Baselines re-recorded against merged main so the gate reflects the code it now guards. The ceilings for the first two drop by ~80x once the reduction kernels land; the third stays until ISSUE-30 is fixed. As before these gate against getting worse, not against the current numbers being good -- README.md already says so and now names all five memory cases. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* Size index arrays by the axis they address, not the array beside them `indices` is the only nnz-sized array in the VCS layout, and nothing ever chose its dtype: `_construct.compress` copied whatever dtype the input scipy array carried, and `write_ivcs_elem` recorded the same. scipy hands out int64 indices for any array with enough nonzeros, so a 33k-gene minor axis was routinely stored -- in memory, on disk, and in every kernel that walks it -- at 8 bytes per nonzero instead of 4. `_rapid_load._filter_and_compact` had the same conflation in a sharper form: one `idx_dtype`, keyed off nnz, applied to both `new_indptr` (indexed by nonzero count, genuinely needs int64 at scale) and `out_indices` (gene indices, bounded by the gene count). Crossing INT32_MAX nonzeros silently doubled the largest allocation in the function for no reason. Adds `_indexutils.smallest_index_dtype(n)` as the single place that rule lives, and applies it at each point an index array is sized: - `_construct.compress` takes `n_minor` and narrows up front, so the wide buffer is never allocated rather than allocated and then shrunk. - `_VCSBase.__init__` narrows as the construction choke point -- after the bounds check, so an out-of-range index is still rejected rather than truncated, and a no-op (no copy) when the dtype is already right. - `write_ivcs_elem` re-derives the dtype a reader will rebuild `indices` as, instead of trusting the array it was handed. - `_filter_and_compact` keys its pointer and column-index dtypes off nnz and the kept-gene count separately. - `transpose_major` now uses the shared helper for the rule it already applied inline. Narrowing never widens: indices already stored in something smaller than int32 are left alone. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Extend dtype-narrowing coverage to main's new construction paths #22/#23 added astype, native minor-axis selection, and elementwise add/sub/multiply (which rebuild through scipy). Every one of them constructs its result with `type(self)(...)`, so all of them inherit the narrowing in `_VCSBase.__init__` for free -- no per-method change was needed, which is the payoff for putting the rule at the construction choke point rather than in `from_scipy`. Verified rather than assumed: an int64-indexed input stays int32 through astype, `v[:, cols]`, both-axes selection, scalar mul/div/neg, add/sub against another VCS array, copy, log1p, `_transpose_major` and `T`. The scipy round-trip in the elementwise path is the one that could plausibly have handed back int64, so it gets its own case. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Trim comments and tests Comments now state what the code does rather than the reasoning behind it. Index-dtype tests drop the cases the code trivially guarantees, keeping the boundary, the truncation risks, the roundtrips, and the paths that build a new array. --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

* Categorical-encode obs/var strings on write, as anndata's writers do `numeric_only_compression` exists because Blosc2 crashes on variable-length-string HDF5 datasets, so every string column in obs/var is written uncompressed. That left the package with no size strategy for metadata at all. Investigating what to do about it turned up a plain bug rather than a codec question. `anndata.AnnData.write_h5ad` converts string columns to categoricals before writing (its own `convert_strings_to_categoricals=True`). `VCSCAnnData` writes field-by-field rather than delegating -- deliberately, for the reasons in `_write_group`'s comment -- and so never did. A low-cardinality annotation (cell type, sample ID, batch) therefore landed as one variable-length string per row, uncompressed and unrecoverable. On 200k cells x 2k genes with three such columns: before (plain strings) file 41.72 MB after (categorical) file 13.48 MB Adds the same parameter with the same name and default to `write_h5ad` and `write_zarr`. anndata only converts columns with fewer categories than rows, so a per-row-unique column (barcodes, gene symbols) is left alone and this can never make a column larger. The codec question was answered too, by writing a 20k-element vlen string dataset under each filter, one subprocess each, against h5py 3.16.0 / HDF5 2.0.0 / hdf5plugin 7.0.0: none, gzip and lzf all fine; blosc2 still dies with SIGFPE. So the workaround stays, and it's specifically Blosc2 rather than HDF5 filters in general. gzip/lzf would be safe but aren't worth adopting -- vlen payloads live in HDF5's global heap where per-dataset compression doesn't reach them well, and the categorical encoding above is worth far more than any string codec could be. Both findings are recorded in `_compression`'s module docstring. No upstream issue filed yet: reproducing this outside the h5py/hdf5plugin combination in use here needs a check against a current hdf5plugin build first. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Merge origin/main into investigate/metadata-compression Clean merge. #23 touched _anndata_class only in _subset_2d's comment; this branch touches _write_group/write_h5ad/write_zarr, so the two don't overlap. Nothing in #22/#23 addresses obs/var metadata encoding, and the categorical conversion is unaffected by the indexing changes. * Trim comments and docstrings --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

* Add a benchmark suite and wire it into CI as a regression gate There was no benchmark suite and no CI job that would notice a cost regression. Correctness tests don't help here: the answers stay right while the memory or the file size quietly doubles, which is exactly the class of defect the v0.2 work is fixing. `benchmarks/` holds a harness, five fast cases run on every PR, one larger case for a scheduled or manual run, and checked-in ceilings. Two decisions worth stating, since a flaky or meaningless gate is worse than none: Timing is never recorded as absolute seconds. Each timing case runs the same work through scipy in the same process and records the *ratio*, which cancels most of the difference between a laptop and a shared runner, and is gated loosely (4x) because it's still the noisy one. The sharp gates are the deterministic numbers: bytes per stored nonzero (1.1x) and memory allocated by an operation (2x). Memory is measured with tracemalloc, not ru_maxrss. RSS is a process-lifetime high-water mark, so an operation staying under the peak set while building its input reports zero however much it allocates -- the first version of this suite duly reported 0 MB for an operation allocating 66 MB. tracemalloc measures allocations and resets between runs. Recorded against main, the two memory cases show what the rest of v0.2 is about: a minor-axis reduction on a 4M-nonzero array allocates 66 MB, and a misaligned-direction matmul allocates 204 MB against a 16 MB array. The ceilings are therefore deliberately generous today and should be re-recorded once those fixes land; README.md says so and names the numbers. Each case runs in its own subprocess, since measurement state and JIT warm-up would otherwise leak between them. The fast set takes ~28s. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Cover the minor-axis operations #22/#23 added, and re-record on merged main Three new fast cases: minor-axis max/min, per-minor-index getnnz, and minor-axis selection. All three are new surface from #22/#23, and all three shipped carrying an nnz-sized temporary -- which is the best argument this gate could have for existing. The pattern this suite was built to catch reappeared in new code within the same week, unnoticed, because nothing measured it: minor_extrema_peak_mb 66.27 MB (max/min: np.repeat + ufunc.at) minor_getnnz_peak_mb 32.02 MB (np.bincount promotes int32 -> intp) minor_selection_peak_mb 62.27 MB (_select_minor, ISSUE-30) for results of length n_minor, on a 4M-nonzero array. Baselines re-recorded against merged main so the gate reflects the code it now guards. The ceilings for the first two drop by ~80x once the reduction kernels land; the third stays until ISSUE-30 is fixed. As before these gate against getting worse, not against the current numbers being good -- README.md already says so and now names all five memory cases. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Document the measured before/after for every gated memory case Names all five memory cases with the number on main, the number with the corresponding fix, and the reason each is high -- so the generous ceilings are self-explaining rather than looking like sloppy thresholds, and so re-recording after each fix is a mechanical check rather than a judgement. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Trim comments, docstrings and the README Drops the dead peak_rss_mb helper and cuts the prose to what the suite does and how to run it. --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

* Reduce over the minor axis without an nnz-sized temporary `_VCSBase._minor_sums` expanded the value-compressed layout back to one float64 per nonzero before reducing it: expanded = np.repeat(self.values.astype(np.float64), group_sizes) return np.bincount(self.indices, weights=expanded, minlength=self.n_minor) That `np.repeat` is scratch space for a reduction that never needs to keep it -- and it undoes, transiently, the exact compression the layout exists to provide. Any `sum(axis=...)` over the minor axis pays it, including the one `NormalizedViewBase.__init__` makes on every view construction. Replaces it with a parallel numba scatter (`_ops.minor_sums`) that walks the value groups directly. Group index ranges are disjoint, so threads take contiguous blocks of groups and accumulate into thread-local rows, summed at the end -- no write hazard, and nothing nnz-sized is ever allocated. The thread-local accumulators are the obvious way to reintroduce the same problem in a new shape (`nthreads * n_minor * 8` bytes), so the thread count is capped to keep that block under a fixed 64 MiB budget: a minor axis wide enough that even two accumulators would blow it runs serially instead. Measured on 1e6 nonzeros, peak allocation drops from 16.00 to 0.20 bytes per nonzero; on 4e7 nonzeros the reduction also runs 11x faster (0.139s -> 0.012s), since it moves far less memory and now uses every core. Also adds the direct `sum()` coverage the axis reductions never had -- both axes against a dense reference, the all-zero case, and a regression test asserting the call path allocates nothing nnz-sized. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Extend the scatter kernels to max/min/getnnz on the minor axis #22 added per-axis max/min and getnnz. Both take the same shape as the `_minor_sums` this branch was already fixing, and both were written the same way -- so the defect this PR exists to remove arrived in three new places at once: _minor_reduce np.repeat(values, group_sizes) then ufunc.at 16.07 B/nnz _minor_nnz np.bincount(int32 indices) promotes to intp 8.00 B/nnz Leaving those while fixing `sum` would make the PR incoherent, so they get the same treatment. `_ops.minor_extrema` computes the extremum over stored values and the per-index count in one parallel pass over the value groups; `_ops.minor_counts` does the count alone. Measured on 2e6 nonzeros, against the implementations they replace: max(axis=0) 32.01 MB -> 0.39 MB (81x) 15.3x faster min(axis=0) 32.01 MB -> 0.39 MB (81x) getnnz(axis=0) 16.00 MB -> 0.20 MB (81x) 8.7x faster Output is identical in every case, and matches a dense reference. #22's own tests for these methods pass unchanged -- the implicit-zero correction stays in `_minor_reduce` rather than moving into the kernel, so the semantics those tests pin are untouched. The accumulator budget now takes a per-element size, since the extrema kernel keeps an extremum *and* a count per slot: 64 MiB still bounds the block, it just buys fewer threads for a wider accumulator. Renames tests/test_reductions.py to test_reduction_memory.py -- #22 added test_reductions_and_arith.py for the semantics, and two files a letter apart covering different concerns is a trap. The docstring now says which is which. Not fixed here: `_select_minor` (new in #23) has the same nnz-sized temporary, at 15.07 B/nonzero. It's an indexing path rather than a reduction, so it's filed separately as ISSUE-30 instead of widening this PR. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Bound the per-gene accumulator blocks in the loader too `_weighted_bincount` and `_gene_detection_counts` size their thread-local block by the full thread count, so it grows with both the feature axis and the machine: 96 MB at 250k features on 48 threads. They now share the same cap as the other scatter kernels. `accumulator_threads` loses its underscore, since it is no longer local to one module. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> * Trim comments and tests, and collapse the kernel wrappers minor_sums and minor_counts were one-line wrappers that only supplied the thread count, so the kernels now take the public name and callers pass it. minor_extrema keeps a wrapper because it reduces the per-thread partials. Tests drop the internal-call and sweep cases, keeping the dense comparisons, the implicit-zero and integer-sentinel edges, the float64 accumulator, the budget, and the memory bound. --------- Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> Co-authored-by: Aaron Meyer <ameyer@ucla.edu>

aarmey merged commit 76db17a into mainSep 2, 2026
5 checks passed

aarmey deleted the add-per-axis-reductions-and-elementwise-arith branch September 2, 2026 18:21

Navigation Menu

Uh oh!

Add per-axis sum/mean/max/min and elementwise arithmetic - #22

Add per-axis sum/mean/max/min and elementwise arithmetic#22
aarmey merged 1 commit into
mainfrom
add-per-axis-reductions-and-elementwise-arith

aarmey commented Sep 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

aarmey commented Sep 2, 2026

Summary

Test plan

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant