perf(module-score): read the matrix once per gene set, not once per gene

CHANGELOG.md

@@ -171,6 +171,38 @@ on `main`; none of it is on PyPI.
 
 ### Fixed
 
+- **`add_module_score` read the expression matrix once per gene, and was 100×
+  slower than it needed to be.** Each gene's row was pulled out on its own
+  (`mat[i, :]` inside a list comprehension) and the stack handed to `np.mean`.
+  Every assay layer here is **CSC**, so slicing a single row walks the whole
+  column-major matrix: on the THP-1 ECCITE data (18,381 × 20,729, 69.5M
+  nonzeros) that is ~22 ms per gene, and the default `ctrl=100` draws a couple
+  of thousand control genes. **50 of 51 profiled seconds** were inside one scipy
+  call, `get_csr_submatrix`, invoked once per gene. Now one row selection per
+  gene set, transposed to CSR so the rows are summed in the order they were
+  asked for:
+
+  | | before | after |
+  |---|---|---|
+  | `cell_cycle_scoring` (THP-1, 2 programs) | 168.4 s | **0.90 s** |
+  | `add_module_score` (30-gene program) | 51.3 s | **0.51 s** |
+
+  R's `CellCycleScoring` takes ~12 s on the same data, so shanuz now runs it an
+  order of magnitude faster rather than 14× slower. The spelling is load-bearing
+  and **bit-identical** to the old arithmetic: an indicator-vector matvec
+  (`ind @ mat`) is faster still but lands 7.5e-16 away, and summing the CSC
+  selection directly 2.5e-15 away, because each accumulates the columns in a
+  different order.
+- **The same seed gave a different module score in a different process.** The
+  control genes were collected in a `set`, and a mean depends on the order its
+  terms are added. Python randomises `str` hashing per process, so iterating
+  that set summed the control expression in a different order every run — the
+  same object at `seed=1` scored 9.7e-16 apart in two processes. A `dict` now,
+  which is also R's semantics (`AddModuleScore` applies `unique()` to the
+  sampled names and indexes the matrix with the result). THP-1 `S.Score`,
+  `G2M.Score` and the interferon program move by at most 2.1e-15 against the old
+  code — which had no fixed value to move from — and all 20,729 `Phase` calls
+  are unchanged.
 - **A reduction that could not use every feature you asked for labelled its
   loadings with the features you asked for anyway.** `_get_scaled_data` filtered
   the request down to what the layer carried and returned only the matrix, so

shanuz/module_score.py

@@ -41,11 +41,32 @@ def _assay_data(seurat, assay: Optional[str], layer: str = "data"):
         return (data if not is_matrix_empty(data) else assay_obj.counts), feats
 
 
-def _row(mat, idx) -> np.ndarray:
-    row = mat[idx, :]
-    if sp.issparse(row):
-        return np.asarray(row.todense()).flatten()
-    return np.asarray(row).flatten()
+def _mean_over_rows(mat, rows: list[int]) -> np.ndarray:
+    """Per-cell mean over ``rows`` — one row selection, not one per row.
+
+    This used to pull each gene out on its own (``mat[i, :]`` in a list
+    comprehension) and hand the stack to ``np.mean``. Every assay layer here is
+    **CSC**, and slicing a single row out of a column-major matrix walks all of
+    it: on the THP-1 ECCITE data (18,381 × 20,729, 69.5M nonzeros) that was
+    ~22 ms per gene, and a default ``ctrl=100`` draws a couple of thousand
+    control genes. 50 of 51 profiled seconds were inside
+    ``scipy.sparse._sparsetools.get_csr_submatrix``, called once per gene.
+
+    Selecting the rows in one go and transposing to CSR first costs 0.12 s for
+    the same 2,268 genes — **410× faster** — and is **bit-identical**, which is
+    the reason for this exact spelling. Two other formulations are faster still
+    and are not: an indicator-vector matvec (``ind @ mat``) differs by 7.5e-16
+    and summing the CSC slice directly by 2.5e-15, because each accumulates the
+    columns in a different order. Summing a *CSR* selection walks the rows in
+    the order they were asked for, which is what the old loop did.
+    """
+    sub = mat[rows, :]
+    if sp.issparse(sub):
+        # sum(axis=0)/k rather than mean(axis=0): scipy's mean divides
+        # elementwise on the way through and lands a few ulps away.
+        return np.asarray(sub.tocsr().sum(axis=0)).ravel() / len(rows)
+    # np.mean over a list of rows *is* this, once the list is stacked.
+    return np.asarray(sub).mean(axis=0)
 
 
 def _alnum(s: str) -> str:
@@ -154,7 +175,13 @@ def add_module_score(
             continue
 
         # Control gene set: per program-gene, sample `ctrl` from its bin.
-        ctrl_genes: set[str] = set()
+        # A dict, not a set: a mean depends on the order its terms are added,
+        # and Python randomises str hashing per process, so iterating a set of
+        # gene names gave a control score that differed in its last bits from
+        # one run to the next. First-seen order is also R's — `AddModuleScore`
+        # applies `unique()` to the sampled names and indexes the matrix with
+        # the result.
+        ctrl_genes: dict[str, None] = {}
         for g in used:
             b = gene_to_bin.get(g)
             if b is None:
@@ -164,11 +191,11 @@ def add_module_score(
                 continue
             size = min(ctrl, len(candidates))
             picked = rng.choice(candidates, size=size, replace=False)
-            ctrl_genes.update(picked.tolist())
+            ctrl_genes.update(dict.fromkeys(picked.tolist()))
 
-        feat_scores = np.mean([_row(mat, feat_idx[g]) for g in used], axis=0)
+        feat_scores = _mean_over_rows(mat, [feat_idx[g] for g in used])
         if ctrl_genes:
-            ctrl_scores = np.mean([_row(mat, feat_idx[g]) for g in ctrl_genes], axis=0)
+            ctrl_scores = _mean_over_rows(mat, [feat_idx[g] for g in ctrl_genes])
         else:
             ctrl_scores = np.zeros(n_cells)
         seurat.meta_data[label] = feat_scores - ctrl_scores

tests/test_module_score_performance.py

@@ -0,0 +1,253 @@
+"""How `add_module_score` reads the matrix, and in what order it adds it up.
+
+Two defects sat in the same two lines.
+
+**Speed.** Each gene's row was pulled out on its own — `mat[i, :]` inside a list
+comprehension — and every assay layer here is CSC, so slicing one row walks the
+whole column-major matrix. On the THP-1 ECCITE data (18,381 × 20,729, 69.5M
+nonzeros) that was ~22 ms per gene, and `ctrl=100` draws a couple of thousand
+control genes: 50 of 51 profiled seconds were inside one scipy call, invoked once
+per gene. `cell_cycle_scoring` took 168 s against R's 12.
+
+**Reproducibility.** The control genes were collected in a `set`, and a mean
+depends on the order its terms are added. Python randomises `str` hashing per
+process, so the same object with the same seed produced a different score in a
+different process. Not a large difference — 9.7e-16 — but a value that moves when
+nothing moved is the kind of thing that gets chased for an afternoon.
+
+The fix is one row selection per gene set, transposed to CSR so the rows are
+walked in the order they were asked for, and a dict in place of the set. Both are
+pinned here: the first as an exact-equality property, the second by running the
+same computation under two hash seeds.
+"""
+import os
+import subprocess
+import sys
+import textwrap
+
+import numpy as np
+import pytest
+import scipy.sparse as sp
+
+from shanuz import create_shanuz_object, normalize_data
+from shanuz.module_score import _mean_over_rows, add_module_score
+
+N_GENES, N_CELLS = 300, 200
+
+
+def _matrix(seed=0, density=0.3):
+    rng = np.random.default_rng(seed)
+    dense = rng.poisson(2.0, size=(N_GENES, N_CELLS)).astype(float)
+    dense[rng.random(dense.shape) > density] = 0.0
+    return dense
+
+
+def _wide_matrix(seed=0, density=0.3):
+    """Values spanning many orders of magnitude, so summation order is visible.
+
+    Small Poisson counts are all exactly representable and their partial sums
+    stay exact, so *every* order gives identical bits — a fixture built from
+    them cannot express an ordering property, and a test written on one passes
+    whatever the code does. Log-normalized expression has the spread this needs;
+    counts do not.
+    """
+    rng = np.random.default_rng(seed)
+    dense = rng.lognormal(0.0, 6.0, size=(N_GENES, N_CELLS))
+    dense[rng.random(dense.shape) > density] = 0.0
+    return dense
+
+
+def _obj(seed=0):
+    obj = create_shanuz_object(
+        counts=sp.csc_matrix(_matrix(seed)),
+        assay="RNA",
+        feature_names=[f"g{i}" for i in range(N_GENES)],
+        cell_names=[f"c{j}" for j in range(N_CELLS)],
+    )
+    normalize_data(obj)
+    return obj
+
+
+def _stack_mean(mat, rows):
+    """The formulation that was replaced, kept as the reference."""
+    out = []
+    for i in rows:
+        r = mat[i, :]
+        out.append(np.asarray(r.todense()).flatten() if sp.issparse(r)
+                   else np.asarray(r).flatten())
+    return np.mean(out, axis=0)
+
+
+# ---------------------------------------------------------------------------
+# 1. The faster path returns the same bits, not merely the same numbers
+# ---------------------------------------------------------------------------
+
+@pytest.mark.parametrize("fmt", ["csc", "csr", "dense"])
+def test_mean_over_rows_is_bit_identical_to_the_per_row_stack(fmt):
+    """Exact equality, deliberately — `assert_allclose` would hide the point.
+
+    Two formulations that are *faster still* are not bit-identical, and the
+    difference is only in the last ulps: an indicator-vector matvec
+    (``ind @ mat``) lands 7.5e-16 away and summing the CSC selection directly
+    2.5e-15 away, because each accumulates the columns in a different order.
+    Nothing downstream would notice, and that is exactly why a tolerance here
+    would let a silent change of estimator through.
+    """
+    dense = _wide_matrix()
+    mat = {"csc": sp.csc_matrix(dense), "csr": sp.csr_matrix(dense),
+           "dense": dense}[fmt]
+    rng = np.random.default_rng(3)
+    rows = sorted(rng.choice(N_GENES, size=64, replace=False).tolist())
+
+    got = _mean_over_rows(mat, rows)
+    want = _stack_mean(mat, rows)
+    assert np.array_equal(got, want), (
+        f"max|diff| = {np.abs(got - want).max():.3e} — same numbers, "
+        f"different arithmetic"
+    )
+
+
+def test_mean_over_rows_respects_the_order_it_was_given():
+    """Row order is not cosmetic: it is the order the terms are summed in.
+
+    Sorting the indices before selecting is the obvious optimization here —
+    scipy indexes sorted rows faster — and it silently returns different
+    numbers. That is also what makes the control-gene ordering below observable
+    at all, so if this property does not hold there is nothing to reproduce.
+    """
+    mat = sp.csc_matrix(_wide_matrix())
+    rng = np.random.default_rng(11)
+    rows = rng.choice(N_GENES, size=64, replace=False).tolist()
+    assert rows != sorted(rows), "fixture must present the rows out of order"
+
+    assert np.array_equal(_mean_over_rows(mat, rows), _stack_mean(mat, rows))
+    # The fixture has to be able to tell the two apart, or the assertion above
+    # holds for any implementation and guards nothing.
+    assert not np.array_equal(
+        _mean_over_rows(mat, rows), _mean_over_rows(mat, sorted(rows))
+    ), "fixture is too well-conditioned to see a reordering"
+
+
+# ---------------------------------------------------------------------------
+# 2. One selection per gene set, not one per gene
+# ---------------------------------------------------------------------------
+
+class _CountingMatrix(sp.csc_matrix):
+    """A CSC matrix that records how many times it is sliced.
+
+    A *subclass*, not a wrapper: `add_module_score` branches on
+    `sp.issparse(mat)`, and a proxy that merely forwards attributes takes the
+    dense branch, where `np.asarray` turns it into a 0-d object array. Same trap
+    that made `fetch_data` return a column of matrices.
+    """
+
+    calls = 0
+
+    def __getitem__(self, key):
+        _CountingMatrix.calls += 1
+        return super().__getitem__(key)
+
+
+def test_add_module_score_slices_per_gene_set_not_per_gene():
+    """The structural property behind the 100× — asserted, not timed.
+
+    A wall-clock threshold would be flaky on a loaded machine and would not say
+    *why* it was slow. The slice count does: one for the expression-bin pool,
+    then one for the program's genes and one for its controls. The old code made
+    one call per gene, which on a default `ctrl=100` is thousands.
+
+    Counted through `add_module_score` rather than by calling `_mean_over_rows`
+    directly, because the defect was the *call site* — a helper that reads the
+    whole selection at once is no use to a loop that hands it one row at a time.
+    """
+    obj = _obj()
+    assay = obj.assays["RNA"]
+    _CountingMatrix.calls = 0
+    assay.layers["data"] = _CountingMatrix(assay.layers["data"])
+
+    genes = [f"g{i}" for i in range(20)]
+    add_module_score(obj, features={"prog": genes}, ctrl=50, seed=1)
+
+    assert _CountingMatrix.calls <= 4, (
+        f"{_CountingMatrix.calls} slices for a 20-gene program — the matrix "
+        f"should be read once per gene set, not once per gene"
+    )
+    # And it did the work: a score per cell, not a column of zeros.
+    scores = obj.meta_data["prog"].to_numpy()
+    assert scores.shape == (N_CELLS,)
+    assert np.abs(scores).sum() > 0
+
+
+def test_scores_still_match_the_per_row_stack_end_to_end():
+    """The whole function, against the arithmetic it used to do.
+
+    `_mean_over_rows` being right in isolation does not make `add_module_score`
+    right: the row indices it is handed, and their order, are chosen at the call
+    site. Recomputed here from the same control genes the function drew.
+    """
+    obj = _obj()
+    mat = obj.assays["RNA"].layers["data"]
+    feats = obj.assays["RNA"].features()
+    idx = {f: i for i, f in enumerate(feats)}
+    genes = [f"g{i}" for i in range(20)]
+
+    add_module_score(obj, features={"prog": genes}, ctrl=50, seed=1)
+    got = obj.meta_data["prog"].to_numpy()
+
+    # Re-derive the program half exactly; the control half is checked by the
+    # reproducibility test, which does not need to know which genes were drawn.
+    prog_mean = _stack_mean(mat, [idx[g] for g in genes])
+    assert np.isfinite(got).all()
+    assert not np.array_equal(got, prog_mean), "controls were not subtracted"
+    assert np.array_equal(
+        _mean_over_rows(mat, [idx[g] for g in genes]), prog_mean
+    )
+
+
+# ---------------------------------------------------------------------------
+# 3. The same seed gives the same score in a different process
+# ---------------------------------------------------------------------------
+
+_CHILD = textwrap.dedent(
+    """
+    import numpy as np, scipy.sparse as sp
+    from shanuz import create_shanuz_object, normalize_data
+    from shanuz.module_score import add_module_score
+
+    rng = np.random.default_rng(0)
+    dense = rng.poisson(2.0, size=(300, 200)).astype(float)
+    dense[rng.random(dense.shape) > 0.3] = 0.0
+    obj = create_shanuz_object(
+        counts=sp.csc_matrix(dense), assay="RNA",
+        feature_names=[f"g{i}" for i in range(300)],
+        cell_names=[f"c{j}" for j in range(200)],
+    )
+    normalize_data(obj)
+    add_module_score(obj, features={"prog": [f"g{i}" for i in range(20)]},
+                     ctrl=50, seed=1)
+    print(obj.meta_data["prog"].to_numpy().tobytes().hex())
+    """
+)
+
+
+def _score_under_hash_seed(seed: str) -> str:
+    env = dict(os.environ, PYTHONHASHSEED=seed)
+    out = subprocess.run(
+        [sys.executable, "-c", _CHILD], capture_output=True, text=True, env=env,
+        cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+    )
+    assert out.returncode == 0, out.stderr
+    return out.stdout.strip()
+
+
+def test_control_scores_do_not_depend_on_pythons_hash_seed():
+    """A `set` of gene names iterates in an order that changes per process.
+
+    The control score is a mean over those genes, and floating-point addition is
+    not associative, so the same object at the same seed scored differently in a
+    different process — 9.7e-16 on THP-1, enough to move a value nothing else
+    moved. Two child processes, two hash seeds, byte-for-byte comparison: the
+    `set` version fails this and no other test in the suite does, because
+    everything else runs in one process where the order is at least stable.
+    """
+    assert _score_under_hash_seed("0") == _score_under_hash_seed("12345")