Merge pull request #695 from OpenFreeEnergy/solvation-prep

Solvation settings prep

docs/reference/api/openmm_rfe.rst

@@ -112,14 +112,3 @@ Below are the settings which can be tweaked in the protocol. The default setting
    :field-list-validators: False
    :inherited-members: SettingsBaseModel
    :member-order: bysource
-
-.. autopydantic_model:: SolvationSettings
-   :model-show-json: False
-   :model-show-field-summary: False
-   :model-show-config-member: False
-   :model-show-config-summary: False
-   :model-show-validator-members: False
-   :model-show-validator-summary: False
-   :field-list-validators: False
-   :inherited-members: SettingsBaseModel
-   :member-order: bysource

docs/reference/api/openmm_solvation_afe.rst

@@ -22,18 +22,6 @@ Protocol Settings
 
 Below are the settings which can be tweaked in the protocol. The default settings (accessed using :meth:`AbsoluteSolvationProtocol.default_settings`) will automatically populate settings which we have found to be useful for running solvation free energy calculations. There will however be some cases (such as when calculating difficult to converge systems) where you will need to tweak some of the following settings.
 
-.. autopydantic_model:: AbsoluteSolvationSettings
-   :model-show-json: False
-   :model-show-field-summary: False
-   :model-show-config-member: False
-   :model-show-config-summary: False
-   :model-show-validator-members: False
-   :model-show-validator-summary: False
-   :field-list-validators: False
-   :inherited-members: SettingsBaseModel
-   :exclude-members: get_defaults
-   :member-order: bysource
-   :noindex:
 
 .. module:: openfe.protocols.openmm_afe.equil_afe_settings
 
@@ -120,15 +108,3 @@ Below are the settings which can be tweaked in the protocol. The default setting
    :inherited-members: SettingsBaseModel
    :member-order: bysource
    :noindex:
-
-.. autopydantic_model:: SolvationSettings
-   :model-show-json: False
-   :model-show-field-summary: False
-   :model-show-config-member: False
-   :model-show-config-summary: False
-   :model-show-validator-members: False
-   :model-show-validator-summary: False
-   :field-list-validators: False
-   :inherited-members: SettingsBaseModel
-   :member-order: bysource
-   :noindex:

openfe/protocols/openmm_afe/base.py

@@ -51,7 +51,7 @@
     SettingsBaseModel,
 )
 from openfe.protocols.openmm_afe.equil_afe_settings import (
-    SolvationSettings,
+    BaseSolvationSettings,
     MultiStateSimulationSettings, OpenMMEngineSettings,
     IntegratorSettings, LambdaSettings, OutputSettings,
     ThermoSettings,
@@ -231,7 +231,7 @@ def _handle_settings(self):
         Get a dictionary with the following entries:
           * forcefield_settings : OpenMMSystemGeneratorFFSettings
           * thermo_settings : ThermoSettings
-          * solvation_settings : SolvationSettings
+          * solvation_settings : BaseSolvationSettings
           * alchemical_settings : AlchemicalSettings
           * lambda_settings : LambdaSettings
           * engine_settings : OpenMMEngineSettings
@@ -289,7 +289,7 @@ def _get_modeller(
         solvent_component: Optional[SolventComponent],
         smc_components: dict[SmallMoleculeComponent, OFFMolecule],
         system_generator: SystemGenerator,
-        solvation_settings: SolvationSettings
+        solvation_settings: BaseSolvationSettings
     ) -> tuple[app.Modeller, dict[Component, npt.NDArray]]:
         """
         Get an OpenMM Modeller object and a list of residue indices
@@ -305,7 +305,7 @@ def _get_modeller(
           List of openff Molecules to add.
         system_generator : openmmforcefields.generator.SystemGenerator
           System Generator to parameterise this unit.
-        solvation_settings : SolvationSettings
+        solvation_settings : BaseSolvationSettings
           Settings detailing how to solvate the system.
 
         Returns

openfe/protocols/openmm_afe/equil_afe_settings.py

@@ -22,7 +22,8 @@
 )
 from openfe.protocols.openmm_utils.omm_settings import (
     MultiStateSimulationSettings,
-    SolvationSettings,
+    BaseSolvationSettings,
+    OpenMMSolvationSettings,
     OpenMMEngineSettings,
     IntegratorSettings,
     OutputSettings,
@@ -137,7 +138,7 @@ def must_be_positive(cls, v):
     thermo_settings: ThermoSettings
     """Settings for thermodynamic parameters"""
 
-    solvation_settings: SolvationSettings
+    solvation_settings: OpenMMSolvationSettings
     """Settings for solvating the system."""
 
     # Alchemical settings

openfe/protocols/openmm_afe/equil_solvation_afe_method.py

@@ -50,7 +50,7 @@
 )
 from openfe.protocols.openmm_afe.equil_afe_settings import (
     AbsoluteSolvationSettings,
-    SolvationSettings, AlchemicalSettings, LambdaSettings,
+    OpenMMSolvationSettings, AlchemicalSettings, LambdaSettings,
     MultiStateSimulationSettings, OpenMMEngineSettings,
     IntegratorSettings, OutputSettings,
     SettingsBaseModel,
@@ -415,7 +415,7 @@ def _default_settings(cls):
                     0.0, 0.0, 0.0, 0.0, 0.0, 0.12, 0.24,
                     0.36, 0.48, 0.6, 0.7, 0.77, 0.85, 1.0],
             ),
-            solvation_settings=SolvationSettings(),
+            solvation_settings=OpenMMSolvationSettings(),
             vacuum_engine_settings=OpenMMEngineSettings(),
             solvent_engine_settings=OpenMMEngineSettings(),
             integrator_settings=IntegratorSettings(),
@@ -739,7 +739,7 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
           A dictionary with the following entries:
             * forcefield_settings : OpenMMSystemGeneratorFFSettings
             * thermo_settings : ThermoSettings
-            * solvation_settings : SolvationSettings
+            * solvation_settings : OpenMMSolvationSettings
             * alchemical_settings : AlchemicalSettings
             * lambda_settings : LambdaSettings
             * engine_settings : OpenMMEngineSettings
@@ -823,7 +823,7 @@ def _handle_settings(self) -> dict[str, SettingsBaseModel]:
           A dictionary with the following entries:
             * forcefield_settings : OpenMMSystemGeneratorFFSettings
             * thermo_settings : ThermoSettings
-            * solvation_settings : SolvationSettings
+            * solvation_settings : OpenMMSolvationSettings
             * alchemical_settings : AlchemicalSettings
             * lambda_settings : LambdaSettings
             * engine_settings : OpenMMEngineSettings

openfe/protocols/openmm_md/plain_md_methods.py

@@ -34,7 +34,7 @@
 )
 from openfe.protocols.openmm_md.plain_md_settings import (
     PlainMDProtocolSettings,
-    SolvationSettings, OpenMMEngineSettings,
+    OpenMMSolvationSettings, OpenMMEngineSettings,
     IntegratorSettings, MDSimulationSettings, MDOutputSettings,
 )
 from openff.toolkit.topology import Molecule as OFFMolecule
@@ -121,7 +121,7 @@ def _default_settings(cls):
                 temperature=298.15 * unit.kelvin,
                 pressure=1 * unit.bar,
             ),
-            solvation_settings=SolvationSettings(),
+            solvation_settings=OpenMMSolvationSettings(),
             engine_settings=OpenMMEngineSettings(),
             integrator_settings=IntegratorSettings(),
             simulation_settings=MDSimulationSettings(
@@ -461,12 +461,9 @@ def run(self, *, dry=False, verbose=True,
             'settings']
         stateA = self._inputs['stateA']
 
-        forcefield_settings: settings.OpenMMSystemGeneratorFFSettings = \
-            protocol_settings.forcefield_settings
-        thermo_settings: settings.ThermoSettings = \
-            protocol_settings.thermo_settings
-        solvation_settings: SolvationSettings = \
-            protocol_settings.solvation_settings
+        forcefield_settings: settings.OpenMMSystemGeneratorFFSettings = protocol_settings.forcefield_settings
+        thermo_settings: settings.ThermoSettings = protocol_settings.thermo_settings
+        solvation_settings: OpenMMSolvationSettings = protocol_settings.solvation_settings
         sim_settings: MDSimulationSettings = protocol_settings.simulation_settings
         output_settings: MDOutputSettings = protocol_settings.output_settings
         timestep = protocol_settings.integrator_settings.timestep

openfe/protocols/openmm_md/plain_md_settings.py

@@ -9,7 +9,9 @@
 """
 from openfe.protocols.openmm_utils.omm_settings import (
     Settings,
-    SolvationSettings, OpenMMEngineSettings, MDSimulationSettings,
+    OpenMMSolvationSettings,
+    OpenMMEngineSettings,
+    MDSimulationSettings,
     IntegratorSettings, MDOutputSettings,
 )
 from gufe.settings import SettingsBaseModel
@@ -36,7 +38,7 @@ def must_be_positive(cls, v):
         return v
 
     # Things for creating the systems
-    solvation_settings: SolvationSettings
+    solvation_settings: OpenMMSolvationSettings
 
     # MD Engine things
     engine_settings: OpenMMEngineSettings

openfe/protocols/openmm_rfe/equil_rfe_methods.py

@@ -53,7 +53,7 @@
 
 from .equil_rfe_settings import (
     RelativeHybridTopologyProtocolSettings,
-    SolvationSettings, AlchemicalSettings, LambdaSettings,
+    OpenMMSolvationSettings, AlchemicalSettings, LambdaSettings,
     MultiStateSimulationSettings, OpenMMEngineSettings,
     IntegratorSettings, OutputSettings,
 )
@@ -454,7 +454,7 @@ def _default_settings(cls):
                 temperature=298.15 * unit.kelvin,
                 pressure=1 * unit.bar,
             ),
-            solvation_settings=SolvationSettings(),
+            solvation_settings=OpenMMSolvationSettings(),
             alchemical_settings=AlchemicalSettings(softcore_LJ='gapsys'),
             lambda_settings=LambdaSettings(),
             simulation_settings=MultiStateSimulationSettings(
@@ -628,7 +628,7 @@ def run(self, *, dry=False, verbose=True,
         thermo_settings: settings.ThermoSettings = protocol_settings.thermo_settings
         alchem_settings: AlchemicalSettings = protocol_settings.alchemical_settings
         lambda_settings: LambdaSettings = protocol_settings.lambda_settings
-        solvation_settings: SolvationSettings = protocol_settings.solvation_settings
+        solvation_settings: OpenMMSolvationSettings = protocol_settings.solvation_settings
         sampler_settings: MultiStateSimulationSettings = protocol_settings.simulation_settings
         output_settings: OutputSettings = protocol_settings.output_settings
         integrator_settings: IntegratorSettings = protocol_settings.integrator_settings

openfe/protocols/openmm_rfe/equil_rfe_settings.py

@@ -8,7 +8,7 @@
 """
 from __future__ import annotations
 
-from typing import Optional, Literal
+from typing import Literal
 from openff.units import unit
 from openff.models.types import FloatQuantity
 
@@ -19,8 +19,10 @@
     ThermoSettings,
 )
 from openfe.protocols.openmm_utils.omm_settings import (
-    SolvationSettings, MultiStateSimulationSettings,
-    OpenMMEngineSettings, IntegratorSettings,
+    IntegratorSettings,
+    MultiStateSimulationSettings,
+    OpenMMEngineSettings,
+    OpenMMSolvationSettings,
     OutputSettings,
 )
 
@@ -128,7 +130,7 @@ def must_be_positive(cls, v):
     """Settings for thermodynamic parameters."""
 
     # Things for creating the systems
-    solvation_settings: SolvationSettings
+    solvation_settings: OpenMMSolvationSettings
     """Settings for solvating the system."""
 
     # Alchemical settings

openfe/protocols/openmm_utils/omm_settings.py

@@ -9,7 +9,7 @@
 """
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, Literal
 from openff.units import unit
 from openff.models.types import FloatQuantity
 
@@ -27,18 +27,23 @@
     from pydantic import validator  # type: ignore[assignment]
 
 
-class SolvationSettings(SettingsBaseModel):
-    """Settings for solvating the system
+class BaseSolvationSettings(SettingsBaseModel):
+    """
+    Base class for SolvationSettings objects
+    """
+    class Config:
+        arbitrary_types_allowed = True
+
+
+class OpenMMSolvationSettings(BaseSolvationSettings):
+    """Settings for controlling how a system is solvated using OpenMM tooling
 
     Note
     ----
     No solvation will happen if a SolventComponent is not passed.
 
     """
-    class Config:
-        arbitrary_types_allowed = True
-
-    solvent_model = 'tip3p'
+    solvent_model: Literal['tip3p', 'spce', 'tip4pew', 'tip5p'] = 'tip3p'
     """
     Force field water model to use.
     Allowed values are; `tip3p`, `spce`, `tip4pew`, and `tip5p`.
@@ -47,16 +52,6 @@ class Config:
     solvent_padding: FloatQuantity['nanometer'] = 1.2 * unit.nanometer
     """Minimum distance from any solute atoms to the solvent box edge."""
 
-    @validator('solvent_model')
-    def allowed_solvent(cls, v):
-        allowed_models = ['tip3p', 'spce', 'tip4pew', 'tip5p']
-        if v.lower() not in allowed_models:
-            errmsg = (
-                f"Only {allowed_models} are allowed solvent_model values"
-            )
-            raise ValueError(errmsg)
-        return v
-
     @validator('solvent_padding')
     def is_positive_distance(cls, v):
         # these are time units, not simulation steps

openfe/protocols/openmm_utils/system_creation.py

@@ -17,8 +17,9 @@
 from gufe import (
     Component, ProteinComponent, SolventComponent, SmallMoleculeComponent
 )
-from ..openmm_rfe.equil_rfe_settings import (
-    SolvationSettings, IntegratorSettings,
+from openfe.protocols.openmm_utils.omm_settings import (
+    IntegratorSettings,
+    OpenMMSolvationSettings,
 )
 
 
@@ -132,11 +133,13 @@ def get_system_generator(
 ModellerReturn = tuple[app.Modeller, dict[Component, npt.NDArray]]
 
 
-def get_omm_modeller(protein_comp: Optional[ProteinComponent],
-                     solvent_comp: Optional[SolventComponent],
-                     small_mols: dict[SmallMoleculeComponent, OFFMol],
-                     omm_forcefield : app.ForceField,
-                     solvent_settings : SolvationSettings) -> ModellerReturn:
+def get_omm_modeller(
+    protein_comp: Optional[ProteinComponent],
+    solvent_comp: Optional[SolventComponent],
+    small_mols: dict[SmallMoleculeComponent, OFFMol],
+    omm_forcefield : app.ForceField,
+    solvent_settings : OpenMMSolvationSettings
+) -> ModellerReturn:
     """
     Generate an OpenMM Modeller class based on a potential input ProteinComponent,
     SolventComponent, and a set of small molecules.

openfe/tests/protocols/test_openmm_settings.py

@@ -19,7 +19,7 @@ def test_system_settings(self):
         assert s.nonbonded_cutoff == 1.1 * unit.nanometer
 
     def test_solvation_settings(self):
-        s = omm_settings.SolvationSettings()
+        s = omm_settings.OpenMMSolvationSettings()
 
         s.solvent_padding = '1.1 nm'
 

openfe/tests/protocols/test_openmmutils.py

@@ -18,7 +18,8 @@
     multistate_analysis, omm_settings
 )
 from openfe.protocols.openmm_rfe.equil_rfe_settings import (
-    SolvationSettings, IntegratorSettings,
+    IntegratorSettings,
+    OpenMMSolvationSettings,
 )
 
 
@@ -335,7 +336,7 @@ def test_get_omm_modeller_complex(self, T4_protein_component,
                 T4_protein_component, openfe.SolventComponent(),
                 {smc: mol},
                 generator.forcefield,
-                SolvationSettings())
+                OpenMMSolvationSettings())
 
         resids = [r for r in model.topology.residues()]
         assert resids[163].name == 'NME'
@@ -363,7 +364,7 @@ def test_get_omm_modeller_ligand_no_neutralize(self, get_settings):
             openfe.SolventComponent(neutralize=False),
             {smc: offmol},
             generator.forcefield,
-            SolvationSettings(),
+            OpenMMSolvationSettings(),
         )
 
         system = generator.create_system(

openfe/tests/protocols/test_rfe_tokenization.py

@@ -39,7 +39,7 @@ def instance(self):
 
 class TestRelativeHybridTopologyProtocol(GufeTokenizableTestsMixin):
     cls = openmm_rfe.RelativeHybridTopologyProtocol
-    key = "RelativeHybridTopologyProtocol-a7946099ab9756c22f3a79fe9082585a"
+    key = "RelativeHybridTopologyProtocol-b6ecbf56c5effbda11d3c10d13f37273"
     repr = f"<{key}>"
 
     @pytest.fixture()

openfe/tests/protocols/test_solvation_afe_tokenization.py

@@ -49,7 +49,7 @@ def protocol_result(afe_solv_transformation_json):
 
 class TestAbsoluteSolvationProtocol(GufeTokenizableTestsMixin):
     cls = openmm_afe.AbsoluteSolvationProtocol
-    key = "AbsoluteSolvationProtocol-9a18332b06a721da1b0fcaf5cc86dd25"
+    key = "AbsoluteSolvationProtocol-0d2c5a0a9c18ec7d2629ce21b983f3fe"
     repr = f"<{key}>"
 
     @pytest.fixture()