Generalize normalized views into declarative normalization recipes

benchmarks/baselines.json

@@ -5,7 +5,14 @@
     "indices_bytes_per_nonzero": 1.1,
     "vs_scipy_ratio": 1.1,
     "peak_alloc_mb": 2.0,
-    "time_ratio_vs_scipy": 4.0
+    "time_ratio_vs_scipy": 4.0,
+    "peak_alloc_mb_view": 2.0,
+    "time_ratio_view_over_materialize": 4.0,
+    "cpu_ratio_1t_cp10k_log1p": 1.5,
+    "cpu_ratio_1t_parafac2": 1.5,
+    "cpu_ratio_1t_pearson": 1.5,
+    "cpu_ratio_1t_raw": 1.5,
+    "cpu_ratio_1t_scanpy": 1.5
   },
   "cases": {
     "layout_bytes_per_nonzero": {
@@ -33,6 +40,123 @@
     },
     "minor_selection_peak_mb": {
       "peak_alloc_mb": 124.5457
+    },
+    "normalized_cp10k_log1p_matmat_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 2.7262
+    },
+    "normalized_cp10k_log1p_matvec_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 0.3858
+    },
+    "normalized_cp10k_log1p_rmatmat_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 3.6521
+    },
+    "normalized_cp10k_log1p_rmatvec_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 0.1321
+    },
+    "normalized_parafac2_matmat_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 2.7262
+    },
+    "normalized_parafac2_matvec_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 0.3858
+    },
+    "normalized_parafac2_rmatmat_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 3.6521
+    },
+    "normalized_parafac2_rmatvec_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 0.1321
+    },
+    "normalized_pearson_matmat_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 2.7262
+    },
+    "normalized_pearson_matvec_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 0.3858
+    },
+    "normalized_pearson_rmatmat_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 3.6521
+    },
+    "normalized_pearson_rmatvec_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 0.1321
+    },
+    "normalized_raw_matmat_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 2.7262
+    },
+    "normalized_raw_matvec_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 0.3858
+    },
+    "normalized_raw_rmatmat_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 3.6521
+    },
+    "normalized_raw_rmatvec_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 0.1321
+    },
+    "normalized_scanpy_matmat_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 2.7262
+    },
+    "normalized_scanpy_matvec_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 0.3858
+    },
+    "normalized_scanpy_rmatmat_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 3.6521
+    },
+    "normalized_scanpy_rmatvec_vs_materialize": {
+      "time_ratio_view_over_materialize": 0.3,
+      "peak_alloc_mb_view": 0.1321
+    },
+    "normalized_cp10k_log1p_matmat_vs_sparse": {
+      "peak_alloc_mb_view": 2.7262
+    },
+    "normalized_cp10k_log1p_matvec_vs_sparse": {
+      "peak_alloc_mb_view": 0.3858
+    },
+    "normalized_parafac2_matmat_vs_sparse": {
+      "peak_alloc_mb_view": 2.7262
+    },
+    "normalized_parafac2_matvec_vs_sparse": {
+      "peak_alloc_mb_view": 0.3858
+    },
+    "normalized_pearson_matmat_vs_sparse": {
+      "peak_alloc_mb_view": 2.7262
+    },
+    "normalized_pearson_matvec_vs_sparse": {
+      "peak_alloc_mb_view": 0.3858
+    },
+    "normalized_raw_matmat_vs_sparse": {
+      "peak_alloc_mb_view": 2.7262
+    },
+    "normalized_raw_matvec_vs_sparse": {
+      "peak_alloc_mb_view": 0.3858
+    },
+    "normalized_scanpy_matmat_vs_sparse": {
+      "peak_alloc_mb_view": 2.7262
+    },
+    "normalized_scanpy_matvec_vs_sparse": {
+      "peak_alloc_mb_view": 0.3858
+    },
+    "normalized_cpu_vs_sparse_1t": {
+      "cpu_ratio_1t_cp10k_log1p": 7.815,
+      "cpu_ratio_1t_parafac2": 8.91,
+      "cpu_ratio_1t_pearson": 2.85,
+      "cpu_ratio_1t_raw": 2.445,
+      "cpu_ratio_1t_scanpy": 7.845
     }
   }
 }

benchmarks/cases.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
+import gc
 from collections.abc import Callable
 
 import numpy as np
 
 from benchmarks.harness import (
+    best_cpu_time,
+    best_time,
     integer_counts_csr,
     peak_alloc_mb,
     ratio_vs_scipy,
@@ -139,6 +142,215 @@ def matmat_vs_scipy() -> dict[str, float]:
     return {"time_ratio_vs_scipy": ratio_vs_scipy(lambda: v @ B, lambda: mat @ B)}
 
 
+# -- normalized views (issue #40 recipes): view-op vs materialize-then-op ---
+#
+# For every recipe, the view-based matmul/matvec should cost less, both in
+# time and in peak allocation, than fully materializing the (dense,
+# implicit-zero-filling) normalized matrix and multiplying that -- the whole
+# point of a *view*. ``time_ratio_view_over_materialize`` < 1 and
+# ``peak_alloc_mb_view`` < ``peak_alloc_mb_materialize`` are the expectation
+# for every case below.
+
+
+def _normalized_bench(recipe: str, *, vector: bool) -> Callable[[], dict[str, float]]:
+    def bench() -> dict[str, float]:
+        from vsparse import VCSRArray
+
+        mat = integer_counts_csr(20_000, 2_000, density=0.05)
+        v = VCSRArray.from_scipy(mat)
+        nv = v.normalized(recipe)
+        rng = np.random.default_rng(0)
+        B = rng.normal(size=mat.shape[1]) if vector else rng.normal(size=(mat.shape[1], 8))
+
+        def via_view() -> np.ndarray:
+            return nv @ B
+
+        def via_materialize() -> np.ndarray:
+            return nv.toarray() @ B
+
+        return {
+            "time_ratio_view_over_materialize": best_time(via_view) / best_time(via_materialize),
+            "peak_alloc_mb_view": peak_alloc_mb(via_view),
+            "peak_alloc_mb_materialize": peak_alloc_mb(via_materialize),
+        }
+
+    bench.__name__ = f"normalized_{recipe}_{'matvec' if vector else 'matmat'}_vs_materialize"
+    return bench
+
+
+def _normalized_rbench(recipe: str, *, vector: bool) -> Callable[[], dict[str, float]]:
+    def bench() -> dict[str, float]:
+        from vsparse import VCSCArray
+
+        mat = integer_counts_csr(20_000, 2_000, density=0.05)
+        v = VCSCArray.from_scipy(mat)
+        nv = v.normalized(recipe)
+        rng = np.random.default_rng(0)
+        B = rng.normal(size=mat.shape[0]) if vector else rng.normal(size=(8, mat.shape[0]))
+
+        def via_view() -> np.ndarray:
+            return B @ nv
+
+        def via_materialize() -> np.ndarray:
+            return B @ nv.toarray()
+
+        return {
+            "time_ratio_view_over_materialize": best_time(via_view) / best_time(via_materialize),
+            "peak_alloc_mb_view": peak_alloc_mb(via_view),
+            "peak_alloc_mb_materialize": peak_alloc_mb(via_materialize),
+        }
+
+    bench.__name__ = f"normalized_{recipe}_{'rmatvec' if vector else 'rmatmat'}_vs_materialize"
+    return bench
+
+
+def _register_normalized_benchmarks() -> None:
+    from vsparse import RECIPES
+
+    for recipe in sorted(RECIPES):
+        for vector in (False, True):
+            fast(_normalized_bench(recipe, vector=vector))
+            fast(_normalized_rbench(recipe, vector=vector))
+
+
+_register_normalized_benchmarks()
+
+
+# -- normalized views vs the *sparse* baseline -------------------------------
+#
+# The cases above compare against `nv.toarray() @ B`, a dense materialization
+# nobody would actually perform -- it makes the view look good for a reason
+# that has nothing to do with the kernels. The honest baseline is the one the
+# prototype benchmarks used: build the sparse delta once, multiply that with
+# scipy, and add the same rank-1 correction the view adds. Both sides then do
+# identical arithmetic and the ratio isolates the kernel.
+#
+# CPU, not just wall, is the point here. Our kernels are `parallel=True` and
+# scipy's are single-threaded, so wall time hides a per-nonzero cost that a
+# shared workstation still pays -- and every recipe but "raw" applies a
+# transcendental to every stored nonzero.
+
+
+def _sparse_delta(nv, mat):
+    """``Delta`` as scipy CSR: what ``to_csr()`` on the view would return.
+
+    ``Delta[i, j] = s[j] * g(x[i, j] / row_scale[i] / gene_scale[j])`` on the
+    stored nonzeros and exactly zero off them, so
+    ``A_norm = Delta + 1 (x) (-c * s)``.
+    """
+    import scipy.sparse as sp
+
+    from vsparse._norm_common import _g_np
+
+    coo = mat.tocoo()
+    gs = nv.gene_scale[coo.col]
+    with np.errstate(divide="ignore", invalid="ignore"):
+        scaled = np.where(gs > 0.0, coo.data / nv.row_scale[coo.row] / gs, 0.0)
+    data = nv.col_post_scale[coo.col] * _g_np(scaled, nv.recipe.g_code)
+    return sp.csr_array((data, (coo.row, coo.col)), shape=mat.shape)
+
+
+def _normalized_vs_sparse(recipe: str, *, vector: bool) -> Callable[[], dict[str, float]]:
+    def bench() -> dict[str, float]:
+        from vsparse import VCSRArray
+
+        mat = integer_counts_csr(20_000, 2_000, density=0.05)
+        v = VCSRArray.from_scipy(mat)
+        nv = v.normalized(recipe)
+        rng = np.random.default_rng(0)
+        B = rng.normal(size=mat.shape[1]) if vector else rng.normal(size=(mat.shape[1], 8))
+
+        delta = _sparse_delta(nv, mat)
+        offset = -(nv.col_mean * nv.col_post_scale)
+
+        def via_view() -> np.ndarray:
+            return nv @ B
+
+        def via_sparse() -> np.ndarray:
+            return delta @ B + (offset @ B)
+
+        np.testing.assert_allclose(via_view(), via_sparse(), rtol=1e-9, atol=1e-9)
+
+        # Context, and core-count dependent: the kernel is `parallel=True`
+        # against a single-threaded scipy. The gated, machine-portable CPU
+        # comparison lives in `normalized_cpu_vs_sparse_1t` below, which is
+        # also the only case here that pays for CPU timing -- doing it in
+        # every case doubled the suite's runtime for a number nothing gates.
+        return {
+            "wall_ratio_view_over_sparse": best_time(via_view) / best_time(via_sparse),
+            "peak_alloc_mb_view": peak_alloc_mb(via_view),
+            "peak_alloc_mb_sparse_delta": peak_alloc_mb(lambda: _sparse_delta(nv, mat)),
+        }
+
+    bench.__name__ = f"normalized_{recipe}_{'matvec' if vector else 'matmat'}_vs_sparse"
+    return bench
+
+
+def _register_sparse_baseline_benchmarks() -> None:
+    from vsparse import RECIPES
+
+    for recipe in sorted(RECIPES):
+        for vector in (False, True):
+            fast(_normalized_vs_sparse(recipe, vector=vector))
+
+
+_register_sparse_baseline_benchmarks()
+
+
+@fast
+def normalized_cpu_vs_sparse_1t() -> dict[str, float]:
+    """Per-nonzero CPU cost of each recipe's kernel, against the same math in scipy.
+
+    Runs with numba pinned to one thread -- `benchmarks.run` gives any case
+    whose name ends in `_1t` a `NUMBA_NUM_THREADS=1` subprocess, which has to
+    happen before numba is imported. Capping the pool from inside the process
+    is not enough: the idle workers still spin, and `process_time()` counts
+    every thread, which made the ratio swing by 2x between runs.
+
+    Pinned, both sides are single-threaded and the ratio isolates what a
+    nonzero costs us over scipy -- almost entirely `g`. It is stable to well
+    under a percent between runs and does not depend on the runner's core
+    count, which is what makes it gateable.
+    """
+    from vsparse import RECIPES, VCSRArray
+
+    # Deliberately larger than the cases above (6M nonzeros, not 2M). Both
+    # sides here land within a small multiple of each other, so the ratio only
+    # settles once each measurement is long enough to swamp scheduling jitter.
+    mat = integer_counts_csr(60_000, 2_000, density=0.05)
+    v = VCSRArray.from_scipy(mat)
+    rng = np.random.default_rng(0)
+    B = rng.normal(size=(mat.shape[1], 8))
+
+    out: dict[str, float] = {}
+    for recipe in sorted(RECIPES):
+        # Each recipe builds its own 6M-nonzero delta (~150 MB). Left to the
+        # allocator, the recipes measured last ran against a progressively more
+        # fragmented heap and their ratios swung by 2x; dropping the previous
+        # one first keeps every recipe on the same footing.
+        gc.collect()
+        nv = v.normalized(recipe)
+        delta = _sparse_delta(nv, mat)
+        offset = -(nv.col_mean * nv.col_post_scale)
+
+        def via_view(nv=nv):
+            return nv @ B
+
+        def via_sparse(delta=delta, offset=offset):
+            return delta @ B + (offset @ B)
+
+        np.testing.assert_allclose(via_view(), via_sparse(), rtol=1e-9, atol=1e-9)
+        # repeat=15: the identity-transform kernels land within ~1.5x of scipy,
+        # so at the default repeat count run-to-run jitter was a large share of
+        # the ratio and the gate was not reproducible across processes.
+        out[f"cpu_ratio_1t_{recipe}"] = best_cpu_time(via_view, repeat=15) / best_cpu_time(
+            via_sparse, repeat=15
+        )
+        del nv, delta, via_view, via_sparse
+
+    return out
+
+
 # -- larger, for the scheduled job -------------------------------------------
 
 

benchmarks/harness.py

@@ -39,6 +39,24 @@ def best_time(fn: Callable[[], Any], repeat: int = 7) -> float:
     return best
 
 
+def best_cpu_time(fn: Callable[[], Any], repeat: int = 7) -> float:
+    """Best *CPU* time over ``repeat`` runs, in seconds. Warms up first.
+
+    Wall time alone flatters every numba kernel here: they are ``parallel=True``
+    while scipy's sparse matmul is single-threaded, so a kernel can be several
+    times faster on the clock while burning an order of magnitude more CPU.
+    On a shared workstation -- the machine this project is aimed at -- that
+    difference is what the user actually pays.
+    """
+    fn()
+    best = float("inf")
+    for _ in range(repeat):
+        start = time.process_time()
+        fn()
+        best = min(best, time.process_time() - start)
+    return best
+
+
 def ratio_vs_scipy(ours: Callable[[], Any], theirs: Callable[[], Any], repeat: int = 7) -> float:
     """``our time / scipy's time`` for the same work."""
     return best_time(ours, repeat) / best_time(theirs, repeat)