Use MBAR bootstrap error

environment.yml

@@ -29,15 +29,15 @@ dependencies:
   - openmm >=8.0.0,!=8.1.0,<8.2.0
   - openmmtools
   - openmmforcefields
-  - perses
+  - perses>=0.10.3
   - pooch
   - py3dmol
   - plugcli
   - tqdm
   - pygraphviz
   - zstandard
   # Issue #443
-  - pymbar<4.0
+  - pymbar>4.0
   # docs
   - autodoc-pydantic<2.0
   - pydata-sphinx-theme

openfe/protocols/openmm_utils/multistate_analysis.py

@@ -10,6 +10,7 @@
 import numpy.typing as npt
 from openmmtools import multistate
 from openff.units import unit, ensure_quantity
+from pymbar import MBAR
 from pymbar.utils import ParameterError
 from openfe.analysis import plotting
 from typing import Optional, Union
@@ -68,7 +69,7 @@ class MultistateEquilFEAnalysis:
     result_units : unit.Quantity
       Units to report results in.
     forward_reverse_samples : int
-      The number of samples to use in the foward and reverse analysis
+      The number of samples to use in the forward and reverse analysis
       of the free energies. Default 10.
     """
     def __init__(self, reporter: multistate.MultiStateReporter,
@@ -198,15 +199,13 @@ def _get_free_energy(
         Parameters
         ----------
         analyzer : multistate.MultiStateSamplerAnalyzer
-          MultiStateSamplerAnalyzer to extract free eneriges from.
+          MultiStateSamplerAnalyzer to extract free energies from.
         u_ln : npt.NDArray
           A n_states x (n_sampled_states * n_iterations)
           array of energies (in kT).
         N_l : npt.NDArray
           An array containing the total number of samples drawn from each
           state.
-        unit_type : unit.Quantity
-          What units to return the free energies in.
 
         Returns
         -------
@@ -220,15 +219,12 @@ def _get_free_energy(
         * Allow folks to pass in extra options for bootstrapping etc..
         * Add standard test against analyzer.get_free_energy()
         """
-        mbar = analyzer._create_mbar(u_ln, N_l)
 
-        try:
-            # pymbar 3
-            DF_ij, dDF_ij = mbar.getFreeEnergyDifferences()
-        except AttributeError:
-            r = mbar.compute_free_energy_differences()
-            DF_ij = r['Delta_f']
-            dDF_ij = r['dDelta_f']
+        # pymbar 4
+        mbar = MBAR(u_ln, N_l, solver_protocol="robust", n_bootstraps=1000, bootstrap_solver_protocol="robust")
+        r = mbar.compute_free_energy_differences(compute_uncertainty=True, uncertainty_method="bootstrap")
+        DF_ij = r['Delta_f']
+        dDF_ij = r['dDelta_f']
 
         DG = DF_ij[0, -1] * analyzer.kT
         dDG = dDF_ij[0, -1] * analyzer.kT
@@ -352,15 +348,8 @@ def get_overlap_matrix(self) -> dict[str, npt.NDArray]:
             * ``matrix``: Estimated overlap matrix of observing a sample from
               state i in state j
         """
-        try:
-            # pymbar 3
-            overlap_matrix = self.analyzer.mbar.computeOverlap()
-            # convert matrix to np array
-            overlap_matrix['matrix'] = np.array(overlap_matrix['matrix'])
-        except AttributeError:
-            overlap_matrix = self.analyzer.mbar.compute_overlap()
+        return self.analyzer.mbar.compute_overlap()
 
-        return overlap_matrix
 
     def get_exchanges(self) -> dict[str, npt.NDArray]:
         """
@@ -373,7 +362,7 @@ def get_exchanges(self) -> dict[str, npt.NDArray]:
           A dictionary containing the following:
             * ``eigenvalues``: The sorted (descending) eigenvalues of the
               lambda state transition matrix
-            * ``matrix``: The transition matrix estimate of a replica switchin
+            * ``matrix``: The transition matrix estimate of a replica switching
               from state i to state j.
         """
         # Get replica mixing statistics

openfe/tests/protocols/test_openmmutils.py

@@ -255,7 +255,8 @@ def test_free_energies(self, analyzer):
         ret_dict = analyzer.unit_results_dict
         assert len(ret_dict.items()) == 7
         assert pytest.approx(ret_dict['unit_estimate'].m) == -47.9606
-        assert pytest.approx(ret_dict['unit_estimate_error'].m) == 0.02396789
+        # more variation when using bootstrap errors so we need a loser tolerance
+        assert pytest.approx(ret_dict['unit_estimate_error'].m, rel=1e4) == 0.0251
         # forward and reverse (since we do this ourselves)
         assert_allclose(
             ret_dict['forward_and_reverse_energies']['fractions'],
@@ -269,23 +270,25 @@ def test_free_energies(self, analyzer):
                       -48.025258, -48.006349, -47.986304, -47.972138, -47.960623]),
             rtol=1e-04,
         )
+        # results generated using pymbar3 with 1000 bootstrap iterations
         assert_allclose(
             ret_dict['forward_and_reverse_energies']['forward_dDGs'].m,
-            np.array([0.07471 , 0.052914, 0.041508, 0.036613, 0.032827, 0.030489,
-                      0.028154, 0.026529, 0.025284, 0.023968]),
-            rtol=1e-04,
+            np.array([0.077645, 0.054695, 0.044680, 0.03947, 0.034822,
+                      0.033443, 0.030793, 0.028777, 0.026683, 0.026199]),
+            rtol=1e-01,
         )
         assert_allclose(
             ret_dict['forward_and_reverse_energies']['reverse_DGs'].m,
             np.array([-47.823839, -47.833107, -47.845866, -47.858173, -47.883887,
                       -47.915963, -47.93319, -47.939125, -47.949016, -47.960623]),
             rtol=1e-04,
         )
+        # results generated using pymbar3 with 1000 bootstrap iterations
         assert_allclose(
             ret_dict['forward_and_reverse_energies']['reverse_dDGs'].m,
-            np.array([0.081209, 0.055975, 0.044693, 0.038691, 0.034603, 0.031894,
-                      0.029417, 0.027082, 0.025316, 0.023968]),
-            rtol=1e-04,
+            np.array([0.088335, 0.059483, 0.046254, 0.041504, 0.03877,
+                      0.035495, 0.031981, 0.029707, 0.027095, 0.026296]),
+            rtol=1e-01,
         )
 
     def test_plots(self, analyzer, tmpdir):