Skip to content

Add general fancy/boolean indexing, per-axis getnnz, and astype - #23

Merged
aarmey merged 3 commits into
mainfrom
general-indexing-nnz-astype
Sep 2, 2026
Merged

Add general fancy/boolean indexing, per-axis getnnz, and astype#23
aarmey merged 3 commits into
mainfrom
general-indexing-nnz-astype

Conversation

@aarmey

@aarmey aarmey commented Sep 2, 2026

Copy link
Copy Markdown
Member

Summary

  • Add _select_minor, mirroring _select_major but for the minor axis (rows for VCSC, columns for VCSR): filters and remaps indices, drops any (major, unique-value) slot left with zero kept indices, and shrinks major_ptr/value_ptr accordingly.
  • Rewire __getitem__ to compose _select_major/_select_minor for any combination of slices/boolean masks/integer arrays/bare ints on either axis. Only the case where both axes are a bare int (a genuine scalar result, which a 2-D VCS array can't represent) still converts via to_scipy(). arr[:, :] now returns self.copy() instead of silently degrading to a raw scipy array.
  • Add getnnz(axis=None) / count_nonzero(), reusing the major/minor reduction split already built for sum/max/min. count_nonzero() correctly excludes stored-but-zero values, unlike nnz/getnnz().
  • Add astype(dtype, copy=True).
  • Updated the now-stale comment in _anndata_class.py's _subset_2d (the re-wrap-to-VCS-type fallback it describes is now unreachable from VCSCAnnData.__getitem__, since _as_slice_index already converts bare ints to slices before calling in — kept as a defensive no-op) and renamed/updated test_general_2d_indexing_falls_back in tests/test_indexing.py since it no longer falls back.

Test plan

  • uv run pytest -q — 1066 passed, 59 skipped
  • uv run ruff check .
  • uv run ty check

Aaron Meyer and others added 3 commits September 2, 2026 11:19
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>
__getitem__ previously only had a fast native path for a full minor
slice + arbitrary major selection; everything else (minor-only
selection, or both axes at once) fell back to converting through
to_scipy(). Add _select_minor (mirrors _select_major but filters and
remaps `indices`, dropping now-empty unique-value slots) and compose
it with _select_major so only the true both-axes-scalar case still
needs to convert.

Also add getnnz(axis)/count_nonzero() (reusing the major/minor
reduction split from sum/max/min) and astype(dtype, copy), rounding
out more of scipy's csr_array/csc_array surface natively.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@aarmey
aarmey merged commit 0711242 into main Sep 2, 2026
5 checks passed
@aarmey
aarmey deleted the general-indexing-nnz-astype branch September 2, 2026 20:14
fishidaho added a commit that referenced this pull request Sep 3, 2026
#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>
fishidaho added a commit that referenced this pull request Sep 3, 2026
#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>
fishidaho added a commit that referenced this pull request Sep 3, 2026
#23 changed what a raw array's __getitem__ returns underneath this branch,
in two ways that matter here:

- A minor-axis selection now stays VCS-native instead of falling out as a
  scipy array. `select`'s re-wrap is therefore no longer the normal path,
  just a defensive one -- reworded to match the same defensive re-wrap in
  `_anndata_class._subset_2d`, and the test comment claiming a VCSCArray row
  selection "has no .normalized() at all" is simply no longer true.

- `arr[rows, cols]` with two index arrays now composes a major- and a
  minor-axis selection, which is the sub-block semantics `select` wants. It
  previously fell through to scipy, which broadcasts two index arrays
  pointwise, and that's why this was written one axis at a time. That
  workaround is now redundant, so it collapses to a single index operation.

Adds a test pinning the outer-vs-pointwise distinction directly, using row
and column selections of different lengths -- a pointwise broadcast would
fail outright rather than return something subtly wrong, which is what makes
this worth keeping a test on.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
fishidaho added a commit that referenced this pull request Sep 3, 2026
The one conflict in the wave, and a benign one: this branch added
`_major_range` immediately before `__getitem__`, #23 added `_select_minor`
in the same place and rewrote `__getitem__`. The two methods are independent
-- a contiguous major range that returns views, versus an arbitrary
minor-axis selection that filters and remaps indices -- so both are kept,
and `__getitem__` is taken from main unchanged.

They now overlap in purpose for one case: `__getitem__` has a native path
for a contiguous slice, which `_major_range` also covers. `_major_range`
stays, because it is the reason the chunked matmul is affordable --
`values`/`indices` come back as views into the parent's buffers, where
`_select_major` gathers. Verified still true after the merge
(`np.shares_memory` on both), and a test now pins the two against each
other: if they ever disagree, the chunked matmul is computing against a
different sub-array than a caller would get.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
fishidaho added a commit that referenced this pull request Sep 3, 2026
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.
fishidaho added a commit that referenced this pull request Sep 3, 2026
…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>
aarmey pushed a commit that referenced this pull request Sep 4, 2026
* 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>
aarmey pushed a commit that referenced this pull request Sep 4, 2026
* Make post-selection normalization explicit and reachable

A normalized view's statistics are computed once, at construction, over the
whole array it wrapped. `__getitem__` then applies those parent statistics
to whatever sub-block is asked for -- while its docstring described the
result as though the block were self-consistently normalized ("the
transform/centering formula is then applied to that small block directly").

It isn't, and the gap is not small. On a mixed population -- two cell types
differing in read depth and marker expression -- indexing the view for one
cell type differs from normalizing those cells by **52% relative Frobenius
norm**. Read as "the normalized data for these cells", that silently fits
downstream analysis to different data than the caller believes.

Rather than redefine what indexing returns, this keeps the two operations
and names them:

- `__getitem__` stays a window into this view's matrix -- exactly
  `toarray()[key]` without materializing the full matrix -- and now says so,
  including what it is *not*.
- `select(rows, cols)` is new: it renormalizes the selected sub-array on its
  own terms, equivalent to `arr[sel].normalized()`, and returns a view so it
  still composes with `@`/`toarray()`.

Recomputation was deliberately not made the default for indexing, because
it is not uniformly more correct: `row_scale` is a per-cell total over the
columns present, so recomputing after a *gene* selection re-derives read
depth from just those genes, which is usually wrong. Making the caller ask
for it keeps that choice explicit; `select`'s docstring spells the caveat
out. Choosing per-selection semantics properly is v0.4's selection-algebra
work.

`select` also fills a real gap: selecting rows of a VCSCArray is a
minor-axis selection, which drops out as scipy with no `.normalized()` at
all, so the raw-array route isn't uniformly available.

Tests pin both contracts against a dense reference, and assert the two
disagree by >10% on a realistic selection -- so the distinction can't be
quietly collapsed by a later change.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* Simplify select() onto #23's native two-axis indexing

#23 changed what a raw array's __getitem__ returns underneath this branch,
in two ways that matter here:

- A minor-axis selection now stays VCS-native instead of falling out as a
  scipy array. `select`'s re-wrap is therefore no longer the normal path,
  just a defensive one -- reworded to match the same defensive re-wrap in
  `_anndata_class._subset_2d`, and the test comment claiming a VCSCArray row
  selection "has no .normalized() at all" is simply no longer true.

- `arr[rows, cols]` with two index arrays now composes a major- and a
  minor-axis selection, which is the sub-block semantics `select` wants. It
  previously fell through to scipy, which broadcasts two index arrays
  pointwise, and that's why this was written one axis at a time. That
  workaround is now redundant, so it collapses to a single index operation.

Adds a test pinning the outer-vs-pointwise distinction directly, using row
and column selections of different lengths -- a pointwise broadcast would
fail outright rather than return something subtly wrong, which is what makes
this worth keeping a test on.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* Trim comments and tests

Docstrings state what the code does. The two both-axes selection tests fold
into one that uses different-length keys, which is what actually pins outer
semantics.

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
aarmey pushed a commit that referenced this pull request Sep 4, 2026
* 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>
aarmey pushed a commit that referenced this pull request Sep 4, 2026
* Fix silent data loss when a minor-axis selection repeats an index

`arr[:, [1, 1, 3]]` returns zeros where column 1's second copy should be.
No error, no warning -- just wrong values, in a shape that looks right.

    v[:, [1, 1, 3]]        want              got
                        [[ 2  2  4]      [[ 0  2  4]
                         [ 6  6  8]       [ 0  6  8]
                         [10 10 12]]      [ 0 10 12]]

This is a regression. Before #23 the same expression fell through to
`to_scipy()[...]`, and scipy fans duplicate fancy indices out correctly;
verified against 45ce15d, where it returns the right answer.

The cause is structural rather than an off-by-one: `_select_minor` builds a
single old -> new lookup array,

    remap = np.full(self.n_minor, -1); remap[idx] = np.arange(len(idx))

and a repeated index writes that slot twice, so only its last destination
survives. One array cannot express one-to-many, so the fix is to invert the
selection instead: `fanout`/`offsets`/`positions` give, for each original
minor index, every output position it maps to, and each stored element emits
one entry per destination. Two passes (count, then fill) via
`_ops.minor_select_counts`/`minor_select_fill`, so each surviving slot writes
a disjoint range and the passes parallelize.

Deriving each major slice's slot count with `np.searchsorted` on `major_ptr`
also drops the `np.repeat(np.arange(n_unique), ...)` this used to build,
which was `nnz`-sized. Peak allocation for a 2e6-nonzero selection falls
from 62.27 MB to 4.78 MB, so this closes ISSUE-30 as well -- not a separate
concern, just what the correct implementation happens not to need.

All 1066 tests from main pass unchanged. New coverage coming from the
fan-out semantics: five duplicate patterns, duplicates on either axis and
both at once, 25 random selections with repeats compared against scipy
directly, plus the selections that already worked (sorted, reversed, empty,
negative, slices, boolean masks) so the fix can't silently narrow them.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* Trim comments and tests

The fan-out explanation moves from a block comment into short notes inside
the kernels. Duplicate-pattern sweep drops to three cases and the type/dtype
assertions go, since construction guarantees them.

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
aarmey pushed a commit that referenced this pull request Sep 4, 2026
* 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>
aarmey added a commit that referenced this pull request Sep 4, 2026
* 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>
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