use bootstrap errors from mbar for pymbar3/4

OpenFreeEnergy · Jan 14, 2025 · abb001b · abb001b
1 parent 2eb0ec8
commit abb001b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 11 deletions.
diff --git a/openfe/protocols/openmm_utils/multistate_analysis.py b/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
@@ -220,13 +221,14 @@ 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()
+            mbar = MBAR(u_ln, N_l, nbootstraps=1000)
+            DF_ij, dDF_ij = mbar.getFreeEnergyDifferences(compute_uncertainty=True, uncertainty_method="bootstrap")
         except AttributeError:
-            r = mbar.compute_free_energy_differences()
+            # 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']
 

diff --git a/openfe/tests/protocols/test_openmmutils.py b/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):