fix(bayesian): stop duplicating physiology parameter defaults - #55
Closed
gmermoud wants to merge 16 commits into
Closed
fix(bayesian): stop duplicating physiology parameter defaults#55gmermoud wants to merge 16 commits into
gmermoud wants to merge 16 commits into
Conversation
MAP calibration returned a single mode with no uncertainty. Approximate the posterior near it by a Gaussian whose covariance inverts the local Hessian of the log posterior, and emit the result as draws so the existing MCMC tooling consumes it unchanged. The approximation is built in PyMC's unconstrained space, expanding the same objective find_MAP optimises, so it sits at a stationary point and every back-transformed draw respects the prior bounds. The Hessian central-differences the exact JAX gradient already exposed through Run3PGLogLikeOp. On the Solling stand several parameters optimise onto their prior bounds, making the MAP a boundary optimum with no local curvature. Rather than fail, condition on those: hold them at their MAP values and report the rest. Which ones to drop is decided on the correlation-scaled precision, since raw curvatures span ten orders of magnitude across parameters and would otherwise discard well-identified ones. Claude-Session: https://claude.ai/code/session_01Swe4wEikddjbJELAmukMSA
Cross-check the gradient-based Hessian against one built from log posterior values alone, and pin the sqrt(n) shrinkage of the standard errors by replicating the observations, which scales the curvature exactly without moving the mode. Claude-Session: https://claude.ai/code/session_01Swe4wEikddjbJELAmukMSA
3PG's compute_dbh derives DBH from a single stand-level "mean tree": it inverts aWS/nWS on the mean stem biomass per tree, and BA/Height are then computed from that same DBH. The ICP observations for all three are instead built by summing per-tree allometric equations over each stand's actual DBH distribution (create_data_inputs.py) — a distribution-aware aggregation the model's single-mean-tree inversion cannot match whenever a stand has real size spread. Confirmed empirically across 34 ICP plots: comparing the field-measured quadratic mean diameter against a "3PG-style" DBH obtained by inverting aWS/nWS on the plot's own observed WS/N (isolating the aggregation effect from model-fit error) gives an average 10.2cm gap, dominating over a 1.9cm average model-fit gap in 33/34 sites. Fitting err_DBH/err_BA/err_Height was therefore pushing the optimizer to trade away real WS/WF/WR accuracy for a target the model cannot correctly represent. Add bayesian_config.DIAGNOSTIC_ONLY_ERROR_NAMES and exclude those sigma priors in run_map_analysis, run_pymc_analysis, and pymc_icp_plots.run_bayesian_for_plot. DBH/BA/Height are still simulated and plotted for reference, just no longer part of the fitted likelihood. Across a 4-restart batch this drops total RMSE by 30% (WS), 62% (WF), and 73% (WR) over 36 ICP sites. The proper fix is giving 3PG a real notion of stand size distribution (e.g. a Weibull shape parameter, 3-PGmix style) instead of a single mean tree; tracked in TODO.md.
plot_comparison's DBH panel now overlays a "3PG-derived" DBH, obtained by inverting the file's default aWS/nWS on the plot's own observed WS and N (stems/ha) — the same mean-tree inversion compute_dbh uses internally, but fed real field data instead of simulated output. Next to the field-measured quadratic mean diameter this isolates the Jensen's-gap aggregation mismatch from model-fit error (see the preceding commit). Needs N, which prepare_plot_input now retains in the observed sheet (renamed from the raw stems_n column) purely for this diagnostic — not required for calibration, so a missing N no longer drops an otherwise-complete observation row. Also: - bayesian_label parameter, so the legend/title read "MAP + Laplace" rather than the previously hardcoded "PyMC (DEz)" when plotting a MAP run. - site_name parameter, adding a figure title with the site name, mean RMSE over the variables actually calibrated on (excluding DIAGNOSTIC_ONLY_ERROR_NAMES), and the MAP log posterior.
Extracts Run3PGLogLikeOp's inline log-likelihood into a standalone build_loglikelihood_fn, reused by a new batched_map_search that runs many optax.lbfgs restarts in parallel via jax.vmap instead of PyMC's sequential, CPU-only n_restarts loop. Wired into run_map_estimation/ run_map_analysis as n_vmap_restarts/n_vmap_steps: the winning restart still gets polished by pm.find_MAP so fit_laplace sees a genuine PyTensor stationary point. Also generalizes plot_comparison's hardcoded "HMC (NUTS)" label to a hmc_label parameter, so a plain MCMC run can share that overlay slot with a correct legend for an upcoming MAP-vs-MCMC comparison.
scripts/run_comparison_site.py runs run_map_analysis (vmap-batched MAP search) then run_pymc_analysis for one site, with param_to_optimize=None so each file's own bounded parameters are used (some ICP plots leave aWS/nWS unbounded, unlike the full Solling file FIT_PARAMS was written for). scripts/slurm/map_mcmc_comparison.sbatch is a 4-task array job on the Dance partition (chacha+disco H100s), one GPU per site.
The account is capped at 1 GPU at a time on the Dance partition's standby QoS, and standby jobs are preempted (SIGTERM) whenever a normal/premium/standard job needs the resources. Without --requeue a preempted task just dies instead of going back in the queue. --array=...%1 keeps us within the 1-GPU cap; MCMC checkpoints, so a requeue only redoes the (fast) MAP+Laplace half.
Progress-bar output otherwise block-buffers once redirected to the log file, making an actively-computing task look stalled for long stretches with nothing new in the log.
DEMetropolisZ (derivative-free) barely mixes on the real 3PG posterior (observed ESS of 3-12 out of 5000 draws on an ICP plot), while the model's log-likelihood already exposes a JAX gradient via Run3PGLogLikeOp.grad (the same one MAP's L-BFGS and the vmap search use). Adds step_method="nuts" to run_pymc_inference/run_pymc_analysis, running pm.NUTS on the identical Uniform-prior model MAP already uses (as opposed to the separate NumPyro/HMC pipeline in parameter_estimation.py, which silently uses a different, peaked prior) so MAP and MCMC stay directly comparable. Verified on the synthetic stand: r_hat ~1.00-1.01 and ESS 111-311/500 draws, vs. DEMetropolisZ's near-total lack of mixing on real data.
--step-method now defaults to nuts (DEMetropolisZ showed near-zero ESS on real 3PG posteriors, see prior commit). --skip-map lets a site whose MAP+Laplace already finished reuse it instead of redoing the expensive vmap search just to redo the MCMC arm with a different sampler.
Benchmarked NUTS as clearly faster on CPU than GPU for this workload (a 21-param site: ~28min on CPU vs. still unfinished after 50+min on GPU) — NUTS only parallelizes ~4x across chains, so GPU's batch throughput design doesn't help the way it does for MAP's 2000-way vmapped restart search. Drops --gres/GPU and forces JAX onto CPU. Also switches from a single array job to one job per site with an explicit --nodelist, since chacha and disco have independent filesystems and each site's map/ only exists on the node it was computed on — the array's implicit node assignment couldn't guarantee that pairing.
run_pymc_analysis loaded priors without the species-derived bound_overrides that run_map_analysis already applied, so MAP and MCMC could silently fit different Tmax/MaxAge bounds on the same site (confirmed divergent for Solling). Also documents the wSx1000/thinPower identifiability finding from the MAP-vs-NUTS comparison as a TODO, with next-step options.
gmermoud
force-pushed
the
fix/consolidate-physiology-param-defaults
branch
from
August 15, 2026 09:45
d9a7c71 to
b1ab9ae
Compare
run_pymc_inference silently overwrote its checkpoint_every argument with max(500, num_samples // 10) on every call, so callers could never request finer-grained checkpoints. This left the 12h solling NUTS job (57 params) with zero checkpoints: the first chunk needed the full tuning phase plus 500 draws across 4 chains to finish before anything was saved. Expose --checkpoint-every on run_comparison_site.py and give solling a much smaller value in the sbatch job so a long CPU run banks progress incrementally instead of all-or-nothing.
Add RESULTS.md summarizing GPU-vmapped MAP search wall-clock, DEMetropolisZ vs NUTS convergence at default and matched wall-clock budgets, and the completed Solling (57-parameter) NUTS run, with supporting plots.
param_bound.default and the parameters sheet's species-value column held an independently-editable copy of the same number for every physiology parameter, on every site file. Verified 0 divergence across the 3 already-compared ICP sites and full_solling_data.xlsx before touching anything, but solling's own param_bound sheet turned out to already be missing defaults for 22 unbounded params that parameters did have — exactly the drift this was meant to prevent. parameters (per-species runtime value) is now the single source of truth for a physiology parameter's default, used both as its fixed value when not optimized and its MAP/MCMC seed when it is. param_bound keeps only min/max. error_param is untouched — sigma priors have no parameters-sheet counterpart, so their default was never duplicated. Migrated all 39 ICP plot files plus full_solling_data.xlsx to drop param_bound's now-redundant default column (data/ is gitignored, so this only shows up in the code and test diff here). gradient_descent.py, morris_sensitivity.py, and pymc_icp_plots.py keep their own separate param_bound/error_param reading for now (untested, out of scope for this change) — backport tracked in TODO.md.
Tracked as #54 instead of TODO.md.
gmermoud
force-pushed
the
fix/consolidate-physiology-param-defaults
branch
from
August 17, 2026 04:43
b1ab9ae to
089dfb0
Compare
Contributor
Author
|
Folded into #52 — the physiology-defaults consolidation now lives in feat/map-laplace-uncertainty directly. |
An error occurred while trying to automatically change base from
feat/map-laplace-uncertainty
to
main
August 18, 2026 12:43
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.
Stacked on #52 (base branch:
feat/map-laplace-uncertainty).Summary
While digging into how the MAP/MCMC comparison's ICP sites and
sollingdiffer (they turn out to share the exact same 82-parameter physiology table — sites just differ in how many of those rows are marked optimizable), found that every site file was carrying the same physiology default in two independently-editable places:param_bound.default— used only to seed MCMC/MAP initparameters's species-value column — the actual value 3PG runs a parameter at when it isn't being optimizedVerified these agreed on every row across
full_solling_data.xlsxand the 3 ICP sites already compared in #52 before touching anything — but the safety check in the migration script caught a real divergence insollingitself: 22 unbounded params had a value inparametersbutNaNinparam_bound.default(harmless in practice, since those params are never requested, but exactly the kind of drift two copies of the same number invites).Changes
load_files.py:param_boundnow carries onlymin/max. A physiology parameter's default (used both as its fixed value when not optimized, and its MAP/MCMC seed when it is) is read from theparameterssheet — the pipeline's existing, necessary single source of truth for per-species runtime values.error_param(sigma priors) is untouched, since those have noparameters-sheet counterpart.full_solling_data.xlsxto dropparam_bound's now-redundantdefaultcolumn.data/is gitignored, so this only shows up in the code/test diff here, not the PR diff.tests/test_load_files.pyfor the new xlsx sourcing path (load_param_defaults_from_filepulling physiology defaults fromparameters, priors unaffected, and a clear error on a multi-speciesparameterssheet — which the pipeline doesn't support yet, deliberately, since every real site file today is single-species).gradient_descent.py,morris_sensitivity.py, andpymc_icp_plots.pyeach have their own separate, untestedparam_bound/error_paramreading and weren't touched here — backport tracked in #54.Test plan
uv run ruff check/uv run ruff format/uvx ty checkuv run pytest tests/test_load_files.py(new + existing cases)test_map_uncertainty.py,test_bayesian_config.py,test_load_files.py,test_map_estimation.py,test_bayesian_likelihood.py,test_pymc_inference.py) — 32 passed, run both before and after migrating the real data filesprepare_data+load_priors_from_file/load_param_defaults_from_fileverified end-to-end against all 40 migrated real site files