Use gufe charge utils. gufe/openfe 1.10+ support

devtools/conda-envs/test_env.yaml

@@ -4,9 +4,9 @@ channels:
   - openeye
 dependencies:
     # Base depends
-  - gufe ~=1.8.0
+  - gufe >=1.9.0,<2
   - numpy
-  - openfe ~=1.8.0   # TODO: Remove once we don't depend on openfe
+  - openfe ~=1.10.0   # TODO: Remove once we don't depend on openfe
   - openff-units
   - openmm
   - openmmforcefields >=0.14.1  # TODO: remove when upstream deps fix this

feflow/protocols/nonequilibrium_cycling.py

@@ -113,7 +113,6 @@ def _execute(self, ctx, *, protocol, state_a, state_b, mapping, **inputs):
             get_alchemical_components,
         )
         from feflow.utils.hybrid_topology import HybridTopologyFactory
-        from feflow.utils.charge import get_alchemical_charge_difference
         from feflow.utils.misc import register_ff_parameters_template
 
         # Get receptor components from systems if found (None otherwise)
@@ -239,17 +238,8 @@ def _execute(self, ctx, *, protocol, state_a, state_b, mapping, **inputs):
         )
 
         # Handle charge corrections/transformations
-        # Get the change difference between the end states
-        # and check if the charge correction used is appropriate
-        try:  # Catch unsupported charges differences and raise protocol error
-            charge_difference = get_alchemical_charge_difference(
-                mapping,
-                forcefield_settings.nonbonded_method,
-                alchemical_settings.explicit_charge_correction,
-                solvent_comp_a,  # Solvent comp in a is expected to be the same as in b
-            )
-        except ValueError as e:
-            raise ProtocolSupportError(str(e))
+        # Get the formal change difference between the end states
+        charge_difference = mapping.get_alchemical_charge_difference()
 
         if alchemical_settings.explicit_charge_correction:
             alchem_water_resids = _rfe_utils.topologyhelpers.get_alchemical_waters(

feflow/settings/integrators.py

@@ -6,7 +6,7 @@
 for the specific integrator settings.
 """
 
-from typing import Annotated, TypeAlias
+from typing import Annotated, TypeAlias, Literal
 
 from pydantic import ConfigDict, field_validator
 
@@ -35,6 +35,15 @@ class PeriodicNonequilibriumIntegratorSettings(SettingsBaseModel):
     """Number of steps for the equilibrium parts of the cycle. Default 12500"""
     nonequilibrium_steps: int = 12500
     """Number of steps for the non-equilibrium parts of the cycle. Default 12500"""
+    barostat: Literal["MonteCarloBarostat", "MonteCarloMembraneBarostat"] = (
+        "MonteCarloBarostat"
+    )
+    """
+    The barostat to be used in the simulations. Default MonteCarloBarostat.
+    Notes
+    -----
+    If the system contains a membrane, use the `MonteCarloMembraneBarostat`.
+    """
     barostat_frequency: TimestepQuantity = 25 * unit.timestep
     """
     Frequency at which volume scaling changes should be attempted.

feflow/tests/test_protein_mutation.py

@@ -636,6 +636,8 @@ def test_double_charge_fails(
         from feflow.utils.exceptions import ProtocolSupportError
 
         settings = NonEquilibriumCyclingProtocol.default_settings()
+        # Change engine platform for tests
+        settings.engine_settings.compute_platform = "CPU"
         # We need to make sure we enable the alchemical charge correction
         settings.alchemical_settings.explicit_charge_correction = True
 

feflow/utils/charge.py

@@ -7,10 +7,6 @@
 from openfe.protocols.openmm_utils.charge_generation import (
     assign_offmol_partial_charges,
 )
-from openfe.protocols.openmm_rfe.equil_rfe_methods import (
-    _get_alchemical_charge_difference,
-)
 
 # TODO: Importing from OpenFE for now, should we migrate them here?
 assign_offmol_partial_charges = assign_offmol_partial_charges
-get_alchemical_charge_difference = _get_alchemical_charge_difference