Add select(recalculate=False): a lazy, stats-preserving row/col selection

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 ------------------------------------------

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 ----------------