Feat/map laplace uncertainty - #52
Merged
Merged
Conversation
gmermoud
force-pushed
the
feat/map-laplace-uncertainty
branch
from
August 13, 2026 21:12
052d471 to
4c28c19
Compare
gmermoud
force-pushed
the
feat/map-laplace-uncertainty
branch
2 times, most recently
from
August 14, 2026 10:04
d02d2f2 to
b106997
Compare
Contributor
Author
|
Rebased this branch onto its current base (`fix/widen-priors` had picked up the species-dependent prior-widening fix after this branch was created, which the new commits below depend on) and force-pushed. No content lost — verified via patch-id that the previous two commits are byte-identical in diff to what's now replayed on top of the base fix. Added two commits on top of the existing Laplace-uncertainty work:
See the updated PR description for details. |
Contributor
Author
|
@gmgivi this PR features both:
Probably good to include some of these results in the paper. I will take a look next week. |
4 tasks
gmermoud
force-pushed
the
feat/map-laplace-uncertainty
branch
from
August 15, 2026 09:45
20c17bc to
c2cfca6
Compare
gmermoud
force-pushed
the
feat/map-laplace-uncertainty
branch
from
August 17, 2026 06:50
1408048 to
a702fee
Compare
5 tasks
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.
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
feat/map-laplace-uncertainty
branch
from
August 18, 2026 12:42
a702fee to
ba15c69
Compare
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.
Summary
map_uncertainty.py, wired intomap_param_est.run_map_analysisvialaplace_draws).err_DBH/err_BA/err_Heightfrom the calibration likelihood (bayesian_config.DIAGNOSTIC_ONLY_ERROR_NAMES). 3PG derives DBH from a single stand-level "mean tree" inversion ofaWS/nWS, while the ICP observations for DBH/BA/Height are built by summing per-tree allometric equations over each stand's actual size distribution — a mismatch confirmed empirically (~10cm average gap, dominant in 33/34 ICP sites) to be mostly a structural aggregation issue, not a model-fit issue. Fitting those three was trading away real WS/WF/WR accuracy for a target the model can't correctly represent. The proper fix (a distributional DBH model) is tracked inTODO.md.bayesian_comparison_plots.plot_comparison: correctbayesian_label(was hardcoded "PyMC (DEz)" even for MAP runs), a QMD-vs-"3PG-derived-DBH" overlay to visualize the aggregation mismatch, and an optionalsite_nametitle with mean RMSE and MAP log posterior.map_param_est.batched_map_search): 2000 random restarts run in parallel viajax.vmap+optax.lbfgs, replacing the old sequential CPU restart loop for the exploration phase (still polished by a finalpm.find_MAPpass sofit_laplace's Hessian lands on a genuine stationary point).scripts/run_comparison_site.py,scripts/slurm/map_mcmc_comparison.sbatch) to run a MAP-vs-MCMC comparison across sites on the Dance SLURM partition, one job per site with explicit node pinning (chacha/disco have independent filesystems).param_bound.defaultand theparameterssheet's species-value column held an independently-editable copy of the same number for every physiology parameter, on every site file (0 divergence verified across all 4 sites tested before touching anything, but Solling'sparam_boundsheet turned out to already be missing defaults for 22 unbounded params thatparametersdid have — exactly the drift this was meant to prevent).parametersis now the single source of truth for a physiology parameter's default;param_boundkeeps onlymin/max. Migrated all 39 ICP plot files plusfull_solling_data.xlsxto drop the now-redundant column.gradient_descent.py,morris_sensitivity.py, andpymc_icp_plots.pykeep their own separate reading for now (untested, out of scope here) — backport tracked insimlab-vs/TrunX#54.Test plan
uv run pytest tests/test_map_uncertainty.py tests/test_bayesian_config.py tests/test_load_files.py tests/test_map_estimation.py tests/test_bayesian_likelihood.pyuv run ruff check/uv run ruff format/uvx ty checkon all touched filesPredictive-accuracy and uncertainty-calibration validation of MAP+Laplace (comparing it against full MCMC) is in #58, stacked on top of this branch — that comparison needs a sampler that actually converges on this posterior (
DEMetropolisZdoesn't, see #58), so it's reported there rather than here.