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

[shanikawm]

The hot spot

add_module_score pulled each gene's row out on its own and handed the stack to np.mean:

feat_scores = np.mean([_row(mat, feat_idx[g]) for g in used], axis=0)
ctrl_scores = np.mean([_row(mat, feat_idx[g]) for g in ctrl_genes], axis=0)

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. The profile is unambiguous:

ncalls  tottime  percall  filename:lineno(function)
  2268   50.129    0.022  {built-in method scipy.sparse._sparsetools.get_csr_submatrix}

50.1 of 51.1 seconds in one call, invoked once per gene.

The fix

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 this goes from 14× slower than R to an order of magnitude faster.

The exact spelling is load-bearing. Measured on the real matrix, all four formulations of the same mean:

formulation	time	bit-identical to the old loop
per-gene loop (old)	49.28 s	—
ind @ mat (indicator matvec)	0.05 s	no, 7.5e-16
CSC selection, .sum(axis=0)	0.10 s	no, 2.5e-15
CSC selection → .tocsr(), .sum(axis=0)	0.12 s	yes, 0.0

The two faster ones accumulate 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 — so the fast path returns the same bits, not merely the same numbers. sum(axis=0)/k rather than .mean(axis=0) for the same reason: scipy's sparse mean divides on the way through and lands a few ulps away.

A reproducibility defect in the same two lines

The control genes were collected in a set. A mean depends on the order its terms are added, and Python randomises str hashing per process, so iterating that set summed the control expression in a different order every run:

BEFORE, two separate processes, same seed=1 -> identical: False  max|d| = 9.71e-16
AFTER,  two separate processes, same seed=1 -> identical: True   max|d| = 0.0

A dict now, which is also R's semantics — AddModuleScore applies unique() to the sampled names and indexes the matrix with the result, i.e. first-seen order.

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.

Verification

• Suite 876 passed, 25 skipped (was 869 — 7 new tests, no regressions).
• ruff check shanuz 51, ruff check . 201, mypy shanuz 46 — all unchanged.

Mutation testing

mutation	tests failed
per-gene row loop restored	1
control genes back in a set	1
sparse .mean(axis=0) instead of sum/k	1
skip the CSR transpose	3
sort the rows before selecting	1

The sort mutation initially survived, and that was the useful finding: the first fixture was Poisson counts, whose partial sums are all exactly representable, so every summation order gives identical bits and an ordering test written on it passes whatever the code does — the same decorative-test failure mode as the SCT test in #59 and the tie-break test in #61. The fixture is now log-normal, where 110 of 200 cells differ under a reordering, and the test asserts that the fixture can tell the two apart before asserting the property.

One further mutation, mean(axis=0) → sum(axis=0)/n in the dense branch, also survives — verified to be a genuine no-op rather than a gap: numpy's mean is literally sum then divide, bit-identical at every size tested.

🤖 Generated with Claude Code