[WIP] Using the robust solver for pyMBAR - avoiding convergence Failu… #735

RiesBen · 2024-07-04T15:38:08Z

Description

In this PR, we would like to propose switching the default solver, if pyMBAR > 4.0.0, such we have an improved convergence rate at the cost of minimal more time. -> less errors thrown.

Todos

Implement feature / fix bug
Add tests
Update documentation as needed
Update changelog to summarize changes in behavior, enhancements, and bugfixes implemented in this PR

Status

Ready to go

Changelog message

…res. ## Description In this PR, we would like to propose switching the default solver, if pyMBAR > 4.0.0, such we have an improved convergence rate at the cost of minimal more time. -> less errors thrown. ## Todos - [ ] Implement feature / fix bug - [ ] Add [tests](https://github.com/choderalab/openmmtools/tree/master/openmmtools/tests) - [ ] Update [documentation](https://github.com/choderalab/openmmtools/tree/master/docs) as needed - [ ] Update [changelog](https://github.com/choderalab/openmmtools/blob/master/docs/releasehistory.rst) to summarize changes in behavior, enhancements, and bugfixes implemented in this PR ## Status - [ ] Ready to go ## Changelog message ``` ```

mikemhenry · 2024-08-15T15:06:54Z

   if pymbar.version.short_version >= "4.0.0" and "solver_protocol" not in self._user_extra_analysis_kwargs:
AttributeError: module 'openmmtools.multistate.pymbar' has no attribute 'version'

mikemhenry · 2024-08-15T15:08:31Z

We should use this to compare versions https://packaging.pypa.io/en/latest/version.html#packaging.version.Version

ijpulidos · 2024-08-15T15:35:37Z

There are a few options here, from our meeting:

Need to benchmark if this affects the performance (is it slower?)
We would change the default behavior and communicate the performance impact in the CHANGELOG (release notes). Would be passed on construction of the analyzer.

Realtime analysis_kwargs in

openmmtools/openmmtools/multistate/multistatesampler.py

Lines 1540 to 1564 in cbff4c8

    
           try:  # Trap errors for MBAR being under sampled and the W_nk matrix not being normalized correctly 
        
               mbar = analysis.mbar 
        
               free_energy, err_free_energy = analysis.get_free_energy() 
        
               n_equilibration_iterations = analysis.n_equilibration_iterations 
        
               statistical_inefficiency = analysis.statistical_inefficiency 
        
           except ParameterError as e: 
        
               # We don't update self._last_err_free_energy here since if it 
        
               # wasn't below the target threshold before, it won't stop MultiStateSampler now. 
        
               logger.debug(f"ParameterError computing MBAR. {e}.") 
        
               bump_error_counter = True 
        
               self._online_error_bank.append(e) 
        
               if len(self._online_error_bank) > 6: 
        
                   # Cache only the last set 
        
                   self._online_error_bank.pop(0) 
        
               free_energy = None 
        
           else: 
        
               self._last_mbar_f_k_offline = mbar.f_k 
        
               free_energy = free_energy[idx, jdx] 
        
               self._last_err_free_energy = err_free_energy[idx, jdx] 
        
               logger.debug("Current Free Energy Estimate is {} +- {} kT".format(free_energy, 
        
                                                                                 self._last_err_free_energy)) 
        
               # Trap a case when errors don't converge (usually due to under sampling) 
        
               if np.isnan(self._last_err_free_energy): 
        
                   self._last_err_free_energy = np.inf 
        
           timer.stop("MBAR")

could stay the same as it was before.

Test with the overlap matrix from example in pymbar issue thread.
Test with "production" simulations. Just to double check how this behaves with longer trajectories and more data.
Make sure we have a CI matrix with pymbar 4 and the robust settings.

@mikemhenry @RiesBen please add any comments to this thread if I'm missing something. Thanks!

IAlibay · 2024-08-21T13:09:57Z

Realtime analysis_kwargs could stay the same as it was before.

I'm not fully clued in on the discussion, but my suggestion would be, if possible, to keep the kwargs equal between realtime analysis and the final analysis. I.e. the failures are more likely to happen at low sample counts.

Need to benchmark if this affects the performance (is it slower?)

I believe from an exchange a little while back with @mrshirts, the relative speed of different solvers is system dependent - i.e. down to the number of iterations run not the speed of each iteration.
It might be good to check what we're testing against as the performance baseline. The default solver in pymbar 3 is adaptive, whilst it's hybr with adaptive fallback in pymbar 4. Robust is adaptive w/ L-BFGS-B fallback in pymbar 4. If we're doing a performance regression test, the baseline probably should be pymbar 3 default vs pymbar 4 robust (although pymbar 4 default vs robust would be good to know too!).

…R-solver

mikemhenry · 2024-09-25T19:22:21Z

I'm not fully clued in on the discussion, but my suggestion would be, if possible, to keep the kwargs equal between realtime analysis and the final analysis. I.e. the failures are more likely to happen at low sample counts.

I think we decided against this since as the simulation goes on, the analysis takes longer so it would be better to do a fast method for the realtime analysis

IAlibay · 2024-09-25T20:15:46Z

the analysis takes longer so it would be better to do a fast method for the realtime analysis

Assuming you're trying to avoid regressing, the "slow" method is the same method as pymbar 3. So you're not actually going "slower", indeed the "slow" method isn't actually clearly slower.

Going for the "fast" method probably increases your chances of getting errored values on fractional datasets.

ijpulidos · 2024-09-26T03:31:19Z

The real time analysis currently instantiates its own MultiStateSamplerAnalyzer and it would always use the default user kwargs specified in the changes in this PR, since we don't really provide a way for the user to change those. That is, it would always be using the "robust" solver with PyMBAR 4. I think this should be fine. I agree with @IAlibay on that with less samples we have more chances to fail, so "robust" here sounds like a good idea.

I agree that we are probably jumping the gun and maybe the "robust" solver is fast enough (maybe even faster than the pymbar 3 implementation we were using). This is something that we should probably double check to keep our sanity.

Other than that, I'd suggest that we probably want to adapt the test script in choderalab/pymbar#419 (comment) to make it a tests for our MultiStateSamplerAnalyzer, just to be sure we are doing things correctly, if that makes sense.

IAlibay · 2024-09-26T16:03:12Z

Maybe to add to this a little bit, what I'm advocating for is closer to:

For now, robust will do - we shouldn't take a performance hit, because it's the same this as before.
In the future, optimizing things would be great.
Getting the "for now" solution out first would be great.

mrshirts · 2024-09-26T16:09:55Z

In the future, optimizing things would be great.

I'm happy to meet with people to discuss what "in the future" means. For now, "Robust" should be the best option, and should fail in relatively few cases. There's some interesting options of creating synthetic data to improve convergence, but that adds bias.

…default-to-robust-pyMBAR-solver

codecov · 2024-09-26T22:02:20Z

Codecov Report

Attention: Patch coverage is 84.21053% with 6 lines in your changes missing coverage. Please review.

Project coverage is 84.95%. Comparing base (c2a13c0) to head (fc0f4ad).

Additional details and impacted files

mikemhenry · 2024-09-26T22:43:08Z

ala-thr.zip

Just uploading this file here (from choderalab/pymbar#419 (comment)) so I can download the file for a test

mikemhenry · 2024-09-27T15:27:11Z

@IAlibay @ijpulidos ready for review!

IAlibay

couple of comments otherwise lgtm

IAlibay · 2024-09-28T14:59:03Z

openmmtools/tests/test_sampling.py

@@ -2612,6 +2615,42 @@ def test_resume_velocities_from_legacy_storage(self):
                    state.velocities.value_in_unit_system(unit.md_unit_system) != 0
                ), "At least some velocity in sampler state from new checkpoint is expected to different from zero."

+@pytest.fixture
+def download_nc_file(tmpdir):


Would using something like pooch be better for this kind of thing?

IAlibay · 2024-09-28T14:59:30Z

openmmtools/tests/test_sampling.py

+    reporter_file = download_nc_file
+    reporter = MultiStateReporter(reporter_file)
+    analyzer = MultiStateSamplerAnalyzer(reporter, max_n_iterations=n_iterations)
+    f_ij, df_ij = analyzer.get_free_energy()


I would encourage doing a number regression check here rather than just a pure smoke test.

RiesBen and others added 2 commits July 4, 2024 17:36

bump ci

d72fe8e

mikemhenry added 2 commits September 25, 2024 12:02

Merge branch 'main' into MultistateAnalyzer---default-to-robust-pyMBA…

d3d211b

…R-solver

doing this a different way

060d5c1

use Version to compare versions

39b683b

fix micromamba, see mamba-org/micromamba-releases#58

50b058f

mikemhenry added 2 commits September 26, 2024 14:33

fix version check

41ce632

Merge remote-tracking branch 'origin/main' into MultistateAnalyzer---…

0d733d2

…default-to-robust-pyMBAR-solver

mikemhenry force-pushed the MultistateAnalyzer---default-to-robust-pyMBAR-solver branch from fa92d11 to 0d733d2 Compare September 26, 2024 21:36

forgot how we import pymbar in this package

c1ac13e

pymbar 3 stores version differently

4916e4b

mikemhenry added 2 commits September 26, 2024 15:59

added test from pymbar issue 419

7d3dc30

re-run flaky tests

fc0f4ad

mikemhenry requested a review from ijpulidos September 27, 2024 15:26

IAlibay reviewed Sep 28, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[WIP] Using the robust solver for pyMBAR - avoiding convergence Failu… #735

[WIP] Using the robust solver for pyMBAR - avoiding convergence Failu… #735

RiesBen commented Jul 4, 2024

mikemhenry commented Aug 15, 2024

mikemhenry commented Aug 15, 2024

ijpulidos commented Aug 15, 2024

IAlibay commented Aug 21, 2024

mikemhenry commented Sep 25, 2024

IAlibay commented Sep 25, 2024

ijpulidos commented Sep 26, 2024

IAlibay commented Sep 26, 2024

mrshirts commented Sep 26, 2024 •

edited

Loading

codecov bot commented Sep 26, 2024 •

edited

Loading

mikemhenry commented Sep 26, 2024 •

edited

Loading

mikemhenry commented Sep 27, 2024

IAlibay left a comment

IAlibay Sep 28, 2024

IAlibay Sep 28, 2024

[WIP] Using the robust solver for pyMBAR - avoiding convergence Failu… #735

Are you sure you want to change the base?

[WIP] Using the robust solver for pyMBAR - avoiding convergence Failu… #735

Conversation

RiesBen commented Jul 4, 2024

Description

Todos

Status

Changelog message

mikemhenry commented Aug 15, 2024

mikemhenry commented Aug 15, 2024

ijpulidos commented Aug 15, 2024

IAlibay commented Aug 21, 2024

mikemhenry commented Sep 25, 2024

IAlibay commented Sep 25, 2024

ijpulidos commented Sep 26, 2024

IAlibay commented Sep 26, 2024

mrshirts commented Sep 26, 2024 • edited Loading

codecov bot commented Sep 26, 2024 • edited Loading

Codecov Report

mikemhenry commented Sep 26, 2024 • edited Loading

mikemhenry commented Sep 27, 2024

IAlibay left a comment

Choose a reason for hiding this comment

IAlibay Sep 28, 2024

Choose a reason for hiding this comment

IAlibay Sep 28, 2024

Choose a reason for hiding this comment

mrshirts commented Sep 26, 2024 •

edited

Loading

codecov bot commented Sep 26, 2024 •

edited

Loading

mikemhenry commented Sep 26, 2024 •

edited

Loading