diff --git a/CHANGELOG.md b/CHANGELOG.md index 18a39aa..400254b 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/shanuz/module_score.py b/shanuz/module_score.py index 26b4b9f..9cef078 100644 --- a/shanuz/module_score.py +++ b/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 diff --git a/tests/test_module_score_performance.py b/tests/test_module_score_performance.py new file mode 100644 index 0000000..082b57f --- /dev/null +++ b/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")