Support recent openfe and pymbar versions (#109)

* Support for latest openfe (>=1.1.0)

* Support for pymbar >4.0

Author: ijpulidos

devtools/conda-envs/test_env.yaml

@@ -6,11 +6,11 @@ dependencies:
     # Base depends
   - gufe >=0.9.5
   - numpy
-  - openfe >=0.15  # TODO: Remove once we don't depend on openfe
+  - openfe >=1.1.0  # TODO: Remove once we don't depend on openfe
   - openff-units
   - openmm
   - openmmforcefields >=0.14.1  # TODO: remove when upstream deps fix this
-  - pymbar <4
+  - pymbar >4.0
   - pydantic >=1.10.17
   - python
 

feflow/protocols/nonequilibrium_cycling.py

@@ -23,7 +23,11 @@
 
 # TODO: Remove/change when things get migrated to openmmtools or feflow
 from openfe.protocols.openmm_utils import system_creation
-from openfe.protocols.openmm_rfe._rfe_utils.compute import get_openmm_platform
+
+try:  # Support openfe < 1.3.0
+    from openfe.protocols.openmm_rfe._rfe_utils.compute import get_openmm_platform
+except ModuleNotFoundError:
+    from openfe.protocols.openmm_utils.omm_compute import get_openmm_platform
 
 from openff.toolkit import Molecule as OFFMolecule
 from openff.units import unit
@@ -864,7 +868,9 @@ def get_estimate(self):
 
         forward_work: npt.NDArray[float] = np.array(forward_work)
         reverse_work: npt.NDArray[float] = np.array(reverse_work)
-        free_energy, error = pymbar.bar.BAR(forward_work, reverse_work)
+        fe_bar = pymbar.bar(forward_work, reverse_work)
+        free_energy = fe_bar["Delta_f"]
+        error = fe_bar["dDelta_f"]
 
         return (
             free_energy * unit.k * self.data["temperature"] * unit.avogadro_constant
@@ -938,7 +944,9 @@ def _do_bootstrap(self, forward, reverse, n_bootstraps=1000):
             indices = np.random.choice(
                 np.arange(traj_size), size=[traj_size], replace=True
             )
-            dg, ddg = pymbar.bar.BAR(forward[indices], reverse[indices])
+            fe_bar = pymbar.bar(forward[indices], reverse[indices])
+            dg = fe_bar["Delta_f"]
+            ddg = fe_bar["dDelta_f"]
             all_dgs[i] = dg
 
         return all_dgs