feat(bayesian): compare NUTS vs DEMetropolisZ for MCMC calibration

[gmermoud]

Summary

Stacked on #52 (MAP+Laplace uncertainty). Adds NUTS as an alternative to PyMC's default DEMetropolisZ sampler, switches the MAP-vs-MCMC comparison job to use it, and reports the resulting comparison in RESULTS.md.

• Add NUTS as an alternative MCMC step method (pymc_param_est.run_pymc_inference(..., step_method="nuts")), alongside the existing DEMetropolisZ — see Results below for why this replaced DEz as the comparison's MCMC arm.
• Switch the comparison job's default to --step-method nuts, add --skip-map so a site with an existing MAP+Laplace run doesn't redo it.
• Run NUTS on CPU: benchmarked as clearly faster than GPU for this workload (a 21-param site: ~28min on CPU vs. still unfinished after 50+min on GPU) — NUTS only parallelizes across chains (4x here), so a GPU's batch-throughput design doesn't pay off the way it does for MAP's 2000-way vmapped restart search.
• Fix a prior-parity bug: run_pymc_analysis built its priors without the species-derived Tmax/MaxAge bound widening that run_map_analysis already applied (load_files.literature_bound_overrides), so MAP and MCMC could silently fit against different bounds on the same site (confirmed divergent for Solling; a no-op on the 3 ICP sites below, where the file's own bound already matched literature). Now both paths pass the same bound_overrides.
• Fix run_pymc_inference silently overwriting its checkpoint_every argument with max(500, num_samples // 10) on every call, so callers could never request finer-grained checkpoints — this had left the Solling job (57 params) with zero checkpoints over a 12h run. Expose --checkpoint-every on run_comparison_site.py.
• Add RESULTS.md, a consolidated write-up of every comparison run (GPU-vmapped MAP search timing, DEMetropolisZ vs. NUTS at default and matched wall-clock budgets, and Solling's 57-parameter NUTS and matched-budget DEMetropolisZ runs), with supporting plots under assets/results/.

Verifying MAP and NUTS actually share priors and likelihood

Before trusting the comparison below, I checked whether MAP and NUTS are really scoring the same objective. They are, with one bug (now fixed, see Summary):

• Log-likelihood: identical by construction — MAP doesn't build its own PyTensor graph, it calls the same pymc_model/Run3PGLogLikeOp/build_loglikelihood_fn that MCMC uses. Both pop DIAGNOSTIC_ONLY_ERROR_NAMES from priors before model construction, so DBH/BA/Height are excluded from the scored likelihood in both paths, leaving only WS/WF/WR — confirmed both by code inspection and by each run's logged Skipping observations with no matching sigma prior: [...].
• Prior bounds: identical for every parameter on all 3 ICP sites in the Results below (wSx1000, thinPower, all err_* sigmas, and — after the fix above — Tmax). The one place they could previously diverge (Tmax/MaxAge's literature-widened bounds) happened to be a no-op for these 3 sites but was real for Solling; now fixed for future runs of any site.

Results: MAP+Laplace vs. full MCMC, 3 sites

Ran the comparison on 3 ICP plots (04.1605, 14.0003, 14.0012). Both methods fit only WS/WF/WR (biomass) — DBH/BA/Height are simulated for reference but never scored, confirmed in each run's logs (Skipping observations with no matching sigma prior: [...]). See RESULTS.md for the full write-up, including the GPU-vmapped MAP search timing, a matched-wall-clock-budget rerun of the DEMetropolisZ comparison below (does increasing its draw count by 27-40x change the outcome? no), and the Solling (57-parameter) NUTS run.

DEMetropolisZ (the prior default sampler) does not converge on this problem. With 3 chains × 5000 draws:

Site	r_hat range	ess_bulk range (of 15,000 draws)
04.1605	1.18 – 3.53	3 – 12
14.0003	1.27 – 2.70	4 – 8
14.0012	1.10 – 2.88	3 – 23

r_hat should be ≈1.0 and ess_bulk a healthy fraction of total draws; this is a near-total mixing failure (effectively 3-23 independent samples out of 15,000). NUTS, run on the same model/priors/likelihood, converges cleanly with 4 chains × 1000 draws:

Site	r_hat range	ess_bulk range (of 4,000 draws)
04.1605	1.000 – 1.000	2,126 – 5,659
14.0003	1.000 – 1.010	1,813 – 6,201
14.0012	1.000 – 1.010	1,617 – 4,869

Results: Solling, 57 parameters

Same comparison at a larger scale — Solling has 57 calibrated parameters vs. 21 on the ICP sites above. NUTS ran first here, so DEMetropolisZ was sized to a matched wall-clock budget (RESULTS.md §3) rather than a fixed draw count:

Method	Wall-clock	Draws	r_hat range	ess_bulk (min)
DEMetropolisZ	89.3 min	200,000 (ceiling)	1.25 – 3.10	3
NUTS	1d 14h04m	4,000 (4 × 1,000)	1.00 – 1.11	25

DEMetropolisZ reached its 200,000-draw ceiling — 40x its default-budget draw count — in a fraction of its 38h10m walltime cap, and still didn't converge (worst parameter fracBB0), matching the pattern on the smaller ICP sites. NUTS converged, but not as cleanly as at 21 parameters (max r_hat 1.11 vs. ≈1.00, worst parameter k), consistent with per-iteration cost and warmup difficulty scaling faster than linearly with dimension; see RESULTS.md's Open runs section for discussion.

Test plan

• uv run pytest tests/test_pymc_inference.py
• uv run ruff check / uv run ruff format / uvx ty check on all touched files
• Ran the full MAP+Laplace vs. NUTS comparison on 3 ICP sites end-to-end on the Dance SLURM cluster (see Results above)
• Ran Solling (57 params) NUTS to completion on the Dance cluster; see RESULTS.md's Open runs section for final convergence diagnostics
• Ran Solling's matched-wall-clock-budget DEMetropolisZ rerun (job trunx-demz-matched on disco) to completion — reached its 200,000-draw ceiling in 1h29m (well under its 38h10m budget) with R-hat 1.25–3.10 and min ESS(bulk) 3, confirming the same non-convergence pattern seen on the 3 ICP sites, at a larger 57-parameter scale; see RESULTS.md §3

[gmermoud]

@gmgivi @CedricTravelletti this demonstrates the advantages of differentiability by running NUTS vs DEzMetropolis side by side. Furthermore, we also run a massive GPU-based multi-init MAP run.

We shall include this in the paper.

[gmermoud]

@gmgivi aa you can see here, we no longer fit DBH directly, but only the biomasses derived from it.

[gmgivi]

@gmgivi aa you can see here, we no longer fit DBH directly, but only the biomasses derived from it.

Yes, thank you @gmermoud . I will review it, and let you know.

[gmgivi]

@gmermoud

I went through the entire PR and the results. In short, the results are correct. But the PR description was slightly misleading.

Here are few conclusions:

• MAP using 2000 restarts would take ~30 mins for 21 params site, whereas ~ 1 hour on 57 param site (solling site). There are two reasons why its taking 1 hour in solling data, first the number of parameters are comparatively large (57 params), and other is solling data is also simulated for more years (30 years).

• Run MAP on solling data simulation with 23 params, and report the results.

• NUTS converges faster compared to DEMetropolitsz, demonstrating the value of differentiability

Method	Data	# params	# monthly time step	Time
DEMetropolitsZ	solling data	23	360	17 hours
NUTS	solling data	23	360	3.5 hours