Add LOO integration tests for MCMC results

Add several integration tests covering leave-one-out (LOO) functionality for MCMC results: mechanistic point/curve LOO, ArviZ consistency check, surrogate-mode LOO (including one-step Euler and reuse-of-noise variants), compare() between fits, and LOO plotting/pointwise outputs. Also import numpy for array checks and set matplotlib to Agg in plot tests. Tests use small MCMC configs and a trained NeuralODE surrogate where applicable to validate elpd, pareto_k shapes, and n_data_points behavior.

Author: JR-1991

tests/integration/test_mcmc.py

@@ -1,6 +1,7 @@
 import random
 
 import jax.numpy as jnp
+import numpy as np
 import optax
 
 import catalax.mcmc as cmc
@@ -77,6 +78,116 @@ def test_surrogate_model(self, generate_data):
             yerrs=1e-5,
         )
 
+    def test_loo_mechanistic(self, generate_data):
+        """Mechanistic LOO returns a valid ELPDData over concentration points."""
+        model, dataset = generate_data
+        config = cmm.MCMCConfig(num_warmup=50, num_samples=100, verbose=0)
+        results = cmm.run_mcmc(model=model, dataset=dataset, config=config, yerrs=1e-2)
+
+        # Reusing the inferred noise keeps every (measurement, time, obs) point.
+        loo_point = results.loo(dataset, leave_out="point")
+        n_obs = int(loo_point.n_data_points)
+        assert (
+            n_obs == dataset.to_jax_arrays(model.get_observable_state_order())[0].size
+        )
+
+        # Leave-one-curve-out collapses each measurement series to one unit.
+        loo_curve = results.loo(dataset, leave_out="curve")
+        assert int(loo_curve.n_data_points) == len(dataset.measurements)
+
+    def test_loo_consistency_check(self, generate_data):
+        """Eval-model reconstruction must match ArviZ native LOO (mechanistic)."""
+        model, dataset = generate_data
+        config = cmm.MCMCConfig(num_warmup=50, num_samples=100, num_chains=2, verbose=0)
+        results = cmm.run_mcmc(model=model, dataset=dataset, config=config, yerrs=1e-2)
+
+        check = results.loo_consistency_check(dataset, yerrs=1e-2)
+        assert check["agree"], check
+
+    def test_loo_surrogate(self, generate_data):
+        """Surrogate-mode posterior still yields concentration-space LOO."""
+        model, dataset = generate_data
+        aug = dataset.augment(n_augmentations=10)
+
+        rbf = ctn.RBFLayer(0.2)
+        neural_ode = ctn.NeuralODE.from_model(
+            model,
+            width_size=8,
+            depth=2,
+            activation=rbf,  # type: ignore
+        )
+        strategy = ctn.Strategy()
+        strategy.add_step(
+            lr=1e-2, length=0.1, steps=100, batch_size=15, loss=optax.log_cosh
+        )
+        neural_ode = ctn.train_neural_ode(
+            model=neural_ode,
+            dataset=aug,
+            strategy=strategy,
+            print_every=1000,
+            weight_scale=1e-7,
+        )
+
+        config = cmm.MCMCConfig(num_warmup=50, num_samples=100, verbose=0)
+        results = cmm.run_mcmc(
+            model=model,
+            dataset=aug,
+            config=config,
+            surrogate=neural_ode,
+            yerrs=1e-2,
+        )
+
+        # Reuse the sampled rates, Euler-integrate, and score against the
+        # *measured* concentrations -- not the surrogate rates. The stored yerrs
+        # is rate-space for a surrogate fit, so pass a concentration-space one.
+        loo_res = results.loo(dataset, yerrs=0.5)
+        assert int(loo_res.n_data_points) > 0
+        # One Pareto-k per held-out data point (the headline diagnostic).
+        assert np.asarray(loo_res.pareto_k).shape[0] == int(loo_res.n_data_points)
+
+        # One-step-ahead integration is also available.
+        loo_onestep = results.loo(dataset, yerrs=0.5, integration="euler_onestep")
+        assert int(loo_onestep.n_data_points) > 0
+
+        # The reuse-the-inferred-noise variant is also still available.
+        loo_reuse = results.loo(dataset, sigma_source="reuse")
+        assert int(loo_reuse.n_data_points) > 0
+
+    def test_loo_compare(self, generate_data):
+        """compare() ranks two fits on the same concentration-space footing."""
+        model, dataset = generate_data
+        config = cmm.MCMCConfig(num_warmup=50, num_samples=100, verbose=0)
+
+        res_a = cmm.run_mcmc(model=model, dataset=dataset, config=config, yerrs=1e-2)
+        res_b = cmm.run_mcmc(model=model, dataset=dataset, config=config, yerrs=1e-2)
+
+        table = res_a.compare({"other": res_b}, dataset)
+        assert set(table.index) == {"self", "other"}
+
+    def test_loo_plots(self, generate_data):
+        """Pointwise mapping and both LOO diagnostic plots render."""
+        import matplotlib
+
+        matplotlib.use("Agg")
+
+        model, dataset = generate_data
+        config = cmm.MCMCConfig(num_warmup=50, num_samples=100, verbose=0)
+        results = cmm.run_mcmc(model=model, dataset=dataset, config=config, yerrs=1e-2)
+
+        pw = results.loo_pointwise(dataset, yerrs=0.5)
+        n_meas = len(dataset.measurements)
+        n_obs = len(model.get_observable_state_order())
+        assert pw.elpd.shape[0] == n_meas
+        assert pw.elpd.shape[2] == n_obs
+        assert pw.pareto_k.shape == pw.elpd.shape
+
+        # Influence overlay (marker size = influence) and both heatmaps.
+        assert results.plot_loo_influence(dataset, yerrs=0.5) is not None
+        assert results.plot_loo_heatmap(dataset, metric="elpd", yerrs=0.5) is not None
+        assert (
+            results.plot_loo_heatmap(dataset, metric="pareto_k", yerrs=0.5) is not None
+        )
+
     def test_initial_estimator(self):
         # Create a simple Michaelis-Menten model
         model = Model(name="test")