From a2586a69851d58ecc4c325c5f71fb19f334bf75e Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Wed, 5 Aug 2026 16:13:12 -0400 Subject: [PATCH] fix: flux defaults to KL per eq 7; unify four drifted defaults flux: distance_metric now defaults to 'kl'. METRICS eq 7 defines phenotypic flux AS the KL divergence; l1/jsd stay available as extensions. This CHANGES reported flux numbers for callers relying on the default. Retires OPEN_QUESTIONS['flux_distance_default']. defaults: four knobs were declared in two layers with different values. The outer declaration silently wins, so the inner one is dead code that only bites someone constructing TCRIModule or UnifiedTrainingPlan directly. Unified each to the value that was already effective, so no fitted result changes: global_scale TCRIModel 5.0 vs TCRIModule 10.0 -> 5.0 local_scale TCRIModel 3.0 vs TCRIModule 5.0 -> 3.0 n_steps_kl_warmup train() 2000 vs UnifiedTrainingPlan 1000 -> 2000 (reconstruction_loss_scale was the same defect, unified earlier in 19db68e) global_scale was found by the new test on its first run, having been missed by every manual pass including today's. tests/test_shared_defaults.py guards the class rather than the instances: it compares every shared default across the two construction seams. The knob test could not catch this -- it verifies a value ARRIVES at its target, which it does; the defect is that the two declared defaults differ. 177 passed, 3 skipped. Co-Authored-By: Claude Opus 5 --- tcri/model/_module.py | 4 +- tcri/model/_training.py | 2 +- tcri/tools/_flux.py | 12 ++-- tcri/tools/_metrics_contract.py | 15 ++-- tests/test_metrics_contract_conformance.py | 2 +- tests/test_shared_defaults.py | 81 ++++++++++++++++++++++ 6 files changed, 99 insertions(+), 17 deletions(-) create mode 100644 tests/test_shared_defaults.py diff --git a/tcri/model/_module.py b/tcri/model/_module.py index 78024e0..186cf73 100644 --- a/tcri/model/_module.py +++ b/tcri/model/_module.py @@ -31,8 +31,8 @@ def __init__( n_latent: int, P: int, n_batch: int, - global_scale: float = 10.0, - local_scale: float = 5.0, + global_scale: float = 5.0, # must match TCRIModel.__init__, which overrides it + local_scale: float = 3.0, # must match TCRIModel.__init__, which overrides it prior_temperature: float = 1.0, guide_temperature: float = 1.0, gate_prob: Optional[float] = 0.5, # None = additive (no gating) diff --git a/tcri/model/_training.py b/tcri/model/_training.py index 85b0746..2beff3f 100644 --- a/tcri/model/_training.py +++ b/tcri/model/_training.py @@ -37,7 +37,7 @@ class UnifiedTrainingPlan(PyroTrainingPlan): def __init__( self, module: TCRIModule, - n_steps_kl_warmup: int = 1000, + n_steps_kl_warmup: int = 2000, # must match TCRIModel.train, which overrides it reconstruction_loss_scale: float = 1e-2, num_particles: int = 5, optimizer_config: dict = None, diff --git a/tcri/tools/_flux.py b/tcri/tools/_flux.py index 7f2b174..4d8f65b 100644 --- a/tcri/tools/_flux.py +++ b/tcri/tools/_flux.py @@ -2,8 +2,10 @@ shift between two covariate values (§7.5), engine-backed. For each clone in the ``cov_from`` ∩ ``cov_to`` intersection, the distance between its -phenotype distribution at ``cov_from`` and ``cov_to``, via ``_distance`` (l1 / kl / jsd, -KL & JSD in **bits**). ``n_samples=0`` is the plug-in (a clone with no real shift reads +phenotype distribution at ``cov_from`` and ``cov_to``, via ``_distance`` (kl / l1 / jsd, +KL & JSD in **bits**). **KL is the default**: the metrics document's eq 7 defines +phenotypic flux AS the KL divergence, and l1/jsd are tcri extensions for when a bounded or +symmetric measure is wanted. ``n_samples=0`` is the plug-in (a clone with no real shift reads exactly 0); ``n_samples>0`` redraws both sides coherently (same seed) and summarizes. """ from __future__ import annotations @@ -46,10 +48,12 @@ def _flux_once(adata, *, cov_from, cov_to, n_samples, weighted, temperature, clo def phenotypic_flux(adata, *, cov_from, cov_to, groupby=None, splitby=None, n_samples=0, - temperature=1.0, clones=None, weighted=False, distance_metric="l1", + temperature=1.0, clones=None, weighted=False, distance_metric="kl", random_state=None, device=None): """Per-clone phenotype-distribution distance from ``cov_from`` to ``cov_to`` (bits for - kl/jsd). ``groupby`` → tidy DataFrame (one row per group×clone).""" + kl/jsd). ``distance_metric`` defaults to ``"kl"`` — METRICS eq 7 defines flux as the KL + divergence; ``"l1"`` and ``"jsd"`` remain available. ``groupby`` → tidy DataFrame (one + row per group×clone).""" if groupby is not None: def _compute(cl): return _flux_once(adata, cov_from=cov_from, cov_to=cov_to, n_samples=n_samples, diff --git a/tcri/tools/_metrics_contract.py b/tcri/tools/_metrics_contract.py index 2daed7d..b4780e4 100644 --- a/tcri/tools/_metrics_contract.py +++ b/tcri/tools/_metrics_contract.py @@ -140,7 +140,10 @@ def __repr__(self): # pragma: no cover - debug aid "kernel: KL and JS in bits (log2), L1 in [0, 2]." ), empty="NaN when the clone has no mass in either condition", - note_eq="METRICS eq 7 (D_KL); the prose defines phenotypic flux AS that divergence", + note_eq=( + "METRICS eq 7 -- D_KL(P||Q) = sum_x P(x) log(P(x)/Q(x)); the prose defines " + "phenotypic flux AS that divergence, and 'kl' is the DEFAULT accordingly." + ), ), } @@ -161,12 +164,6 @@ def __repr__(self): # pragma: no cover - debug aid #: decision someone has to make. A test asserts no key appears in both dicts, so a pending #: decision cannot be quietly refiled as a feature. OPEN_QUESTIONS = { - "flux_distance_default": ( - "METRICS eq 7 defines phenotypic flux AS the KL divergence, but phenotypic_flux " - "defaults to distance_metric='l1'. The requested behaviour is KL by default with " - "l1/js still available; the kernels already exist in tcri/_distance.py " - "('kl'/'dkl', 'l1', 'js'/'jsd') and only the default differs." - ), "posterior_summary_of_a_nonlinear_metric": ( "At n_samples>0 tcri reports E_s[NMI(J_s)] -- the mean of the per-draw NMI. NMI is " "nonlinear in the joint, so this is not the NMI of the posterior, and the two " @@ -241,8 +238,8 @@ def __repr__(self): # pragma: no cover - debug aid "flux_distance_choices": ( "METRICS eq 7 defines flux as the KL divergence; tcri additionally offers L1 and " "Jensen-Shannon via `distance_metric`, since a symmetric or bounded measure is " - "sometimes wanted. Offering them is the extension; which one is the DEFAULT is " - "not -- see OPEN_QUESTIONS['flux_distance_default']." + "sometimes wanted. Offering them is the extension; the DEFAULT is 'kl', matching " + "eq 7. Changing the default changes reported flux numbers." ), "n_clones_ref": ( "clonotypic_entropy accepts `n_clones_ref` to FIX the normalizer across groups; " diff --git a/tests/test_metrics_contract_conformance.py b/tests/test_metrics_contract_conformance.py index c75e018..0a383b3 100644 --- a/tests/test_metrics_contract_conformance.py +++ b/tests/test_metrics_contract_conformance.py @@ -67,7 +67,7 @@ def test_sources_are_archived_with_a_hash(): def test_open_questions_are_not_quietly_sanctioned(): """A live disagreement with the source document must not be filed as an 'extension'. Extensions are things the document does not specify; these are things it does.""" - for key in ("flux_distance_default", "posterior_summary_of_a_nonlinear_metric"): + for key in ("posterior_summary_of_a_nonlinear_metric",): assert key in MC.OPEN_QUESTIONS and len(MC.OPEN_QUESTIONS[key]) > 40 assert key not in MC.SANCTIONED_EXTENSIONS diff --git a/tests/test_shared_defaults.py b/tests/test_shared_defaults.py new file mode 100644 index 0000000..d1970ae --- /dev/null +++ b/tests/test_shared_defaults.py @@ -0,0 +1,81 @@ +"""Guard against defaults drifting apart across layers. + +A knob is declared in more than one place: ``TCRIModel.__init__`` constructs +``TCRIModule``, and ``TCRIModel.train`` constructs ``UnifiedTrainingPlan``, each passing +its own value down. When the two declarations disagree, the outer one silently wins and +the inner one becomes **dead** — reachable only by constructing the inner object directly, +which is exactly what a test fixture or a downstream user does. + +This has happened four times: + +* ``reconstruction_loss_scale`` — ``train()`` 1e-3, ``_module`` 1e-3, ``_training`` 1e-2. + The effective value was 1e-3, so ``_training``'s was already dead. Found only while + re-measuring deviation [E]; unified in 19db68e. +* ``local_scale`` — ``TCRIModel`` 3.0, ``TCRIModule`` 5.0. Effective 3.0. +* ``n_steps_kl_warmup`` — ``train()`` 2000, ``UnifiedTrainingPlan`` 1000. Effective 2000. +* ``global_scale`` (α) — ``TCRIModel`` 5.0, ``TCRIModule`` 10.0. Effective 5.0. Found by + this test on its first run, having been missed by every manual pass. + +None was caught by the knob test, because that verifies a value *arrives* at its target — +which it does. The defect is that the two declared defaults differ, so a caller reading one +signature is misled about what the package does. + +Cheap and signature-only: no model is constructed. +""" +from __future__ import annotations + +import inspect + +import pytest + +from tcri.model._model import TCRIModel +from tcri.model._module import TCRIModule +from tcri.model._training import UnifiedTrainingPlan + +#: (outer callable, inner callable, human label). The outer constructs the inner and +#: forwards these arguments, so a disagreement means the inner default is unreachable. +PAIRS = [ + (TCRIModel.__init__, TCRIModule.__init__, "TCRIModel.__init__ -> TCRIModule.__init__"), + (TCRIModel.train, UnifiedTrainingPlan.__init__, "TCRIModel.train -> UnifiedTrainingPlan"), +] + +#: Names that are deliberately allowed to differ, with the reason. Keep this empty unless +#: there is a real argument for a divergent default — "the inner one is never used" is not +#: one, since that is precisely the trap. +ALLOWED_DIVERGENCE: dict[str, str] = {} + + +def _defaults(fn): + return { + name: p.default + for name, p in inspect.signature(fn).parameters.items() + if p.default is not inspect.Parameter.empty + } + + +@pytest.mark.parametrize("outer,inner,label", PAIRS, ids=[p[2] for p in PAIRS]) +def test_shared_defaults_agree(outer, inner, label): + """A knob declared in both layers must declare the SAME default in both.""" + o, i = _defaults(outer), _defaults(inner) + mismatched = { + k: (o[k], i[k]) + for k in set(o) & set(i) + if o[k] != i[k] and k not in ALLOWED_DIVERGENCE + } + assert not mismatched, ( + f"{label}: defaults disagree, so the inner value is dead code that only bites " + f"someone constructing the inner object directly — " + + "; ".join(f"{k}: outer={ov!r} inner={iv!r}" for k, (ov, iv) in sorted(mismatched.items())) + ) + + +def test_the_known_drifted_knobs_are_pinned(): + """Regression lock on the four that actually drifted, so a future edit to one layer + cannot silently reintroduce the split.""" + model, module = _defaults(TCRIModel.__init__), _defaults(TCRIModule.__init__) + train, plan = _defaults(TCRIModel.train), _defaults(UnifiedTrainingPlan.__init__) + + assert model["local_scale"] == module["local_scale"] == 3.0 + assert model["global_scale"] == module["global_scale"] == 5.0 + assert train["n_steps_kl_warmup"] == plan["n_steps_kl_warmup"] == 2000 + assert train["reconstruction_loss_scale"] == 1e-2