using GufeQuantity instead of deprecated OpenFF models

devtools/conda-envs/test_env.yaml

@@ -4,7 +4,7 @@ channels:
   - openeye
 dependencies:
     # Base depends
-  - gufe >=1.4, <2
+  - gufe >=1.7.0, <2
   - numpy
   - openfe >=1.1.0  # TODO: Remove once we don't depend on openfe
   - openff-units

feflow/settings/integrators.py

@@ -6,11 +6,20 @@
 for the specific integrator settings.
 """
 
+from typing import Annotated, TypeAlias
+
 from pydantic.v1 import validator
 
 from openff.units import unit
-from gufe.vendor.openff.models.types import FloatQuantity
 from gufe.settings import SettingsBaseModel
+from gufe.settings.typing import GufeQuantity, specify_quantity_units
+
+FemtosecondQuantity: TypeAlias = Annotated[
+    GufeQuantity, specify_quantity_units("femtoseconds")
+]
+TimestepQuantity: TypeAlias = Annotated[
+    GufeQuantity, specify_quantity_units("timestep")
+]
 
 
 class PeriodicNonequilibriumIntegratorSettings(SettingsBaseModel):
@@ -19,17 +28,15 @@ class PeriodicNonequilibriumIntegratorSettings(SettingsBaseModel):
     class Config:
         arbitrary_types_allowed = True
 
-    timestep: FloatQuantity["femtosecond"] = 4 * unit.femtoseconds
+    timestep: FemtosecondQuantity = 4 * unit.femtoseconds
     """Size of the simulation timestep. Default 4 fs."""
     splitting: str = "V R H O R V"
     """Operator splitting"""
     equilibrium_steps: int = 12500
     """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_frequency: FloatQuantity["timestep"] = (
-        25 * unit.timestep
-    )  # todo: IntQuantity
+    barostat_frequency: TimestepQuantity = 25 * unit.timestep
     """
     Frequency at which volume scaling changes should be attempted.
     Note: The barostat frequency is ignored for gas-phase simulations.

feflow/settings/nonequilibrium_cycling.py

@@ -14,6 +14,7 @@
 from openfe.protocols.openmm_utils.omm_settings import (
     OpenMMSolvationSettings,
     OpenMMEngineSettings,
+    ThermoSettings,
 )
 from openfe.protocols.openmm_rfe.equil_rfe_settings import AlchemicalSettings
 
@@ -65,14 +66,17 @@ class Config:
     """Settings for assigning partial charges to small molecules."""
 
     # Lambda settings
-    lambda_functions = DEFAULT_ALCHEMICAL_FUNCTIONS
+    lambda_functions: dict[str, str] = DEFAULT_ALCHEMICAL_FUNCTIONS
 
     # alchemical settings
     alchemical_settings: AlchemicalSettings = AlchemicalSettings(softcore_LJ="gapsys")
 
     # integrator settings
     integrator_settings: PeriodicNonequilibriumIntegratorSettings
 
+    # Thermodynamic settings
+    thermo_settings: ThermoSettings
+
     # platform and serialization
     engine_settings: OpenMMEngineSettings  # This defines platform
     # TODO: Need to do validation checking on these values related to IntegratorSettings eq/neq steps

feflow/tests/conftest.py

@@ -130,6 +130,7 @@ def short_settings():
 
     settings = NonEquilibriumCyclingProtocol.default_settings()
 
+    settings.engine_settings.compute_platform = "CPU"  # CPU platform for tests
     settings.thermo_settings.temperature = 300 * unit.kelvin
     settings.integrator_settings.equilibrium_steps = 250
     settings.integrator_settings.nonequilibrium_steps = 250
@@ -155,6 +156,7 @@ def short_settings_multiple_cycles():
 
     settings = NonEquilibriumCyclingProtocol.default_settings()
 
+    settings.engine_settings.compute_platform = "CPU"  # CPU platform for tests
     settings.thermo_settings.temperature = 300 * unit.kelvin
     settings.integrator_settings.equilibrium_steps = 1000
     settings.integrator_settings.nonequilibrium_steps = 1000

feflow/tests/test_nonequilibrium_cycling.py

@@ -587,6 +587,8 @@ def _assign_random_partial_charges(offmol: Molecule, seed: int = 42):
         state_b = ChemicalSystem({"ligand": small_comp_b})
 
         settings = NonEquilibriumCyclingProtocol.default_settings()
+        # Make sure we use CPU platform
+        settings.engine_settings.compute_platform = "CPU"
         protocol = NonEquilibriumCyclingProtocol(settings=settings)
 
         setup = SetupUnit(
@@ -617,8 +619,8 @@ def test_settings_round_trip():
     related to <https://github.com/OpenFreeEnergy/feflow/issues/87>.
     """
     neq_settings = NonEquilibriumCyclingProtocol.default_settings()
-    neq_json = json.dumps(neq_settings.dict(), cls=JSON_HANDLER.encoder)
-    neq_settings_2 = NonEquilibriumCyclingSettings.parse_obj(
+    neq_json = json.dumps(neq_settings.model_dump(), cls=JSON_HANDLER.encoder)
+    neq_settings_2 = NonEquilibriumCyclingSettings.model_validate(
         json.loads(neq_json, cls=JSON_HANDLER.decoder)
     )
     assert neq_settings == neq_settings_2