Sampling-based estimation with NUTS by kiante-fernandez · Pull Request #324 · GWeindel/hmp

kiante-fernandez · 2026-08-05T16:22:19Z

Builds on #244, merge after that one.

Adds MCMCEstimator next to EMEstimator. Same likelihood, reimplemented in JAX so it can be differentiated, wrapped in a PyTensor Op with a hand-written VJP. It is registered as an observed distribution rather than added as a potential, which populates idata.log_likelihood, so the number of events can be picked by cross validation with a standard error instead of by walking up local maxima.

New modules: jax_likelihood.py (forward-backward), pytensor_op.py (PyMC wrapper, compiles to JAX under numpyro), mcmc.py (model and posterior summary). fit_likelihoods and max_explored_scale in base.py let the search strategies read likelihoods without touching EM internals; both take estimator= in their constructor. Sampling dependencies are an mcmc extra, and MCMCEstimator is imported lazily so hmp.estimators works without them.

Checks: likelihood matches numpy, per-trial Op matches the scalar one, gradient matches finite differences. Eight replicates on simulated data put 187/200 true values inside their 94% interval. LOO picks the true event count when used for model comparisons. On SAT speed it picks three events, as shown. All six tutorial operations run under both estimators. Suite at 132.

Not supported: any duration distribution but the gamma, and fixed_channel_pars/fixed_time_pars. Both raise rather than fit something else.

jitter defaults off. The perturbation is uniform on (-1, 1) unconstrained, about 2.7x on a log scale, and cold starts do not converge with it on. It also spreads the chains, which is what separated_modes reads, so turn it on when the question is multimodality. Need to think about this more.

The default backend runs one chain at a time: where multiprocessing forks its workers, JAX in the parent can deadlock them. numpyro runs chains in parallel. Another issue that will need more work.

An r hat near 1.8 here would come up sometimes; wrong event count, too few components, group differences that are not there. More draws do not help.

Priors are provisional. estim_probs gained an intermediates argument used only by the reference generator, worth a look in review.

Resolves MergeError when concatenating MNE-Python epochs with conflicting metadata columns (e.g., 'block' and 'Block'). This commit modifies `hmp/io.py` to drop both 'block' and 'Block' columns from metadata before concatenation, preventing the error.

- Added `utils.py` with convergence checkers and parameter validation functions. - Enhanced `EventModel` to support new estimator interface, allowing for MCMC and EM fitting. - Introduced `MCMCEstimator` for Bayesian parameter estimation using PyMC. - Updated simulation functions to handle event width and ensure directory creation. - Created comprehensive tests for estimation utilities, EM and MCMC estimators. - Improved plotting functions to visualize MCMC results. - Added optional dependencies for MCMC support in `pyproject.toml`.

…ng points and improve parameter handling in MCMC fitting. Update tests to reflect changes and remove unnecessary imports.

…mpatibility! - Integrated JAX-based likelihood computation into MCMCEstimator. - Added parameters to fix channel and time parameters during estimation. - Improved handling of initial time parameters to ensure positivity and replace NaNs with defaults. - Updated EventModel to support backward compatibility for existing model parameters. - Refactored tests for MCMCEstimator to reduce data size and improve execution speed. - Adjusted likelihood assertions in tests to allow for small numerical differences (? is this okay?). - Updated optional dependencies in pyproject.toml to include JAX and JAXlib. - Enhanced cumulative estimation tests to verify backward compatibility.

- Merged latest devel branch from upstream repository - Preserved estimation interface refactor: BaseEstimator, EstimationResult, EMEstimator - Added estimator parameter to EventModel.fit() for new estimation framework - Maintained backward compatibility for traditional EM approach - Excluded MCMC implementation to focus only on interface changes - Resolved merge conflicts in __init__.py, distributions.py, simulations.py, pyproject.toml - Added ruff dependency from devel branch for linting

- Added tests for EstimationResult and EMEstimator classes - Tests verify proper initialization, parameter validation, and basic functionality - Excluded MCMC-specific tests to focus on core estimation interface

- Added hmp/estimators/utils.py with general estimation utilities: - ConvergenceChecker abstract base class - RelativeLikelihoodConvergence and ParameterConvergence checkers - Parameter validation and initialization functions - Log-likelihood computation utilities - Updated estimators/__init__.py to export utility functions - Restored complete test suite in test_estimators.py with all utility tests - These utilities support the estimation interface without MCMC dependencies

- Fixed null pointer bug in EMEstimator: added proper null checks for fixed_channel_pars and fixed_time_pars - Added missing docstring to EMEstimator.__init__ method - Verified all error handling, type hints, and documentation - Confirmed backward compatibility and integration works correctly - All existing tests pass, estimation interface is production ready

…stimator

…fit method for improved parameter estimation

- Added tests for the new estimation interface in `tests/test_estimation_interface.py`.

Resolves conflicts against upstream's io package, PatternData, and projectors refactors. Takes upstream's rewritten event.py as the base; the estimator extraction in hmp/estimators/ still needs porting to the new EM.

The estimator now owns the starting-point loop, its parallelization and the EM iterations; the model keeps data preparation, grouping and parameter formatting. EM is moved verbatim from EventModel so results are bit-identical, and EventModel.fit() gains an estimator argument to swap the method. Addresses review: drops the unused estimators/utils.py, annotates PatternData and model, makes model positional, exposes fitted, and removes the EstimationResult array coercion.

The tests targeted TrialData, hmp.preprocessing and fit() kwargs that no longer exist. They now cover what the refactor actually needs guarding: an injected EMEstimator reproduces the default fit exactly, and a coarser estimator demonstrably changes it, so the injection cannot pass by being ignored. EventModel.fit() returns the EstimationResult and stores it as estimation_result, exposing convergence, iteration count and uncertainty, which Bayesian estimators will need.

self.params was written and never read, and the **kwargs chain that fed it let constructor typos pass silently. n_cor sat on the model only to hand a constant to a parameter with the same default. The best-of-starting-points branch became unreachable once the count came from the array length, and fit() re-copied five arrays em() had already built. tests/test_estimators.py was a strict subset of test_estimation_interface.py.

data() re-read both .fif recordings and re-ran the projection on every call, which was seven times and about 1.2s each. Splitting it into a module-scoped projection and a per-test PatternData keeps the tests isolated and takes the file from 10.4s to 3.1s.

With the default single starting point, fit(cpus=N) spawned N workers to run one task and then forced that task serial anyway, costing seconds of process startup for no parallelism. Run in-process in that case and size the pool to the number of starting points otherwise. Adds the first test of the cpus>1 path, which nothing covered before. It passes explicit starting points because gen_random_stages draws from an unseeded default_rng.

estim_probs already computed the log-likelihood of each trial and summed it away on the same line. It now keeps that vector, _estim_probs_groups scatters each group's values back into the original trial order, and two public methods expose it: log_likelihood(..., per_trial=) and event_probabilities(), the latter evaluable at any parameters rather than only at a fitted model. Pointwise values are what LOO and WAIC need to compare numbers of events, and per-draw event probabilities are what turns a posterior over parameters into a posterior over by-trial event times. Comparing per-trial vectors is also a much stronger check between implementations than comparing a sum, where sign and indexing errors cancel. EM now scores parameters through the public surface, and fit() no longer requires iteration traces in the diagnostics, which an estimator that does not iterate cannot supply.

The data is float32 by default, so a float32 running sum and a float64 one differ by about 1e-7 whatever the number of trials. The assertion only held because a single participant here is two identical noiseless trials, and doubling a float is exact. Using all the trials in float64 makes the invariant exact, and means the check fails if the float64 path ever falls back to float32, which sampling-based estimators will need to compare against.

The likelihood is reimplemented in JAX so it can be differentiated, wrapped as a PyTensor Op, and sampled through PyMC. Registering the Op with jax_funcify lets the graph compile to JAX, so numpyro differentiates it and no gradient has to be derived by hand; a gradient Op is provided anyway so the default backend also works. Checked against reference values captured from the numpy implementation in float64, over 18 parameter settings spanning 1 to 3 events and scales short and long enough to push mass past the support. Values agree to 2.5e-15 at every stage of the forward-backward pass, and gradients match finite differences. Priors are placeholders. They are weakly informative and centred on the data, enough to sample but not yet agreed.

Covers what was checked by hand while porting: the JAX likelihood against reference values stage by stage, gradients against finite differences, the Op against the function it wraps on both backends, and the PyMC model logp against the likelihood plus its priors. Two of them record behaviour rather than requirements. Parameters stored as float32 cap agreement at about 1e-7, and EM stops where the likelihood is still improvable in the scale direction, because the scale update inverts the mean of an untruncated gamma while the likelihood uses a truncated and renormalised one. Both should start failing if either is changed.

az.summary rounds to two decimals, so taking the maximum from it and comparing against 1.01 was testing a rounded number. On 62 trials that reported 1.0100 and declared the fit unconverged when the actual value is 1.0056 with 7367 effective samples and no divergences.

The substantive claim for a sampler is that it estimates the same thing the existing estimator does, so this asserts every scale is within one posterior standard deviation of EM's. Absolute recovery of the simulated parameters is deliberately not asserted. The posterior means sit slightly above the values dataset_c was generated from, but EM shows the same offset on the same data, so it is a property of the method rather than of this estimator.

A fit reports one set of event probabilities, at one set of parameters. posterior_event_probabilities evaluates them at a sample of posterior draws instead, so the uncertainty in the parameters reaches the quantity the method exists to estimate. Computed on request rather than during sampling, since the result is (draws, trials, samples, events) and usually only part of it is wanted.

Each distinct parameter is one random variable indexed into the positions that share it, so the sampler moves in the free space and the sharing holds by construction. EM instead averages tied parameters after each update, which is a projection and would not leave a sampler targeting the right posterior. Codes in the maps are compared down each column: groups carrying the same code at a given event share that event's parameter, while the same code at another event is a separate parameter. Reading them as global identifiers would collapse every event onto one parameter whenever the map is all zeros, which is the default for an ungrouped model. Groups that omit events are still refused rather than silently mishandled.

The grouped path summed one Op per group without ever being compared against the numpy grouped likelihood. It does match, to 3.7e-16, including the case where the groups have different max_duration and so normalise the pmf over different supports. That should have been checked before relying on it. The previous commit said EM's residual scale gradient comes from mean_to_scale inverting an untruncated mean while the likelihood uses a truncated one. That is not supported. Both facts hold separately, but across stages the size of the moment gap does not track the gradient at all: the stage with no truncation and a negligible gap carries the largest gradient, and pooled correlation is -0.09. EM stopping short of a stationary point is still established; the reason is not.

…ation gen_random_stages tested its loop condition before drawing anything, against durations initialised to zero. Whenever every location is zero the condition was false immediately, so no draw happened and every random starting point came back as all zeros, which mean_to_scale turns into a zero scale. That is the case for a single event, and for any number of events when location=0 is passed. EM hid this by keeping whichever starting point scored best and discarding the degenerate ones. Sampling cannot, since the initial point has to be valid, which is how it surfaced. Also requires at least one sample per stage: a zero-length stage has a zero scale whatever the censoring asks for.

A misspecified number of events makes the posterior multimodal, since the events can attach to different features of the data. Chains then settle at log probabilities far apart and r_hat rises, but r_hat alone does not distinguish that from slow mixing, and the two call for opposite responses: more draws help one and not the other. The result now carries whether the chains are separated, the spread of their log probabilities, and the per-chain levels. On this data the correct number of events gives a single mode, and one event gives chains around -90 and -130.

PyMC runs chains in forked processes, and forking after JAX has been used deadlocks instead of failing, so nuts_sampler="pymc" hung indefinitely with no error. Chains now run sequentially there. The numpyro sampler stays inside JAX and is unaffected. That backend is also about ten times slower per draw, since the Op is called through Python rather than compiled into the graph, which is now noted where the option is documented.

Taking the prior standard deviation from the spread of the data made it informative: on the simulated data the true magnitudes reach 3.78 while that spread is 1.16, and the posterior came back pulled towards zero, with a larger mean error than EM, which has no prior at all. Widening it removes both. Mean absolute error goes 0.198 at the data scale, 0.182 at three times, 0.180 at ten, against 0.181 for EM, and the magnitudes move outward towards the simulated values throughout. The residual error is unchanged by widening further and matches EM's, so it belongs to the method rather than to the prior. The priors are still provisional.

Forking after JAX has run deadlocks instead of raising, so that path once hung with no error and no output. Chains run sequentially there now, and this asserts a fit completes and returns both chains. Recording the comparison it came from, on 15 trials with every sampler given enough draws to converge: NUTS through numpyro, NUTS through the default backend, and Slice all agree, the latter two within 0.029 and 0.086 posterior standard deviations. Slice never evaluates the gradient, so its agreement is the check that does not share a failure mode with comparing the gradient against finite differences. Relative cost was 1, 17 and 81.

PyTensor merges Ops whose props compare equal, so the key hashes the arrays rather than their shapes. Only the gamma density is implemented; other duration distributions are refused.

- register the likelihood as an observed distribution, so every fit carries per-trial log-likelihoods for cross validation - start every chain from the best starting point instead of one each; chains are pooled, so spreading them mixes worse basins into the result - settle the sign of lp by probing the model, not assuming the backend - route jitter and init per backend; jitter off by default, it moves log-scale parameters by up to 2.7x and breaks cold starts - refuse fixed parameters rather than silently estimate them - keep the stored posterior paired with the model that produced it - a single chain has no r_hat; do not fail convergence on it

fit_likelihoods and max_explored_scale replace direct reads of EM internals, and both methods take the estimator in their constructor.

The import is deferred so hmp.estimators works without the sampling stack installed.

Starting points are no longer spread by default, and chains that start together tend to reach the same mode, so the separated-modes check has to request the spread it reads.

kiante-fernandez · 2026-08-05T16:25:03Z

To be clear about scope: this is a provisional first pass at getting an MCMC estimator working, not a finished feature. The diff looks large because it sits on top of #244 and GitHub is showing both until that one merges; the part that is actually new here is thirteen files.

Everything in it is what was needed to get the prototype running end to end, so some of it is groundwork rather than the sampler itself. Happy to split anything out if it is easier to review that way.

GWeindel · 2026-08-05T17:30:14Z

I'll need quite some time to go through this but this seems like a perfect first take on estimating HMP using MCMC. I'm not so surprised on the jitter effect. I do think that for a first time the best option is to have a mixture between EM and MCMC at least to ensure that the chains do not get stuck in weird parameter space while still ensuring that it is explored appropriately. So yes definitely some thinking needed here.

I suggest we merge #244 once you accept kiante-fernandez#1. Then I think this can stay a single PR.

I'll be away from GitHub from Friday up to the 23rd of August but then I look forward to dig into this more deeply.

init_from="em" fits with expectation maximization first and samples from its result. A maximizer gets into a sensible region of this likelihood more reliably than a sampler started elsewhere; what the sampler adds is the spread around it. The EM likelihood is recorded alongside so the two can be compared.

Trimming on the likelihood prefers the largest model, because adding an event cannot lower it. select() compares the submodels by leave-one-out cross validation and returns the smallest one whose distance from the best is within the error on that distance. It refuses when only one number per fit is available, as with EM.

# Conflicts: # hmp/estimators/__init__.py # hmp/models/event.py # tests/test_estimation_interface.py

JoKra1 · 2026-08-24T13:51:33Z

This looks really cool, and I would be very happy with JAX as an extra dependency - this will make some of my work easier as well! 👍🏻

One suggestion regarding efficiency @kiante-fernandez: when you only need the log-likelihood but not the event probs (as for NUTS) you do not need to evaluate the backward recursion for a HMP-like model. Rather than evaluating two convolutions in the loop, computing the log-likelihood only requires multiplying the last event forward variables by the flipped pmf of the last interval (see trial_likelihood2 below - this is for the basic event model which was easier for me to modify now since I have not checked your code in detail yet). The trial-likelihood is then the sum over the last event not the first!

for event in np.arange(
            1, n_events
        ):  # Following stage transitions integrate previous transitions
            add_b = backward[:, :, event - 1] * probs_b[:, :, event - 1]  # Next stage in back
            for trial in np.arange(n_trials):
                # convolution between gamma * gains at previous event and event
                forward[:, trial, event] = np.convolve(forward[:, trial, event - 1], pmf[:, event])[
                    :max_duration
                ]
                # same but backwards
                backward[:, trial, event] = np.convolve(add_b[:, trial], pmf_b[:, event])[
                    :max_duration
                ]
            forward[:, :, event] = forward[:, :, event] * probs[:, :, event]

        trial_likelihood2 = forward.copy()
        for trial in np.arange(n_trials):
            trial_likelihood2[: durations[trial], trial, -1] *= np.flip(pmf[: durations[trial], -1]) # Must flip only the part up to dur
        
        # re-arranging backward to the expected variable
        backward = backward[:, :, ::-1]  # undoes stage inversion
        for trial in np.arange(n_trials):  # Undoes sample inversion
            backward[: durations[trial], trial, :] = backward[: durations[trial], trial, :][::-1]
        eventprobs_raw = forward * backward
        trial_likelihood2 = np.clip(trial_likelihood2, 0, None) 
        eventprobs = np.clip(eventprobs_raw, 0, None)  # floating point precision error
        trial_likelihood = np.log(
            eventprobs[:, :, 0].sum(axis=0)
        )  # sum over max_samples to avoid 0s in log
        likelihood = np.sum(trial_likelihood)
        trial_likelihood2 = np.log(
            trial_likelihood2[:, :, -1].sum(axis=0) # MUST SUM OVER LAST EVENT!
        )
        print(likelihood, np.sum(trial_likelihood2))

likelihood and np.sum(trial_likelihood2) are equivalent up to rounding error occurring during convolution of backward (so np.sum(trial_likelihood2)) should even be more stable.

kiante-fernandez and others added 30 commits July 15, 2025 14:42

Enhance MCMCEstimator and EventModel: Add support for multiple starti…

e0d7e81

…ng points and improve parameter handling in MCMC fitting. Update tests to reflect changes and remove unnecessary imports.

Add estimation framework tests

4ca187a

- Added tests for EstimationResult and EMEstimator classes - Tests verify proper initialization, parameter validation, and basic functionality - Excluded MCMC-specific tests to focus on core estimation interface

Add test for estimation interface integration

869c15c

Remove estimation interface test file

58ffcb9

Add integration test for estimation interface with TrialData setup

a8b3754

Remove obsolete integration tests for estimation interface and MCMC e…

1c4b050

…stimator

Enhance EMEstimator to support multiple starting points and refactor …

e128257

…fit method for improved parameter estimation

Remove deprecated JAX and MCMC estimator implementations

d0f1729

Remove MCMC optional dependency from project configuration

32bbe5b

lintting and testing

d93e088

- Added tests for the new estimation interface in `tests/test_estimation_interface.py`.

Run release workflow only from core repo + node 20 deprecation bumps

d3c85db

Merge branch 'GWeindel:devel' into devel

99681a7

Merge upstream/devel into estimation-interface-merge

00ea919

Resolves conflicts against upstream's io package, PatternData, and projectors refactors. Takes upstream's rewritten event.py as the base; the estimator extraction in hmp/estimators/ still needs porting to the new EM.

Drop redundant inline comments in estimation interface tests

01c0dfe

Keep only the noqa codes that fire

1937973

Removed leftover

19f0018

Remove event_width

43c2283

kiante-fernandez added 19 commits August 2, 2026 02:49

Add a per-trial likelihood Op, jitted and keyed by its data

fe012fb

PyTensor merges Ops whose props compare equal, so the key hashes the arrays rather than their shapes. Only the gamma density is implemented; other duration distributions are refused.

Let search strategies read likelihoods through the estimator interface

e7a2230

fit_likelihoods and max_explored_scale replace direct reads of EM internals, and both methods take the estimator in their constructor.

Raise a clear error for per-iteration diagnostics on sampled fits

9ab8863

Export MCMCEstimator and declare the mcmc extra

cc909e8

The import is deferred so hmp.estimators works without the sampling stack installed.

Cover Op identity, likelihood semantics, and the estimator passthrough

9d2b7d7

Ask for jitter where the test needs the chains to disagree

d585dac

Starting points are no longer spread by default, and chains that start together tend to reach the same mode, so the separated-modes check has to request the spread it reads.

GWeindel mentioned this pull request Aug 5, 2026

Paralellized the cpus>1, starting_points=1 situation (CumulativeModel) #323

Merged

kiante-fernandez and others added 6 commits August 5, 2026 15:43

Cover EM seeding and selection by cross validation

aaf744a

Merge remote-tracking branch 'upstream/devel' into mcmc_dev

899f06c

# Conflicts: # hmp/estimators/__init__.py # hmp/models/event.py # tests/test_estimation_interface.py

Merge branch 'GWeindel:devel' into devel

f23e651

Merge branch 'multilevel' into mcmc_dev

e751280

Navigation Menu

Sampling-based estimation with NUTS - #324

Sampling-based estimation with NUTS#324
kiante-fernandez wants to merge 55 commits into
GWeindel:develfrom
kiante-fernandez:mcmc_dev

kiante-fernandez commented Aug 5, 2026

Uh oh!

kiante-fernandez commented Aug 5, 2026

Uh oh!

GWeindel commented Aug 5, 2026

Uh oh!

JoKra1 commented Aug 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

kiante-fernandez commented Aug 5, 2026

Uh oh!

kiante-fernandez commented Aug 5, 2026

Uh oh!

GWeindel commented Aug 5, 2026

Uh oh!

JoKra1 commented Aug 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants