OpenMMSolvation: Support for non-cubic boxes and defined box properties

environment.yml

@@ -27,7 +27,7 @@ dependencies:
   - openfe-analysis>=0.2.0
   - click
   - typing-extensions
-  - openmm !=8.1.0
+  - openmm >=8.0.0,!=8.1.0
   - openmmtools
   - openmmforcefields
   - perses

openfe/protocols/openmm_afe/equil_solvation_afe_method.py

@@ -683,6 +683,11 @@ def _create(
                       "passed")
             raise ValueError(errmsg)
 
+        # Validate solvation settings
+        settings_validation.validate_openmm_solvation_settings(
+            self.settings.solvation_settings
+        )
+
         # Check vacuum equilibration MD settings is 0 ns
         nvt_time = self.settings.vacuum_equil_simulation_settings.equilibration_length_nvt
         if nvt_time is not None:

openfe/protocols/openmm_md/plain_md_methods.py

@@ -171,6 +171,11 @@ def _create(
         # Validate protein component
         system_validation.validate_protein(stateA)
 
+        # Validate solvation settings
+        settings_validation.validate_openmm_solvation_settings(
+            self.settings.solvation_settings
+        )
+
         # actually create and return Units
         # TODO: Deal with multiple ProteinComponents
         solvent_comp, protein_comp, small_mols = system_validation.get_components(stateA)

openfe/protocols/openmm_rfe/equil_rfe_methods.py

@@ -528,6 +528,11 @@ def _create(
         nonbond = self.settings.forcefield_settings.nonbonded_method
         system_validation.validate_solvent(stateA, nonbond)
 
+        # Validate solvation settings
+        settings_validation.validate_openmm_solvation_settings(
+            self.settings.solvation_settings
+        )
+
         # Validate protein component
         system_validation.validate_protein(stateA)
 

openfe/protocols/openmm_utils/omm_settings.py

@@ -11,7 +11,8 @@
 
 from typing import Optional, Literal
 from openff.units import unit
-from openff.models.types import FloatQuantity
+from openff.models.types import FloatQuantity, ArrayQuantity
+from openff.interchange.components._packmol import _box_vectors_are_in_reduced_form
 
 from gufe.settings import (
     Settings,
@@ -26,38 +27,183 @@
 
 class BaseSolvationSettings(SettingsBaseModel):
     """
-    Base class for SolvationSettings objects
+    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
+    """Settings for controlling how a system is solvated using OpenMM tooling.
+
+    Defining the number of waters
+    -----------------------------
+
+    The number of waters is controlled by either:
+      a) defining a solvent padding (``solvent_padding``) in combination
+         with a box shape
+      b) defining the number of solvent molecules
+         (``number_of_solvent_molecules``)
+         alongside the box shape (``box_shape``)
+      c) defining the box directly either through the box vectors
+         (``box_vectors``) or rectangular box lengths (``box_size``)
+
+    When using ``solvent_padding``, ``box_vectors``, or ``box_size``,
+    the exact number of waters added is determined automatically by OpenMM
+    through :meth:`openmm.app.Modeller.addSolvent` internal heuristics.
+    Briefly, the necessary volume required by a single water is estimated
+    and then the defined target cell is packed with waters avoiding clashes
+    with existing solutes and box edges.
+
+
+    Defining the periodic cell size
+    -------------------------------
+
+    The periodic cell size is defined by one, and only one, of the following:
+      * ``solvent_padding`` in combination with ``box_shape``,
+      * ``number_of_solvent_molecules`` in combination with ``box_shape``,
+      * ``box_vectors``,
+      * ``box_size``
+
+    When using ``number_of_solvent_molecules``, the size of the cell is
+    defined by :meth:`openmm.app.Modeller.addSolvent` internal heuristics,
+    automatically selecting a padding value that is large enough to contain
+    the number of waters based on a geometric estimate of the volume required
+    by each water molecule.
+
+
+    Defining the periodic cell shape
+    ---------------------------------
+
+    The periodic cell shape is defined by one, and only one, of the following:
+      * ``box_shape``,
+      * ``box_vectors``,
+      * ``box_size``
+
+    Default settings will create a cubic box, although more space efficient
+    shapes (e.g. ``dodecahedrons``) are recommended to improve simulation
+    performance.
+
+
+    Notes
+    -----
+    * The number of water molecules added will be affected by the number of
+      ions defined in SolventComponent. For example, the value of
+      ``number_of_solvent_molecules`` is the sum of the number of counterions
+      added and the number of water molecules added.
+    * Solvent addition does not account for any pre-existing waters explicitly
+      defined in the :class:`openfe.ChemicalSystem`. Any waters will be added
+      in addition to those pre-existing waters.
+    * No solvation will happen if a SolventComponent is not passed.
+
+
+    See Also
+    --------
+    :mod:`openmm.app.Modeller`
+    Base class for SolvationSettings objects
+    """
+    solvent_model: Literal['tip3p', 'spce', 'tip4pew', 'tip5p'] = 'tip3p'
+    """
+    Force field water model to use when solvating and defining the model
+    properties (e.g. adding virtual site particles).
+
+    Allowed values are; `tip3p`, `spce`, `tip4pew`, and `tip5p`.
+    """
+    solvent_padding: Optional[FloatQuantity['nanometer']] = 1.2 * unit.nanometer
+    """
+    Minimum distance from any solute bounding sphere to the edge of the box.
 
     Note
     ----
-    No solvation will happen if a SolventComponent is not passed.
+    * Cannot be defined alongside ``number_of_solvent_molecules``,
+      ``box_size``, or ``box_vectors``.
+    """
+    box_shape: Optional[Literal['cube', 'dodecahedron', 'octahedron']] = 'cube'
+    """
+    The shape of the periodic box to create.
 
+    Notes
+    -----
+    * Must be one of `cube`, `dodecahedron`, or `octahedron`.
+    * Cannot be defined alongside ``box_vectors`` or ``box_size``.
     """
-    solvent_model: Literal['tip3p', 'spce', 'tip4pew', 'tip5p'] = 'tip3p'
+    number_of_solvent_molecules: Optional[int] = None
     """
-    Force field water model to use.
-    Allowed values are; `tip3p`, `spce`, `tip4pew`, and `tip5p`.
+    The number of solvent molecules (water + ions) to add.
+
+    Note
+    ----
+    * Cannot be defined alongside ``solvent_padding``, ``box_size``,
+      or ``box_vectors``.
+    """
+    box_vectors: Optional[ArrayQuantity['nanometer']] = None
     """
+    `OpenMM reduced form box vectors <http://docs.openmm.org/latest/userguide/theory/05_other_features.html#periodic-boundary-conditions>`.
 
-    solvent_padding: FloatQuantity['nanometer'] = 1.2 * unit.nanometer
-    """Minimum distance from any solute atoms to the solvent box edge."""
+    Notes
+    -----
+    * Cannot be defined alongside ``solvent_padding``,
+      ``number_of_solvent_molecules``, or ``box_size``.
+
+    See Also
+    --------
+    :mod:`openff.interchange.components.interchange`
+    :mod:`openff.interchange.components._packmol`
+    """
+    box_size: Optional[ArrayQuantity['nanometer']] = None
+    """
+    X, Y, and Z lengths of the unit cell for a rectangular box.
+
+    Notes
+    -----
+    * Cannot be defined alongside ``solvent_padding``,
+      ``number_of_solvent_molecules``, or ``box_vectors``.
+    """
+
+    @validator('box_vectors')
+    def supported_vectors(cls, v):
+        if v is not None:
+            if not _box_vectors_are_in_reduced_form(v):
+                errmsg = f"box_vectors: {v} are not in OpenMM reduced form"
+                raise ValueError(errmsg)
+        return v
 
     @validator('solvent_padding')
     def is_positive_distance(cls, v):
         # these are time units, not simulation steps
+        if v is None:
+            return v
+
         if not v.is_compatible_with(unit.nanometer):
             raise ValueError("solvent_padding must be in distance units "
                              "(i.e. nanometers)")
         if v < 0:
             errmsg = "solvent_padding must be a positive value"
             raise ValueError(errmsg)
+
+        return v
+
+    @validator('number_of_solvent_molecules')
+    def positive_solvent_number(cls, v):
+        if v is None:
+            return v
+
+        if v <= 0:
+            errmsg = f"number_of_solvent molecules: {v} must be positive"
+            raise ValueError(errmsg)
+
+        return v
+
+    @validator('box_size')
+    def box_size_properties(cls, v):
+        if v is None:
+            return v
+
+        if v.shape != (3, ):
+            errmsg = (f"box_size must be a 1-D array of length 3 "
+                      f"got {v} with shape {v.shape}")
+            raise ValueError(errmsg)
+
         return v
 
 

openfe/protocols/openmm_utils/settings_validation.py

@@ -10,6 +10,38 @@
     IntegratorSettings,
     MultiStateSimulationSettings,
 )
+from openfe.protocols.openmm_utils.omm_settings import OpenMMSolvationSettings
+
+
+def validate_openmm_solvation_settings(
+    settings: OpenMMSolvationSettings
+) -> None:
+    """
+    Checks that the OpenMMSolvation settings are correct.
+
+    Raises
+    ------
+    ValueError
+      If more than one of ``solvent_padding``, ``number_of_solvent_molecules``,
+      ``box_vectors``, or ``box_size`` are defined.
+      If ``box_shape`` is defined alongside either ``box_vectors``,
+      or ``box_size``.
+    """
+    unique_attributes = (
+        settings.solvent_padding, settings.number_of_solvent_molecules,
+        settings.box_vectors, settings.box_size,
+    )
+    if len([x for x in unique_attributes if x is not None]) > 1:
+        errmsg = ("Only one of solvent_padding, number_of_solvent_molecules, "
+                  "box_vectors, and box_size can be defined in the solvation "
+                  "settings.")
+        raise ValueError(errmsg)
+
+    if settings.box_shape is not None:
+        if settings.box_size is not None or settings.box_vectors is not None:
+            errmsg = ("box_shape cannot be defined alongside either box_size "
+                      "or box_vectors in the solvation settings.")
+            raise ValueError(errmsg)
 
 
 def validate_timestep(hmass: float, timestep: unit.Quantity):

openfe/protocols/openmm_utils/system_creation.py

@@ -210,17 +210,32 @@ def _add_small_mol(comp,
 
     # Add solvent if neeeded
     if solvent_comp is not None:
-        conc = solvent_comp.ion_concentration
-        pos = solvent_comp.positive_ion
-        neg = solvent_comp.negative_ion
+        # Do unit conversions if necessary
+        solvent_padding = None
+        box_size = None
+        box_vectors = None
+
+        if solvent_settings.solvent_padding is not None:
+            solvent_padding = to_openmm(solvent_settings.solvent_padding)
+
+        if solvent_settings.box_size is not None:
+            box_size = to_openmm(solvent_settings.box_size)
+
+        if solvent_settings.box_vectors is not None:
+            box_vectors = to_openmm(solvent_settings.box_vectors)
 
         system_modeller.addSolvent(
             omm_forcefield,
             model=solvent_settings.solvent_model,
-            padding=to_openmm(solvent_settings.solvent_padding),
-            positiveIon=pos, negativeIon=neg,
-            ionicStrength=to_openmm(conc),
+            padding=solvent_padding,
+            positiveIon=solvent_comp.positive_ion,
+            negativeIon=solvent_comp.negative_ion,
+            ionicStrength=to_openmm(solvent_comp.ion_concentration),
             neutralize=solvent_comp.neutralize,
+            boxSize=box_size,
+            boxVectors=box_vectors,
+            boxShape=solvent_settings.box_shape,
+            numAdded=solvent_settings.number_of_solvent_molecules,
         )
 
         all_resids = np.array(

openfe/tests/protocols/test_openmm_afe_solvation_protocol.py

@@ -2,6 +2,7 @@
 # For details, see https://github.com/OpenFreeEnergy/openfe
 import itertools
 import json
+from math import sqrt
 import sys
 import pytest
 from unittest import mock
@@ -463,6 +464,57 @@ def test_dry_run_solv_benzene_tip4p(benzene_modifications, tmpdir):
         assert sol_sampler.is_periodic
 
 
+def test_dry_run_solv_benzene_noncubic(
+    benzene_modifications, tmpdir
+):
+    s = AbsoluteSolvationProtocol.default_settings()
+    s.solvation_settings.solvent_padding = 1.5 * offunit.nanometer
+    s.solvation_settings.box_shape = 'dodecahedron'
+
+    protocol = AbsoluteSolvationProtocol(settings=s)
+
+    stateA = ChemicalSystem({
+        'benzene': benzene_modifications['benzene'],
+        'solvent': SolventComponent()
+    })
+
+    stateB = ChemicalSystem({
+        'solvent': SolventComponent(),
+    })
+
+    # Create DAG from protocol, get the vacuum and solvent units
+    # and eventually dry run the first solvent unit
+    dag = protocol.create(
+        stateA=stateA,
+        stateB=stateB,
+        mapping=None,
+    )
+    prot_units = list(dag.protocol_units)
+
+    sol_unit = [u for u in prot_units
+                if isinstance(u, AbsoluteSolvationSolventUnit)]
+
+    with tmpdir.as_cwd():
+        sampler = sol_unit[0].run(dry=True)['debug']['sampler']
+        system = sampler._thermodynamic_states[0].system
+
+        vectors = system.getDefaultPeriodicBoxVectors()
+        width = float(from_openmm(vectors)[0][0].to('nanometer').m)
+
+        # dodecahedron has the following shape:
+        # [width, 0, 0], [0, width, 0], [0.5, 0.5, 0.5 * sqrt(2)] * width
+
+        expected_vectors = [
+            [width, 0, 0],
+            [0, width, 0],
+            [0.5 * width, 0.5 * width, 0.5 * sqrt(2) * width],
+        ] * offunit.nanometer
+        assert_allclose(
+            expected_vectors,
+            from_openmm(vectors)
+        )
+
+
 def test_dry_run_solv_user_charges_benzene(benzene_modifications, tmpdir):
     """
     Create a test system with fictitious user supplied charges and