From 9a7e98600fa43dba23cec4de85e3641fbaedd408 Mon Sep 17 00:00:00 2001 From: Nicholas Ceglia Date: Tue, 2 Jun 2026 00:45:44 -0400 Subject: [PATCH] fix: warn when pyro posterior params fall back to uniform _ensure_pyro_posterior_params silently re-initialized 'q_p_ct_raw' to a uniform 1/P simplex whenever the Pyro param store lacked it. That masked failed/missing param-store loads: downstream posterior metrics then ran on an uninformative prior with no signal to the caller (Notion #24). Emit a RuntimeWarning before the fallback, naming the affected metrics and the likely cause (param store not persisted/restored). Behavior is otherwise unchanged; the early-return path when 'q_p_ct_raw' is present stays silent. Adds tests/test_pyro_params.py: warns on empty store (and still populates the param), silent when already present. Tests snapshot/restore the global Pyro param store to avoid cross-test leakage. Closes Notion #24. Co-Authored-By: Claude Opus 4.8 --- tcri/utils/_utils.py | 11 +++++++++ tests/test_pyro_params.py | 52 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+) create mode 100644 tests/test_pyro_params.py diff --git a/tcri/utils/_utils.py b/tcri/utils/_utils.py index 1c1e8c5..55e3b3e 100644 --- a/tcri/utils/_utils.py +++ b/tcri/utils/_utils.py @@ -120,6 +120,17 @@ def _ensure_pyro_posterior_params(model, adata) -> None: else: raise RuntimeError("Could not infer P from classifier or adata.") + _warnings.warn( + "Pyro param store has no 'q_p_ct_raw'; re-initializing it to a uniform " + "1/P simplex. Downstream posterior metrics (joint_distribution_posterior, " + "phenotypic/clonotypic entropy, mutual information) will run on this " + "uninformative prior instead of the trained posterior. This usually means " + "the Pyro param store failed to load or was never saved; verify the model's " + "Pyro params were persisted and restored (e.g. via save_tcri_session / " + "load_tcri_session).", + RuntimeWarning, + stacklevel=2, + ) init = torch.full((ct_count, P), 1.0 / P, device=device) pyro.param("q_p_ct_raw", init, constraint=constraints.simplex) diff --git a/tests/test_pyro_params.py b/tests/test_pyro_params.py new file mode 100644 index 0000000..19b9384 --- /dev/null +++ b/tests/test_pyro_params.py @@ -0,0 +1,52 @@ +"""Tests for the _ensure_pyro_posterior_params fallback warning (Notion #24). + +When the Pyro param store has no 'q_p_ct_raw', the helper re-initializes it to a +uniform 1/P simplex. That silent fallback used to mask failed param-store loads, so +downstream posterior metrics ran on an uninformative prior with no signal to the +caller. It must now emit a warning. The param store is global state, so each test +snapshots and restores it to avoid leaking into other tests. +""" + +import warnings + +import pyro + +from tcri.utils._utils import _ensure_pyro_posterior_params + + +def test_ensure_pyro_params_warns_on_empty_store(trained_model): + """An empty param store triggers the uniform fallback, which must warn.""" + model, adata = trained_model + store = pyro.get_param_store() + saved = store.get_state() + try: + store.clear() + with warnings.catch_warnings(record=True) as rec: + warnings.simplefilter("always") + _ensure_pyro_posterior_params(model, adata) + + runtime = [w for w in rec if issubclass(w.category, RuntimeWarning)] + assert runtime, "expected a RuntimeWarning on the uniform fallback" + assert "q_p_ct_raw" in str(runtime[0].message) + # fallback still populates the param so callers can proceed + assert "q_p_ct_raw" in store + finally: + store.set_state(saved) + + +def test_ensure_pyro_params_silent_when_present(trained_model): + """When 'q_p_ct_raw' already exists the helper early-returns without warning.""" + model, adata = trained_model + store = pyro.get_param_store() + saved = store.get_state() + try: + _ensure_pyro_posterior_params(model, adata) # guarantee it is present + assert "q_p_ct_raw" in store + + with warnings.catch_warnings(record=True) as rec: + warnings.simplefilter("always") + _ensure_pyro_posterior_params(model, adata) + + assert not any(issubclass(w.category, RuntimeWarning) for w in rec) + finally: + store.set_state(saved)