From a3b5305ad401b7c99aa7214558c9521d8a1b0e51 Mon Sep 17 00:00:00 2001 From: Richard Gowers Date: Mon, 18 Dec 2023 16:29:34 +0000 Subject: [PATCH] Issue 335 float quantity in settings (#657) * use FloatQuantity type annotation in Settings this is different to `unit.Quantity` as it then allows the coercion of strings to correct values through a pydantic validator will eventually aid in ingesting CLI inputs to Settings as strings can be passed directly onto Settings * tests for setting fields in Settings via strings --------- Co-authored-by: Mike Henry <11765982+mikemhenry@users.noreply.github.com> --- .../openmm_rfe/equil_rfe_settings.py | 3 +- openfe/protocols/openmm_utils/omm_settings.py | 21 +++--- .../tests/protocols/test_openmm_settings.py | 68 +++++++++++++++++++ 3 files changed, 81 insertions(+), 11 deletions(-) create mode 100644 openfe/tests/protocols/test_openmm_settings.py diff --git a/openfe/protocols/openmm_rfe/equil_rfe_settings.py b/openfe/protocols/openmm_rfe/equil_rfe_settings.py index 6a9770b28..8619d04ee 100644 --- a/openfe/protocols/openmm_rfe/equil_rfe_settings.py +++ b/openfe/protocols/openmm_rfe/equil_rfe_settings.py @@ -10,6 +10,7 @@ from typing import Optional from openff.units import unit +from openff.models.types import FloatQuantity import os from gufe.settings import ( @@ -91,7 +92,7 @@ class Config: Default False. """ - explicit_charge_correction_cutoff = 0.8 * unit.nanometer + explicit_charge_correction_cutoff: FloatQuantity['nanometer'] = 0.8 * unit.nanometer """ The minimum distance from the system solutes from which an alchemical water can be chosen. Default 0.8 * unit.nanometer. diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py index 063560814..74ddc970a 100644 --- a/openfe/protocols/openmm_utils/omm_settings.py +++ b/openfe/protocols/openmm_utils/omm_settings.py @@ -11,6 +11,7 @@ from typing import Optional from openff.units import unit +from openff.models.types import FloatQuantity import os from gufe.settings import ( @@ -38,7 +39,7 @@ class Config: Method for treating nonbonded interactions, currently only PME and NoCutoff are allowed. Default PME. """ - nonbonded_cutoff = 1.0 * unit.nanometer + nonbonded_cutoff: FloatQuantity['nanometer'] = 1.0 * unit.nanometer """ Cutoff value for short range nonbonded interactions. Default 1.0 * unit.nanometer. @@ -80,7 +81,7 @@ class Config: Allowed values are; `tip3p`, `spce`, `tip4pew`, and `tip5p`. """ - solvent_padding = 1.2 * unit.nanometer + solvent_padding: FloatQuantity['nanometer'] = 1.2 * unit.nanometer """Minimum distance from any solute atoms to the solvent box edge.""" @validator('solvent_model') @@ -152,7 +153,7 @@ class Config: Default `250`. """ - online_analysis_target_error = 0.0 * unit.boltzmann_constant * unit.kelvin + online_analysis_target_error: FloatQuantity = 0.0 * unit.boltzmann_constant * unit.kelvin """ Target error for the online analysis measured in kT. Once the free energy is at or below this value, the simulation will be considered complete. A @@ -242,11 +243,11 @@ class IntegratorSettings(SettingsBaseModel): class Config: arbitrary_types_allowed = True - timestep = 4 * unit.femtosecond + timestep: FloatQuantity['femtosecond'] = 4 * unit.femtosecond """Size of the simulation timestep. Default 4 * unit.femtosecond.""" - collision_rate = 1.0 / unit.picosecond + collision_rate: FloatQuantity['1/picosecond'] = 1.0 / unit.picosecond """Collision frequency. Default 1.0 / unit.pisecond.""" - n_steps = 250 * unit.timestep + n_steps = 250 * unit.timestep # todo: IntQuantity """ Number of integration timesteps between each time the MCMC move is applied. Default 250 * unit.timestep. @@ -263,7 +264,7 @@ class Config: """ constraint_tolerance = 1e-06 """Tolerance for the constraint solver. Default 1e-6.""" - barostat_frequency = 25 * unit.timestep + barostat_frequency = 25 * unit.timestep # todo: IntQuantity """ Frequency at which volume scaling changes should be attempted. Default 25 * unit.timestep. @@ -311,7 +312,7 @@ class Config: minimization_steps = 5000 """Number of minimization steps to perform. Default 5000.""" - equilibration_length: unit.Quantity + equilibration_length: FloatQuantity['nanosecond'] """ Length of the equilibration phase in units of time. The total number of steps from this equilibration length @@ -319,7 +320,7 @@ class Config: must be a multiple of the value defined for :class:`IntegratorSettings.n_steps`. """ - production_length: unit.Quantity + production_length: FloatQuantity['nanosecond'] """ Length of the production phase in units of time. The total number of steps from this production length (i.e. @@ -341,7 +342,7 @@ class Config: Selection string for which part of the system to write coordinates for. Default 'not water'. """ - checkpoint_interval = 250 * unit.timestep + checkpoint_interval = 250 * unit.timestep # todo: Needs IntQuantity """ Frequency to write the checkpoint file. Default 250 * unit.timestep. """ diff --git a/openfe/tests/protocols/test_openmm_settings.py b/openfe/tests/protocols/test_openmm_settings.py new file mode 100644 index 000000000..f73a2907b --- /dev/null +++ b/openfe/tests/protocols/test_openmm_settings.py @@ -0,0 +1,68 @@ +# This code is part of OpenFE and is licensed under the MIT license. +# For details, see https://github.com/OpenFreeEnergy/openfe + +import pytest +from openff.units import unit + +from openfe.protocols.openmm_rfe import equil_rfe_settings +# afe settings currently have no FloatQuantity values +from openfe.protocols.openmm_utils import omm_settings + + +class TestOMMSettingsFromStrings: + # checks that we can set Settings fields via strings + def test_system_settings(self): + s = omm_settings.SystemSettings() + + s.nonbonded_cutoff = '1.1 nm' + + assert s.nonbonded_cutoff == 1.1 * unit.nanometer + + def test_solvation_settings(self): + s = omm_settings.SolvationSettings() + + s.solvent_padding = '1.1 nm' + + assert s.solvent_padding == 1.1 * unit.nanometer + + def test_alchemical_sampler_settings(self): + # todo: online_analysis_target_error is in kT, how to pass this as string? + pass + + def test_integator_settings(self): + s = omm_settings.IntegratorSettings() + + s.timestep = '3 fs' + + assert s.timestep == 3.0 * unit.femtosecond + + s.collision_rate = '1.1 / ps' + + assert s.collision_rate == 1.1 / unit.picosecond + + # todo: nsteps, barostat frequency require IntQuantity + + def test_simulation_settings(self): + s = omm_settings.SimulationSettings( + equilibration_length=2.0 * unit.nanosecond, + production_length=5.0 * unit.nanosecond, + ) + + s.equilibration_length = '2.5 ns' + s.production_length = '10 ns' + + assert s.equilibration_length == 2.5 * unit.nanosecond + assert s.production_length == 10.0 * unit.nanosecond + + # todo: checkpoint_interval IntQuantity + + +class TestEquilRFESettingsFromString: + def test_alchemical_settings(self): + s = equil_rfe_settings.AlchemicalSettings() + + s.explicit_charge_correction_cutoff = '0.85 nm' + + assert s.explicit_charge_correction_cutoff == 0.85 * unit.nanometer + +