feat: GPU device seam + synthetic benchmark harness - #45
Merged
Conversation
The `_compute/_xp` device seam has existed since PR5 (resolve_device /
torch_device, CPU + torch-CUDA), and `joint_distribution` accepted `device=` --
but NO metric exposed it, and `joint_draws` hard-coded `device=None` when calling
the engine. So the GPU path was documented and completely unreachable: every
metric ran on CPU no matter what the caller asked for.
Threaded `device=` from the four public metrics down to `_joint_draws`:
tl.mutual_information / clonotypic_entropy / phenotypic_entropy / phenotypic_flux
-> _entropy_metric / _flux_once
-> joint_draws
-> _engine_blocks -> _joint_draws
Verified by spying on `_joint_draws`: all four now deliver the value (flux twice,
once per covariate). Results are unchanged -- device is a placement detail, not a
numerical one -- pinned by test_device_does_not_change_results.
The API contract caught this correctly: adding a public parameter is frozen
interface drift, and test_contract_conformance failed on all four signatures until
`_contract.pyi` was updated deliberately. It also caught a parameter-ORDER
inconsistency -- `device` had landed before `random_state` in flux but after it in
the other three -- so the ordering is now uniform (`device` last everywhere).
This unblocks GPU work: without it there was no code path to profile.
Full suite 166 passed, 3 skipped.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Grid runner reproducing Supplementary Note 1's benchmark design: sweep fuzziness
x N x K-at-inference x seeds, compare TCRi's normalized MI against the closed-form
truth from tcri.datasets.simulate_tcri, plus a KMeans/GMM baseline.
Two traps the harness handles explicitly:
- NORMALIZATION. eq 6 uses the MEAN denominator; tcri defaults to 'min'.
Comparing a min-normalized estimate to a mean-normalized truth silently
inflates it, so --normalize-mode is explicit and defaults to 'average' (eq 6).
- WHICH ORACLE. true_* is the population value, empirical_* is what a perfect
estimator returns on the realized sample; they differ by the plug-in bias at
small N. MAE is reported against both.
Presets: smoke (2 cells) / reduced (24) / full (1500, the note's grid).
--profile attaches the torch profiler to one training run.
Local smoke: TCRi AE 0.034-0.037 vs KMeans 0.116-0.189 -- the baseline-dominance
claim reproduces. HDI covered the realized value 2/2.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
docs/contract/ss-small-team-handbook.md (1998 lines) was an untracked local file that my 'git add -A' picked up into the metrics-contract commit. It is unrelated to the package and was never meant to be committed. Removed from the tree and added to .gitignore so it cannot be swept up again. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
simulate_from_fit_params() reproduces sc_simulator.simulate_dataset from a fitted (pi, omega, gamma_params, V), plus temperature_scale() for the omega^(1/T) axis the published benchmark sweeps. Needed because the synthetic generator CANNOT stand in for the empirical one. The benchmark's true-NMI anchors are 0.520/0.316/0.182 at T=0.1/0.5/1.0; the best symmetric-Dirichlet fit (omega_concentration=0.763) gives 0.633/0.320/0.150, and its sharpening ratio T=0.1->T=1.0 is 4.22x versus the true 2.86x. The mismatch is in the SHAPE of omega's response to temperature, so no reparameterization fixes it. The empirical fit matches all three exactly. Because gamma_params/V are in the pickle, the expression generator is reproducible too, so estimator difficulty matches -- not only the ground truth. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The reference sampler does np.dot(U_i, params['V'].T), i.e. V is (n_genes, n_factors). I had assumed (L, D), which raised a matmul dimension error on the real pickle (984 vs 5). Detect the orientation from L rather than assuming. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Adds --fit-params plus a temperature axis, and the 'published'/'published_quick'
presets matching the figures' axes (K in {8,10,12} at inference, T in {0.1,0.5,1.0},
N in {250..5000}, fuzziness 0.0-0.9, 10 seeds).
Guards the trap: a preset that sweeps temperature REQUIRES --fit-params, because the
synthetic Dirichlet omega cannot reproduce the published true-NMI anchors (its
sharpening response is 4.22x vs the true 2.86x). Verified end to end -- the generator
reproduces 0.520/0.316/0.182 at T=0.1/0.5/1.0 exactly, with 47 clones/984 genes/K=10.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ded by K The presets pin k_infer=10. Sweeping different fitted omegas (K=8/10/12) without this would conflate 'different fixture' with 'wrong K at inference'. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
…ll overrides The metric reports E_s[NMI(J_s)]. NMI is nonlinear in the joint, so that is not the NMI of the posterior. Read the same draws both ways and record the gap; it needs no ground truth, so it also works on real data. --temperature/--n-cells let a single grid cell be driven to convergence without paying for the whole sweep. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Synthetic-only finding, mechanism identified and magnitude bounded, fix undecided. Written up so the open questions (posterior summary convention, eq-2 concentration) are on record before anyone acts on the benchmark numbers. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Bugbot is not enabled for your account, so this pull request was not reviewed. Enable Bugbot in the Cursor dashboard to get automatic reviews on future PRs. |
Per-entry Dirichlet concentration is local_scale/P. Below 1 the posterior draws are corner-seeking, which is the proposed source of the upward NMI bias. This exposes the knob so the prediction can be tested: raising it should collapse the Jensen gap and the error floor together. Moves both the guide's posterior and (via uns) the metric's draw, so a result here is not attributable to one side without a follow-up that pins the metric separately. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Threads
device=through to the metrics engine (CU-01) and adds the benchmarkharness that has been used to calibrate the NMI estimator against known truth.
Contents
GPU seam.
device=was in the frozen_contract.pyibutjoint_drawshard-coded
device=None, so no public metric could reach it. Now wired through.Measured: metrics ~6× faster on an A100.
benchmarks/run_grid.py. Sweeps coupling strength, sample size, K and seedagainst a known population truth. Each cell carries three reference points —
population truth, a label oracle (plug-in NMI over the realized true labels), and
a clustering baseline. Presets:
smoke,reduced,full,published,published_quick. Overrides for--k-infer,--temperature,--n-cellsso asingle cell can be driven to convergence without paying for the whole sweep.
simulate_from_fit_params()intcri/datasets/_simulate.py. Generates from afitted parameter set rather than a synthetic Dirichlet ω, so benchmark numbers are
comparable to published ones instead of to a toy problem. Reproduces the published
ground-truth anchors exactly (0.520 / 0.316 / 0.182).
Both posterior summaries. Every row now carries
E[NMI(J)](what the packagereports) alongside
NMI(E[J])computed from the same draws, plus their difference.The gap needs no ground truth, so the same diagnostic ports to real data.
Why the harness matters
It found a coupling-dependent upward bias in the NMI estimate that survives to
convergence and reproduces on three independently fitted parameter sets. Written up
in
docs/nmi_temperature_bias.md. The finding is synthetic-only and the fix isundecided — nothing in this PR changes estimator behavior.
Tests
166 passed, 3 skippedon the pinned venv.🤖 Generated with Claude Code