Add replica exchange attempts during equilibration phase

[jfennick]

Description

The equilibration protocol currently propagates replicas only; it does not perform replica exchange. As such, it is equilibrating w.r.t. an (alchemically-indexed) family of independent canonical ensembles, rather than the grand canonical ensemble. This PR makes the equilibration loop identical to the main replica exchange loop, excluding disk IO.

Failing to equilibrate w.r.t. the same ensemble can cause transient errors in the free energies of the replicas in the soft-core asymptote region, which manifests as cryptic normalization errors with pymbar, i.e. https://github.com/choderalab/yank/issues?q=W_nk

Note that in https://github.com/choderalab/yank/blob/master/Yank/reports/YANK_Health_Report_Template.ipynb, 2 out of 3 of the free energy call sites take the discard_from_start parameter, however report.generate_free_energy() uses all of the data written to the netcdf file and thus triggers the error anyway. The only way to emulate the discard_from_start parameter for this third call site is to simply not write the bad initial equilibration data to disk in the first place, i.e. to perform equilibration w.r.t. the grand canonical ensemble.

Also note that choderalab/pymbar#419 claims that the normalization error is due to a change in the numerical solver in pymbar. While that may be part of the problem, I think the fundamental issue is feeding bad data to the solver: if you comment out report.generate_free_energy() and increase discard_from_start (to about 5-10), the norm of the weights will rapidly approach 1 until pymbar successfully runs.

Alternatively, using this PR and setting number_of_equilibration_iterations to 5-10 in Yank, you can keep discard_from_start=1 and avoid the normalization error, even when default_number_of_iterations is arbitrarily small.

Todos

• Implement feature / fix bug
• Add tests
• Update documentation as needed
• Update changelogNotable points that this PR has either accomplished or will accomplish.

Status

• Ready to go

[mikemhenry]

Thanks! It looks there are some tests errors:

Traceback (most recent call last):
  File "/usr/local/miniconda/envs/test/lib/python3.8/site-packages/nose/case.py", line 197, in runTest
    self.test(*self.arg)
  File "/Users/runner/work/openmmtools/openmmtools/openmmtools/tests/test_sampling.py", line 1169, in test_equilibrate
    sampler.equilibrate(n_iterations=10, mcmc_moves=equilibration_move)
  File "/Users/runner/work/openmmtools/openmmtools/openmmtools/multistate/multistatesampler.py", line 680, in equilibrate
    self._mix_replicas()
  File "/usr/local/miniconda/envs/test/lib/python3.8/site-packages/mpiplus/mpiplus.py", line 271, in _wrapper
    return run_single_node(rank, task, *args, **kwargs)
  File "/usr/local/miniconda/envs/test/lib/python3.8/site-packages/mpiplus/mpiplus.py", line 220, in run_single_node
    result = task(*args, **kwargs)
  File "/Users/runner/work/openmmtools/openmmtools/openmmtools/multistate/sams.py", line 430, in _mix_replicas
    self._update_logZ_estimates(replicas_log_P_k)
  File "/Users/runner/work/openmmtools/openmmtools/openmmtools/multistate/sams.py", line 635, in _update_logZ_estimates
    gamma = self.gamma0 * min(pi_star, t**(-beta_factor)) # Eq. 15 of [1]
ZeroDivisionError: 0.0 cannot be raised to a negative power

They don't show up on main so I think they are coming from the changes in the PR.

[jfennick]

pytest openmmtools/tests/test_sampling.py now works on my machine.

[mikemhenry]

I will review this today!

[mikemhenry]

Sorry for the delay in reviewing this PR! We are working on improving our maintenance of this package and are working through the issue and PR backlog. Thanks again for your well done pull requests!

[jfennick]

Great! If you guys are getting back into active maintenance, then I've got a few more PRs I can submit.

[jchodera]

Great! If you guys are getting back into active maintenance, then I've got a few more PRs I can submit.

Thanks so much for your patience! We're finally in a position to be able to start maintaining openmmtools properly again at this point, so please do send them our way.