New Estimation interface Implementation - #244
Conversation
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
…fit method for improved parameter estimation
- Added tests for the new estimation interface in `tests/test_estimation_interface.py`.
GWeindel
left a comment
GWeindel
left a comment
There was a problem hiding this comment.
Very impressive! The new architecture allows for future estimation methods and also fixes some previous issues and inconsistencies. I'll still need some time for testing but overall this looks very promising and the code is very clean. I let @qubixes also take a look as he might also have some comments. @jelmerborst you might also want to take a look to the new implementation.
@GWeindel Thanks for the comments, super helpful!! I 100% need to do another pass, as I rushed hacking the refactor after removing all the MCMC stuff. I will take a look in coming two weeks. |
qubixes
left a comment
qubixes
left a comment
There was a problem hiding this comment.
I have looked a little bit through the code, but it's a little bit too much to go completely into every detail. Overall, the code looks very solid and abstracting estimators and convergence criteria makes a lot of sense to me. So overall, good job and I leave it to @GWeindel to check that everything is still correct!
| EstimationResult | ||
| Results of the EM estimation | ||
| """ | ||
| if model is None: |
There was a problem hiding this comment.
If you make model a positional argument, then you don't need to test whether it is None.
|
Hi @kiante-fernandez any news on this PR? If needed I can also give it a try in the upcoming weeks. |
@GWeindel thanks for the ping! Got lost in the stack. Let me circle back to this next weekend! |
|
Perfect no rush, just for checking whether you wanted to finish the PR |
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.
|
This seems to work like a charm. |
GWeindel
left a comment
GWeindel
left a comment
There was a problem hiding this comment.
Congrats! This is really a surgical PR.
There will be some work for later PRs as the legacy in estimator declaration and results collection you kept to avoid breaking changes might make less sense with this new architecture but we can definitely merge!
| estimator : BaseEstimator, optional | ||
| Estimation method used to optimize the parameters. Defaults to | ||
| :class:`~hmp.estimators.em.EMEstimator` built from the tolerance and | ||
| iteration limits given to the model. |
There was a problem hiding this comment.
What's your opinion on this? I guess in another PR I could make the default a bit more proper to adapt to the new interface
There was a problem hiding this comment.
I think that the estimator should be a constructor argument, EventModel(..., estimator=...), defaulting to EMEstimator, with the model-level tolerance, max_iteration and min_iteration deprecated. That also fixes something the current shape cannot do which is the Eliminative and Cumulative construct EventModels internally across call sites, so there is no way to pass an estimator through them today. Worth doing before we plug on alternatives to EM estimators, otherwise it only works with a bare EventModel.
| # Iteration traces are specific to iterative estimators; a sampler has none. | ||
| self.traces = result.diagnostics.get("traces") | ||
| self.traces_group = result.diagnostics.get("traces_group") | ||
| self.time_pars_dev = result.diagnostics.get("time_pars_dev") |
There was a problem hiding this comment.
Right, I'll handle this in a new PR. My opinion is that all these should only be in the estimation_result
| _, groups, _ = self.group_constructor(pattern_data.durations) | ||
| _, eventprobs = self._estim_probs_groups( | ||
| pattern_data, channel_pars, time_pars, groups, cpus=cpus | ||
| ) |
There was a problem hiding this comment.
Is there a reason here to branch but not do the other branch? Ofc I can handle this later just wondering
Thanks! I agreed on the legacy stuff. Having those (tolerance, max_iteration and min_iteration) on the model made sense when it ran EM itself, less so now that it delegates. I only kept them because EliminativeMethod.get_event_model and CumulativeMethod both forward them for now, so removing them is a cross-cutting change. Another PR can clean that up. |
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.
|
Small test hygiene commit. The per-trial sum check was passing, but I worried it was gonna create an issue. I moved it onto all trials in float64 where the invariant holds exactly. Kept it on float64 rather than just loosening the tolerance because a future JAX likelihood will need to be validated against this one in float64, and I want that path covered. No production code touched otherwise. |
Updated with latest devel prior to merge
3 participants
This PR refactors the HMP estimation system with a cleaner architecture. The changes improve code maintainability and make it easier to add new estimation methods in the future.
Added
BaseEstimatorinterface andEMEstimatorimplementationCreated
EstimationResultclass for consistent output formatSeparated model logic from estimation algorithms
Rewrote
fit()method to use new frameworkMaintained backward compatibility - existing code still works
Preserved all current functionality and parameters
Removed old EM methods
Existing functionality verified
Backward compatibility confirmed
Example