Point estimation: trustworthy find_MAP, new find_MLE - #1192
Closed
fmuia wants to merge 9 commits into
Closed
Conversation
- New `hssm.optimize` module holding the machinery shared by point estimation: `PointEstimate` (a `dict` subclass, so it still works anywhere a plain point dict did, notably `sample(initvals=...)`), the start-value/jitter helpers, the convergence audit, the `float32` safety net, an `scipy.optimize.minimize` driver and standard errors. - Jitter is applied in transformed space, where the transforms leave no bound to cross; a relative jitter on the constrained scale takes `z = 0.95` to `1.045`, which PyMC then rejects as a bad start. - Gradients go through `rewrite_pregrad` first, as `pymc.Model.dlogp` does, so HSSM differentiates the same stabilized graph PyMC does. - `PointEstimate` is exported at the top level. Refs #1102
#1102) `find_MAP` delegated to `pm.find_MAP` with no starting point, so PyMC rebuilt its own (`t=2.0`, `a=2.0`) and ignored HSSM's processed initial values. For many models -- including a plain hierarchical DDM on `cavanagh_theta` at float64 -- the *gradient* of the log-density is non-finite there while the logp is finite, so PyMC's start check passed, L-BFGS-B aborted at `nit=0`, and the untouched start was stored as `model.map` with no warning. - `start` now defaults to `model.initvals`. - The optimizer result is audited; a failed run warns and returns `None` (or raises under `strict=True`) and is never cached. - `sample(initvals="map")` raises instead of silently falling back to PyMC's default point when the MAP did not converge. - New `n_starts`, `strict`, `se`, `seed` and `method` parameters; the method that *actually* ran is recorded, since a likelihood without a gradient is switched to Powell behind the caller's back. - Only `SamplingError` is treated as "PyMC rejected this start", so a bad `method=` surfaces as the `ValueError` SciPy raised rather than a convergence warning. - `sample(sampler="laplace")` warns: bambi's laplace path offers no initval plumbing, so it hits the same bad start. - `model.map`'s guard moves from a falsy check to `is None`, so a computed-but-empty estimate no longer reads as "not computed". Refs #1102
Maximizes the *observed* log-likelihood -- the priors are dropped -- so the result is a frequentist point estimate rather than a posterior mode. - Raises on hierarchical models: dropping the priors leaves the group-level scale unidentified (under HSSM's default non-centered parameterization the likelihood is invariant under rescaling the offsets against sigma; under the centered one sigma does not enter the likelihood at all), so the optimum is a flat ridge. Use `find_MAP`, or pass `allow_unidentified=True`. - Raises on models with `pm.Potential` terms, which `observedlogp` would silently drop. - Optimizes in transformed space; the transforms are bijective, so no Jacobian correction is needed and the parameters stay inside the range the likelihood network was trained on. - The objective, its gradient and the point expander are compiled once and reused across starts; the scorer reuses the compiled objective and the method is resolved from the compile rather than a second gradient probe. - New `model.mle` property. Refs #1102
- New `tests/test_optimize.py`: recovery, the non-finite-gradient regression that motivated the work, failure reporting and caching, multi-start, `float32` fallback, standard errors, `PointEstimate` (pickle/copy/ArviZ round trips), `sample()` integration and helper units. Slow markers cover the full hierarchical fit, cloudpickle, save/load and the `approx_differentiable` backends. - Regression tests for the defects found in review: a bad `method=` raises instead of warning, jitter near an upper bound no longer burns starts, integer-valued starts are still jittered, an untouched zero reads as "did not move", the failure message does not advertise `return_raw=True`, and `find_MLE` neither probes the gradient twice nor compiles the objective or the point expander per start. - `tests/test_initvals.py` and the RL suite assert the new `model.map` contract and the extra-field refresh. Refs #1102
- New "Find MAP and MLE point estimates" tutorial under How-to guides > Sampling and diagnostics, alongside variational inference as the other non-MCMC route. It ships with stored outputs and is listed in mkdocs' `execute_ignore`, so docs builds do not re-run its hierarchical fit and MCMC draw; notebook CI still executes it. - New `hssm.PointEstimate` reference page, with attribute docs enabled because for this class the attributes are the API. - `find_MAP`/`find_MLE` added to the HSSM and rl API pages. - Changelog entry. Refs #1102
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…1102) mypy does not carry the `method is None` narrowing through the intermediate `fallback` boolean, so the float32 branch returned `str | None` from a `-> str` function and failed CI. pyrefly narrows it, which is why only the mypy step caught it. No behavior change: both branches are unchanged, and each is covered by test_float32_warns_and_falls_back_to_powell and test_float32_honours_an_explicit_method. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FeXVhc4ccwyCiSkdK7EpVV
…points (#1102) Two findings from an adversarial review of the point-estimation work. A failed `find_MAP`/`find_MLE` returned None without clearing `_map_dict`/ `_mle_dict`, so the "a failed estimate is never cached" contract held only on a first run. After an earlier successful call, `model.map` returned the stale estimate instead of raising, and `sample(initvals="map")` slipped past its `_map_dict is None` guard to initialize NUTS from it -- the silent wrong-start failure that guard exists to prevent. Clearing before `report_failure` covers the `strict=True` path, where that call raises. `make_scorer` indexed every value variable without checking, so scoring an estimate produced with `include_transformed=False` raised a bare KeyError past the `(ValueError, FloatingPointError)` handler. It now raises with the cause named. Not `-inf`: that is the answer for a point of zero density, and reusing it would report an unscoreable estimate as infinitely improbable. Each of the three new tests was confirmed to fail without the corresponding fix. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FeXVhc4ccwyCiSkdK7EpVV
`ddm_data` called `simulator()` without `random_state`, so every CI job drew a different 500-trial dataset while the recovery tests asserted the estimate lands within 0.2 of the truth. At 500 trials `v` is the least well-identified parameter: across 12 seeds its estimate sits near 0.585 (SD 0.055), leaving the assertion about two SDs from its bound, so the suite failed a small fraction of runs. `run_tests (3.12)` hit it on d6c4cf0, recovering v = 0.877 while 3.13 and 3.14 passed on identical code. Seed 10 chosen for margin: worst deviation from truth 0.058 across both MAP and MLE. Verified the whole fast module still passes (73 passed), since the fixture is module-scoped and feeds every test in the file. Note: `test_find_map_approx_differentiable` has the same unseeded draw, left alone here because its ONNX likelihood cannot be exercised in this environment to confirm a seed choice. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FeXVhc4ccwyCiSkdK7EpVV
…1102) Findings from an independent adversarial review of the branch. - `standard_errors` differentiated outside the model's support. A bound enters the graph through `pt.switch`, so past it the log-density is `-inf` while the gradient stays finite: the Hessian came back well conditioned, every guard passed, and an estimate pinned to a bound got a confident standard error where the docstring and tutorial promise `NaN`. Observed `se[p_outlier] = 0.0061` at `p_outlier = 0`, implying a 95% interval straddling negative probability. Each stencil point is now checked against the density itself, which is the only thing that reveals the boundary. All entries go `NaN`, not just the offending one, because the errors inverting the joint Hessian are all corrupted by one unusable column. - `report_failure` shared one message between both paths, so the `strict=True` RuntimeError claimed the call "returned None" and recommended the mode already in force. - `find_MAP` compiled a second logp evaluator for a number `pm.find_MAP` already returns as `opt_result.fun`. Verified equal bit-for-bit under L-BFGS-B, Powell, Nelder-Mead and BFGS; a scorer is now built only if a method leaves `fun` unset. Worth ~10% of a run on an analytical DDM and more on an ONNX likelihood. - Multi-start expanded each jittered candidate to constrained space for a `start=` that accepts transformed names directly, costing a compiled function and, on regressions, every trial-wise deterministic per candidate. - `dims_and_coords` copied the model's coords wholesale, hanging the per-trial `__obs__` index off every estimate for nothing to read. - Four docstrings corrected: `success` is hardwired `True`, `opt_result` is never `None` on a returned estimate, the raw mapping loses its transformed entries under `include_transformed=False`, and that kwarg is consumed by HSSM rather than forwarded to PyMC. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01FeXVhc4ccwyCiSkdK7EpVV
Contributor
Author
|
Superseded by #1218, which carries the same work on a clean linear history rebased onto current Rewriting the history here wasn't possible (the branch ruleset blocks force-pushes) so this went to a new branch instead. |
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
2 participants
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.
find_MAPdelegated topm.find_MAPwith no starting point, so PyMC used its own (t=2.0,a=2.0) and ignored HSSM's processed initial values. For many models — including a plain hierarchical DDM oncavanagh_thetaat float64 — the gradient is non-finite there while the logp is finite, so PyMC's start check passed, L-BFGS-B aborted atnit=0, and the untouched start was cached asmodel.mapwith no warning.find_MAPnow starts frommodel.initvals, audits the optimizer result, and never caches a failed run (warns and returnsNone, or raises understrict=True).sample(initvals="map")raises instead of silently falling back to PyMC's default point.find_MLE()maximizes the observed log-likelihood only. It refuses hierarchical models (without the priors the group-level scale is unidentified — a flat ridge, not a point) andpm.Potentialmodels, withallow_unidentified=Trueas an escape hatch.PointEstimate— adictsubclass, so existinginitvals=code keeps working — carrying.params,.logp,.se,.to_dataframe()and.to_datatree(). Newn_starts,strict,se,seed,methodandprogressbarparameters; the optimizer that actually ran is recorded.rewrite_pregrad, aspymc.Model.dlogpdoes.sample(sampler="laplace")now warns: bambi's laplace path offers no initval plumbing, so it hits the same bad start.hssm.PointEstimatereference page.Verification: 699 fast tests, 6 slow tests, ruff/pyrefly/mypy and
mkdocs build --strictall pass.Closes #1102
🤖 Generated with Claude Code