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,<2   # 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

@@ -30,6 +30,7 @@
 from openff.units.openmm import to_openmm, from_openmm
 
 from ..settings import NonEquilibriumCyclingSettings
+from ..utils.charge import validate_charge_difference
 from ..utils.data import serialize, deserialize
 from ..utils.exceptions import ProtocolSupportError
 from ..utils.misc import (
@@ -113,7 +114,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,10 +239,9 @@ 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(
+        # Get the formal change difference between the end states
+        try:
+            charge_difference = validate_charge_difference(
                 mapping,
                 forcefield_settings.nonbonded_method,
                 alchemical_settings.explicit_charge_correction,

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,13 @@ 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"] = "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

@@ -617,12 +617,12 @@ def test_double_charge_fails(
     ):
         """
         Test that attempting a mutation with a double charge change between lysine and glutamate
-        systems raises a `NotSupportedError`.
+        systems raises a `ProtocolSupportError`.
 
         This test verifies that the `NonEquilibriumCyclingProtocol` correctly raises an error when trying to
-        create a directed acyclic graph (DAG) for an invalid mutation involving a double charge change.
-        The test expects the `NotSupportedError` to be raised with a message indicating that
-        double-charge transformations are not supported.
+        execute the setup of an invalid mutation involving a double charge change. The charge
+        validation happens in the `SetupUnit`, so we only execute that unit directly instead of
+        running the full DAG (which would also run the, much more expensive, simulation units).
 
         Parameters
         ----------
@@ -633,9 +633,12 @@ def test_double_charge_fails(
         lys_to_glu_mapping : LigandAtomMapping
             Atom mapping defining the correspondence between atoms in the lysine and glutamate systems.
         """
+        from gufe.protocols.protocolunit import Context
         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
 
@@ -644,11 +647,14 @@ def test_double_charge_fails(
         dag = protocol.create(
             stateA=lys_capped_system,
             stateB=glu_capped_system,
-            name="Invalid proline mutation",
+            name="Invalid double charge mutation",
             mapping=lys_to_glu_mapping,
         )
 
-        # Expect an error when trying to create the DAG with this invalid transformation
+        # Charge validation happens in the setup unit -- run only that unit
+        setup_unit = dag.protocol_units[0]
+
+        # Expect an error when trying to execute the setup for this invalid transformation
         with pytest.raises(ProtocolSupportError):
             with tmpdir.as_cwd():
                 shared = Path("shared")
@@ -657,4 +663,7 @@ def test_double_charge_fails(
                 scratch = Path("scratch")
                 scratch.mkdir()
 
-                execute_DAG(dag, shared_basedir=shared, scratch_basedir=scratch)
+                context = Context(shared=shared, scratch=scratch)
+                setup_unit.execute(
+                    context=context, raise_error=True, **setup_unit.inputs
+                )

feflow/utils/charge.py

@@ -4,13 +4,109 @@
 or vice versa for charge-changing alchemical transformations.
 """
 
+import logging
+import warnings
+from gufe import LigandAtomMapping, SolventComponent
 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
+
+logger = logging.getLogger(__name__)
+
+
+# TODO: Re-evaluate if we want a more global utility function for this in the openfe "ecosystem"
+# Vendored from openfe protocol method in https://github.com/OpenFreeEnergy/openfe/blob/75cb2e85a46514633ecfe33353dfa5e9dc22e729/src/openfe/protocols/openmm_rfe/hybridtop_protocols.py#L373
+def validate_charge_difference(
+    mapping: LigandAtomMapping,
+    nonbonded_method: str,
+    explicit_charge_correction: bool,
+    solvent_component: SolventComponent | None,
+) -> int:
+    """
+    Validates the net charge difference between the two states.
+
+    Useful for uses in Hybrid Topology protocols where alchemical changes
+    of 2 or more charge units are not supported, and/or not using PME
+    when there is charge correction is not supported.
+
+    Parameters
+    ----------
+    mapping : LigandAtomMapping
+      Mapping object between transforming components.
+    nonbonded_method : str
+      The OpenMM nonbonded method used for the simulation.
+    explicit_charge_correction : bool
+      Whether to use an explicit charge correction.
+    solvent_component : openfe.SolventComponent | None
+      The SolventComponent of the simulation.
+
+    Returns
+    -------
+    int
+      The alchemical charge difference between the two states.
+
+    Raises
+    ------
+    ValueError
+      * If an explicit charge correction is attempted and the
+        nonbonded method is not PME.
+      * If the absolute charge difference is greater than one
+        and an explicit charge correction is attempted.
+      * If an explicit charge correction is attempted and there is no
+        solvent present.
+    UserWarning
+      * If there is any charge difference and no explicit charge
+        correction has been requested.
+    """
+    difference = mapping.get_alchemical_charge_difference()
+
+    if abs(difference) == 0:
+        return difference
+
+    if not explicit_charge_correction:
+        wmsg = (
+            f"A charge difference of {difference} is observed "
+            "between the end states. No charge correction has "
+            "been requested, please account for this in your "
+            "final results."
+        )
+        logger.warning(wmsg)
+        warnings.warn(wmsg)
+        return difference
+
+    if solvent_component is None:
+        errmsg = "Cannot use explicit charge correction without solvent"
+        raise ValueError(errmsg)
+
+    # We implicitly check earlier that we have to have pme for a solvated
+    # system, so we only need to check the nonbonded method here
+    if nonbonded_method.lower() != "pme":
+        errmsg = (
+            "Explicit charge correction when not using PME is not currently supported."
+        )
+        raise ValueError(errmsg)
+
+    if abs(difference) > 1:
+        errmsg = (
+            f"A charge difference of {difference} is observed "
+            "between the end states and an explicit charge  "
+            "correction has been requested. Unfortunately "
+            "only absolute differences of 1 are supported."
+        )
+        raise ValueError(errmsg)
+
+    ion = {-1: solvent_component.positive_ion, 1: solvent_component.negative_ion}[
+        difference
+    ]
+
+    wmsg = (
+        f"A charge difference of {difference} is observed "
+        "between the end states. This will be addressed by "
+        f"transforming a water into a {ion} ion"
+    )
+    logger.info(wmsg)
+
+    return difference