discrimination_score recomputes an n_pert x n_pert distance matrix one row at a time

[FarzanT]

Summary

discrimination_score (src/cell_eval/metrics/_anndata.py) loops over every
perturbation and calls sklearn.metrics.pairwise_distances once per
perturbation to produce a single row of an n_pert x n_pert distance matrix:

for p_idx, p in enumerate(data.perts):
    ...
    distances = skm.pairwise_distances(
        real_effects[:, include_mask],
        pred_effects[p_idx, include_mask].reshape(1, -1),
        metric=metric,
    ).flatten()

That is n_pert separate sklearn dispatches that together compute the full
matrix. For l2/cosine the per-row work is a BLAS matrix-vector product that
could be a single matrix-matrix product; the loop forfeits that. With three
registered variants (discrimination_score_l1/l2/cosine) and datasets that
have thousands of perturbations, this is a large part of the anndata-metric
runtime.

Why it isn't a one-line pairwise_distances(real, pred)

When exclude_target_gene=True on expression data (the default), each
perturbation excludes a different feature column (the gene named like the
perturbation), so a single unmasked call does not reproduce the per-row masked
distances. The fix has to compute the full matrix once and then remove each
row's target-gene contribution exactly.

Proposal

Compute the full distance matrix once per metric and derive ranks via
argsort. Apply an exact, vectorized rank-1 column correction for the
target-gene-exclusion path (l1/l2/cosine), falling back to exact per-row masked
distances for other metrics. Output stays numerically identical (same ranks).

Local numbers (Apple M2 Pro, Python 3.12, scikit-learn 1.8): bit-identical
ranks across a large synthetic sweep, with speedups that grow with n_pert —
at n_pert=10000, ~8x (l1), ~51x (l2), ~61x (cosine).

I have the implementation, tests, and a benchmark ready and will open a PR.

[FarzanT]

While profiling discrimination_score at scale we hit a second, distinct DE
hotspot that dwarfs it — flagging it here since #238 now carries the fix for
both, and the evidence is useful context.

Evidence (separate metric: de_overlap_metric)

An OmniPert report_only validate-only run over a context with ~18,000
perturbations (16 CPU, 24h wall) ran for 8.8 hours and never reached
results.csv. Faulthandler sampling of the live process (524 periodic stack
dumps) put 513/524 ≈ 98% of all samples inside one metric:

de_overlap_metric
  -> DEComparison.compute_overlap
     -> DEResults.get_top_genes
        hottest leaf: polars frame.py ... in `columns`

The DE compute itself (pdex) was already fast (~135s). This is purely the
overlap metric pass, and it consumed the entire budget before the rest of
the DE metrics and the anndata-pair metrics (discrimination_score included)
even started.

Two compounding inefficiencies in cell_eval/_types/_de.py

1. Redundant rank-matrix pivots. compute_overlap calls
   get_top_genes(sort_by, fdr_threshold) for both sides on every invocation,
   and that builds a polars .pivot() with one column per perturbation
   (~18k-wide). The de/full profile registers 10 variants
   ({overlap, precision} x k in {None, 50, 100, 200, 500}), all with the same
   default sort_by/fdr_threshold — k only truncates the per-pert gene list
   downstream, never the matrix. So the identical pair of 18k-column pivots is
   rebuilt 10 times.
2. O(perts²) column membership. The per-pert loop tested
   pert not in real_sig_rank_matrix.columns or pert not in pred_sig_rank_matrix.columns.
   polars rebuilds a fresh ~18k-element list on every .columns access, twice
   per pert over ~18k perts (~3×10⁸ list rebuilds — the columns leaf the
   sampler kept landing on).

Fix (in #238)

Memoize get_top_genes keyed by (sort_by, fdr_threshold) on the per-side
DEResults (→ one pivot per side instead of 10), and hoist the column names
into sets once before the loop (→ O(perts) membership). Pure performance;
outputs are bit-identical, covered by tests/test_de_overlap_equivalence.py.

Measured old-vs-new with benchmarks/bench_de_overlap.py (one full 10-variant
pass, synthetic real/pred DE, outputs verified identical; n_genes=2000,
n_sig=100):

n_pert	old (s)	new (s)	speedup
1000	2.34	1.05	2.2x
2000	6.79	2.18	3.1x
4000	22.32	4.54	4.9x
8000	80.84	9.78	8.3x

Old grows ~quadratically, new ~linearly, so the speedup climbs with n_pert —
exactly the axis that made the 18k-pert run intractable.

Still unprofiled at this scale

The anndata-pair metrics clustering_agreement and discrimination_score
remain unprofiled at ~18k perts — that run never got past the overlap pass.
#238 already vectorizes discrimination_score; clustering_agreement at this
scale is a separate potential cliff we haven't measured.

[FarzanT]

Same pathology, third instance — now downstream in the DE metric pass

Update on the per-perturbation-scan pattern: the overlap fix (previous comment)
worked — on the same ~18k-perturbation OmniPert report run, the
de_overlap phase that previously stalled ~8.6h now completes. But clearing it
exposed the next bottleneck, the same for pert in iter_perturbations():
full-table-scan shape in two more DE metrics that the old run never reached.

Evidence (faulthandler, ~18k perts, 371k-row DE table) — the runner now sits
continuously in:

DENsigCounts.__call__ -> DEResults.get_significant_genes -> polars filter().collect()

The two cliffs (cell_eval/metrics/_de.py):

1. DENsigCounts.__call__ loops over every pert calling
   get_significant_genes(pert) for the real and pred side — each a full scan
   of the entire DE table (filter((target==pert) & (fdr<thr))), ~18k × 2 ≈ 36k
   scans, where only the significant-gene count is used.
2. compute_pr_auc / compute_roc_auc → compute_generic_auc builds merged
   once (good) but then does merged.filter(target == pert) per pert — ~18k
   more scans, ×2 for the two AUC variants.

Fix (in #238, bit-exact):

• DENsigCounts: one filter_to_significant().group_by(target).len() per side,
  reindexed over the full perturbation universe (0 for perts with no significant
  genes).
• compute_generic_auc: one merged.partition_by(target, maintain_order=True)
  before the loop; absent perts → nan (matching the old empty-slice branch).
  maintain_order keeps each partition in the row order the per-pert filter
  produced, so the arrays handed to average_precision_score / roc_curve are
  bit-identical. The per-pert sklearn calls are unchanged — only the slicing
  goes from O(perts × rows) to O(rows) total.

Covered by tests/test_de_perpert_scan_equivalence.py (verbatim pre-fix
references; exact-equality assertions including the nan and zero-significant
paths).

Audit: the remaining DE metrics — DESpearmanSignificant, DESpearmanLFC,
DEDirectionMatch, DESigGenesRecall — already use a single group_by/join
and don't have the per-pert loop. The one DE metric still unprofiled at this
scale is the anndata-pair clustering_agreement.

So three instances of the same per-pert-scan pattern have now surfaced at high
perturbation count: discrimination_score (this issue), de_overlap, and now
DENsigCounts + the AUC metrics.