diff --git a/src/vsparse/_norm_common.py b/src/vsparse/_norm_common.py index 9b6d406..e5d8904 100644 --- a/src/vsparse/_norm_common.py +++ b/src/vsparse/_norm_common.py @@ -985,18 +985,50 @@ def slice_norms(self, condition_idxs: Any, n_cond: int) -> np.ndarray: # -- selection --------------------------------------------------------------- - def select(self, rows: Any = slice(None), cols: Any = slice(None)) -> Any: - """A normalized view of the selected sub-array, with statistics recomputed for it. + def select( + self, rows: Any = slice(None), cols: Any = slice(None), *, recalculate: bool = True + ) -> Any: + """A normalized view of the selected sub-array. Returns a view, not a dense array, so it still composes with - ``@``/:meth:`toarray`. + ``@``/:meth:`toarray` without ever materializing the selection. + + Parameters + ---------- + recalculate + If ``True`` (the default), statistics are recomputed fresh from + the selected sub-array -- a column selection then re-derives + read depth from only the selected columns, which is rarely what + a caller wants (select genes first and normalize after if it + matters). If ``False``, this view's *existing* statistics are + reused instead: ``row_scale``/``a`` sliced by ``rows`` (a + per-cell quantity, so subsetting is exact, not an + approximation), and the per-gene statistics (``gene_scale``/ + ``col_mean``/``col_post_scale``, i.e. ``b``/``c``/``s``) sliced + by ``cols`` unchanged. This is the right choice for evaluating a + model fit on a held-out slice against statistics computed on + the whole (or a different) array -- e.g. bi-cross-validation + scoring a test block against train-derived statistics -- and, + unlike bracket indexing (:meth:`__getitem__`), stays a lazy view + no matter how large the selection is: nothing gets materialized + until (and unless) the caller calls :meth:`toarray` or `@`s it, + and even then the O(nnz) work is bounded, not an eager dense + allocation sized to the selection. """ - # A column selection re-derives read depth from only the selected - # columns, which is rarely what a caller wants; select genes first - # and normalize after if it matters. - sub: Any = self._arr[_prep_key(rows), _prep_key(cols)] + row_key, col_key = _prep_key(rows), _prep_key(cols) + sub: Any = self._arr[row_key, col_key] if not isinstance(sub, type(self._arr)): sub = type(self._arr).from_scipy(sub) + if not recalculate: + return type(self).from_stats( + sub, + self.recipe, + np.asarray(self.a)[row_key], + np.asarray(self.b)[col_key], + np.asarray(self.c)[col_key], + np.asarray(self.s)[col_key], + stale=self.stale, + ) return type(self)(sub, self.recipe) # -- on-the-fly elementwise access ------------------------------------------ diff --git a/tests/test_norm_selection.py b/tests/test_norm_selection.py index bc4921f..48c9cd1 100644 --- a/tests/test_norm_selection.py +++ b/tests/test_norm_selection.py @@ -102,6 +102,58 @@ def test_select_returns_a_view_that_still_composes(vcls): np.testing.assert_allclose(sub @ B, _reference(dense[mask]) @ B, atol=1e-8) +# -- select(recalculate=False): a window that stays a lazy view ------------- + + +def test_select_no_recalculate_matches_getitem_values(vcls): + """select(recalculate=False) keeps the parent's statistics, like __getitem__.""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + + want = np.asarray(nv[mask, :]) + got = nv.select(mask, recalculate=False).toarray() + + np.testing.assert_allclose(got, want, atol=1e-10) + + +def test_select_no_recalculate_returns_a_lazy_view(vcls): + """Unlike __getitem__, this never eagerly materializes the selection.""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + + sub = nv.select(mask, recalculate=False) + assert isinstance(sub, _norm_cls(vcls)) + assert sub.shape == (int(mask.sum()), dense.shape[1]) + + rng = np.random.default_rng(5) + B = rng.normal(size=(dense.shape[1], 4)) + Bl = rng.normal(size=(3, int(mask.sum()))) + want = np.asarray(nv[mask, :]) + np.testing.assert_allclose(sub @ B, want @ B, atol=1e-8) + np.testing.assert_allclose(Bl @ sub, Bl @ want, atol=1e-8) + + +def test_select_no_recalculate_composes_with_column_selection_too(vcls): + """A column selection just slices the existing per-gene stats, not recomputed.""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + cols = np.array([0, 3, 7, 15, 40]) + + want = np.asarray(nv[mask, :])[:, cols] + got = nv.select(mask, cols, recalculate=False).toarray() + np.testing.assert_allclose(got, want, atol=1e-10) + + +def test_select_no_recalculate_default_recalculate_still_recomputes(vcls): + """Bare select(mask) is unaffected by the new keyword-only argument.""" + dense, mask = _mixed_population() + nv = vcls.from_scipy(_scipy_for(vcls, dense)).normalized() + + np.testing.assert_allclose( + nv.select(mask).toarray(), nv.select(mask, recalculate=True).toarray(), atol=1e-12 + ) + + # -- __getitem__: a window that keeps the parent's statistics ----------------