diff --git a/gufe/components/solventcomponent.py b/gufe/components/solventcomponent.py index 7bb6b4d8..e29d7774 100644 --- a/gufe/components/solventcomponent.py +++ b/gufe/components/solventcomponent.py @@ -2,13 +2,47 @@ # For details, see https://github.com/OpenFreeEnergy/gufe from __future__ import annotations +import numpy as np from openff.units import unit -from typing import Optional, Tuple +from typing import Optional, Tuple, Literal from .component import Component _CATIONS = {'Cs', 'K', 'Li', 'Na', 'Rb'} _ANIONS = {'Cl', 'Br', 'F', 'I'} +_ALLOWED_BOX_TYPES = {'cube', 'dodecahedron', 'octahedron'} + + +def _box_vectors_are_in_reduced_form(box_vectors: unit.Quantity) -> bool: + """ + Return ``True`` if the box is in OpenMM reduced form; ``False`` otherwise. + + These conditions are shared by OpenMM and GROMACS and greatly simplify + working with triclinic boxes. Any periodic system can be represented in this + form by rotating the system and lattice reduction. + See http://docs.openmm.org/latest/userguide/theory/05_other_features.html#periodic-boundary-conditions + + Acknowledgement + --------------- + Taken from openff.interchange's openff.interchange.components._packmol + """ + if box_vectors.shape != (3, 3): + return False + + a, b, c = box_vectors.m + ax, ay, az = a + bx, by, bz = b + cx, cy, cz = c + return ( + [ay, az] == [0, 0] + and bz == 0 + and ax > 0 + and by > 0 + and cz > 0 + and ax >= 2 * np.abs(bx) + and ax >= 2 * np.abs(cx) + and by >= 2 * np.abs(cy) + ) # really wanted to make this a dataclass but then can't sort & strip ion input @@ -22,39 +56,94 @@ class SolventComponent(Component): by specific MD engine methods. """ _smiles: str - _positive_ion: Optional[str] - _negative_ion: Optional[str] + _positive_ion: str + _negative_ion: str _neutralize: bool _ion_concentration: unit.Quantity def __init__(self, *, # force kwarg usage smiles: str = 'O', + num_solvent: Optional[int] = None, positive_ion: str = 'Na+', negative_ion: str = 'Cl-', neutralize: bool = True, - ion_concentration: unit.Quantity = 0.15 * unit.molar): + ion_concentration: Optional[unit.Quantity] = 0.15 * unit.molar, + solvent_padding: Optional[unit.Quantity] = 1.2 * unit.nanometer, + solvent_density: Optional[unit.Quantity] = None, + box_shape: Optional[Literal['cube', 'dodecahedron', 'octahedron']] = 'cube', + box_vectors: Optional[unit.Quantity] = None,): """ Parameters ---------- - smiles : str, optional + smiles : str smiles of the solvent, default 'O' (water) + num_solvent : int, optional + number of solvent molecules present, if ``None``, will be determined + by either `solvent_density` if defined, or the utilized solvation + backend (Protocol specific). Cannot be defined as the same time + as both `solvent_density` and either one of `solvent_padding` or + `box_vectors`. Default `None` positive_ion, negative_ion : str the pair of ions which is used to neutralize (if neutralize=True) and bring the solvent to the required ionic concentration. Must be a positive and negative monoatomic ions, defaults "Na+", "Cl-" - neutralize : bool, optional + neutralize : bool if the net charge on the chemical state is neutralized by the ions in this solvent component. Default `True` - ion_concentration : openff-units.unit.Quantity, optional + ion_concentration : openff.units.unit.Quantity, optional ionic concentration required, default 0.15 * unit.molar this must be supplied with units, e.g. "1.5 * unit.molar" + solvent_padding : openff.units.unit.Quantity, optional + padding distance away from solute to use in constructing the + solvent box. Must be supplied with units, e.g. + "1.2 * unit.nanometer". Cannot be defined at the same time as + `box_vectors`. Cannot be defined alongside both `num_solvent` + and `solvent_density`. Default 1.2 * unit.nanometer + solvent_density : openff.units.unit.Quantity, optional + Target density of solvated systems. Must be defined with units + compatible with g / mL, e.g. "850 * unit.kilogram / unit.meter**3". + Cannot be defined alongside `num_solvent` if either `box_vector` + or `solvent_padding` are defined. Default `None` + box_shape : str, optional + Defined the shape of the solvent box being built. Can be one of + 'cube', 'dodecahedron', and 'octahedron'. Cannot be defined alongside + `box_vectors`. Default 'cube' + box_vectors : openff.units.unit.Quantity, optional + Vectors defining the solvent box. Box vectors must be provided in + `OpenMM reduced form `_ + as a unit.Quantity array of shape (3, 3). + Cannot be defined alongside `box_shape`. + Cannot be defined alongside both `num_solvent` and + `solvent_density`. Default `None` + + Raises + ------ + ValueError + * If an unknown positive or negative ion type is passed. + * If `ion_concentration` is given in units incompatible with + openff.units.unit.molar. + * If `ion_concentration` is not positive. + * If `num_solvent` is not `None` and not positive. + * If `num_solvent` and `density` are defined alongside either + `solvent_padding` or `box_vectors`. + * If `solvent_padding` is defined alongside `box_vectors`. + * If `solvent_padding` is given in units incompatible with + openff.units.unit.nanometer. + * If `solvent_padding` is not `None` and not positive. + * If `solvent_density` is defined and not compatible with units of + "g / mL". + * If `solvent_desnity` is defined and not positive. + * If an unknown box shape is passed. Examples -------- - To create a sodium chloride solution at 0.2 molar concentration:: + To create a sodium chloride solution at 0.2 molar concentration with + a cubic box with edges up to 1.2 nm from the solute:: >>> s = SolventComponent(position_ion='Na', negative_ion='Cl', - ... ion_concentration=0.2 * unit.molar) + ... ion_concentration=0.2 * unit.molar, + ... solvent_padding=1.2 * unit.nanometer, + ... box_shape='cube') """ self._smiles = smiles @@ -76,14 +165,92 @@ def __init__(self, *, # force kwarg usage raise ValueError(f"ion_concentration must be given in units of" f" concentration, got: {ion_concentration}") if ion_concentration.m < 0: - raise ValueError(f"ion_concentration must be positive, " + raise ValueError("ion_concentration must be positive, " f"got: {ion_concentration}") self._ion_concentration = ion_concentration + # Validate num_solvent + if num_solvent is not None and num_solvent < 1: + errmsg = "num_solvent must be greater than zero, got {num_solvent}" + raise ValueError(errmsg) + + if (num_solvent is not None and solvent_density is not None and + (box_vectors is not None or solvent_padding is not None)): + errmsg = ("Cannot define the number of solvent molecules " + f"({num_solvent}) alongside the solvent density " + f"({solvent_density}) and either of the box vectors " + f"({box_vectors}) or the solvent padding " + f"({solvent_padding})") + raise ValueError(errmsg) + + self._num_solvent = num_solvent + + # Validate the solvent padding + if solvent_padding is not None: + + if box_vectors is not None: + errmsg = (f"solvent_padding ({solvent_padding}) cannot be " + f"defined alongside box_vectors ({box_vectors})") + raise ValueError(errmsg) + + if (not isinstance(solvent_padding, unit.Quantity) or + not solvent_padding.is_compatible_with(unit.nanometer)): + errmsg = ("solvent_padding must be given in units of " + f"distance, got: {solvent_padding}") + raise ValueError(errmsg) + + if solvent_padding.m < 0: + errmsg = ("solvent_padding must be positive, " + f"got: {solvent_padding}") + raise ValueError(errmsg) + + self._solvent_padding = solvent_padding + + # Validate solvent density + if solvent_density is not None: + if (not isinstance(solvent_density, unit.Quantity) or + not solvent_density.is_compatible_with(unit.gram / unit.milliliter)): + errmsg = ("solvent_density must be given in units compatible " + f"with g/mL, got: {solvent_density}") + raise ValueError(errmsg) + + if solvent_density.m < 0: + errmsg = ("solvent_density cannot be negative, " + f"got: {solvent_density}") + raise ValueError(errmsg) + + self._solvent_density = solvent_density + + # Validate box shape + if box_shape is not None: + if box_shape.lower() not in _ALLOWED_BOX_TYPES: + errmsg = (f"Unknown box_shape passed, got: {box_shape}, " + f"must be one of {', '.join(_ALLOWED_BOX_TYPES)}") + raise ValueError(errmsg) + + self._box_shape = box_shape + + # Validate box vectors + if box_vectors is not None: + if not isinstance(box_vectors, unit.Quantity): + errmsg = ("box_vectors must be defined as a unit.Quantity, " + f"got: {box_vectors}") + raise ValueError(errmsg) + + if not _box_vectors_are_in_reduced_form(box_vectors): + errmsg = ("box_vectors are not in reduced form, " + f"got: {box_vectors}") + raise ValueError(errmsg) + + self._box_vectors = box_vectors + @property def name(self) -> str: - return f"{self.smiles}, {self.positive_ion}, {self.negative_ion}" + return (f"{self.smiles}, {self.positive_ion}, {self.negative_ion}, " + f"{self.num_solvent}, {self.ion_concentration}, " + f"{self.solvent_padding}, {self.solvent_density}, " + f"{self.box_shape}, {self.box_vectors}") @property def smiles(self) -> str: @@ -115,6 +282,31 @@ def total_charge(self): """Solvents don't have a formal charge defined so this returns None""" return None + @property + def num_solvent(self) -> Optional[int]: + """Number of solvent molecules, if defined""" + return self._num_solvent + + @property + def solvent_padding(self) -> Optional[unit.Quantity]: + """The solvent padding distance, if defined""" + return self._solvent_padding + + @property + def solvent_density(self) -> Optional[unit.Quantity]: + """The solvent density, if defined""" + return self._solvent_padding + + @property + def box_shape(self) -> Optional[str]: + """The solvated system box shape, if defined""" + return self._box_shape + + @property + def box_vectors(self) -> Optional[unit.Quantity]: + """The solvated system box vectors, if defined""" + return self._box_vectors + @classmethod def _from_dict(cls, d): """Deserialize from dict representation""" @@ -130,7 +322,12 @@ def _to_dict(self): return {'smiles': self.smiles, 'positive_ion': self.positive_ion, 'negative_ion': self.negative_ion, 'ion_concentration': ion_conc, - 'neutralize': self._neutralize} + 'neutralize': self._neutralize, + 'num_solvent': self._num_solvent, + 'solvent_padding': self._solvent_padding, + 'solvent_density': self._solvent_density, + 'box_shape': self._box_shape, + 'box_vectors': self._box_vectors} @classmethod def _defaults(cls): diff --git a/gufe/tests/test_chemicalsystem.py b/gufe/tests/test_chemicalsystem.py index 73c05cf1..25605367 100644 --- a/gufe/tests/test_chemicalsystem.py +++ b/gufe/tests/test_chemicalsystem.py @@ -119,7 +119,7 @@ class TestChemicalSystem(GufeTokenizableTestsMixin): cls = ChemicalSystem key = "ChemicalSystem-92176395ceb86ecd7787ce2585b24218" - repr = "ChemicalSystem(name=, components={'solvent': SolventComponent(name=O, K+, Cl-), 'ligand': SmallMoleculeComponent(name=toluene)})" + repr = "ChemicalSystem(name=, components={'solvent': SolventComponent(name=O, K+, Cl-, None, 0.0 molar, 1.2 nanometer, 1.2 nanometer, cube, None), 'ligand': SmallMoleculeComponent(name=toluene)})" @pytest.fixture def instance(self, solv_comp, toluene_ligand_comp): diff --git a/gufe/tests/test_solvents.py b/gufe/tests/test_solvents.py index e63e672c..d1bc1844 100644 --- a/gufe/tests/test_solvents.py +++ b/gufe/tests/test_solvents.py @@ -1,5 +1,5 @@ import pytest - +import numpy as np from gufe import SolventComponent from openff.units import unit @@ -29,18 +29,49 @@ def test_hash(pos, neg): assert s2.negative_ion == 'Cl-' -def test_neq(): - s1 = SolventComponent(positive_ion='Na', negative_ion='Cl') - s2 = SolventComponent(positive_ion='K', negative_ion='Cl') +@pytest.mark.parametrize('args1, args2', [ + [{'positive_ion': 'Na', 'negative_ion': 'Cl'}, + {'positive_ion': 'K', 'negative_ion': 'Cl'}], + [{'num_solvent': 2}, {'num_solvent': 4}], + [{'solvent_padding': 1.2*unit.nanometer}, + {'solvent_padding': 2.0*unit.nanometer}], + [{'solvent_density': 850*unit.kilogram/unit.meter**3}, + {'solvent_density': 750*unit.kilogram/unit.meter**3}], + [{'box_shape': 'cube'}, {'box_shape': 'dodecahedron'}], + [{'solvent_padding': None, + 'box_vectors': 20 * np.identity(3) * unit.angstrom}, + {'solvent_padding': None, + 'box_vectors': 10 * np.identity(3) * unit.angstrom}], +]) +def test_neq(args1, args2): + s1 = SolventComponent(**args1) + s2 = SolventComponent(**args2) assert s1 != s2 -@pytest.mark.parametrize('conc', [0.0 * unit.molar, 1.75 * unit.molar]) -def test_from_dict(conc): - s1 = SolventComponent(positive_ion='Na', negative_ion='Cl', - ion_concentration=conc, - neutralize=False) +@pytest.mark.parametrize('args', [ + {'positive_ion': 'Na', + 'negative_ion': 'Cl', + 'ion_concentration': 0.0 * unit.molar, + 'neutralize': False, + 'num_solvent': 20, + 'solvent_padding': 1.5 * unit.nanometer, + 'solvent_density': None, + 'box_shape': 'dodecahedron', + 'box_vectors': None}, + {'positive_ion': 'Na', + 'negative_ion': 'Cl', + 'ion_concentration': 1.75 * unit.molar, + 'neutralize': False, + 'num_solvent': None, + 'solvent_padding': None, + 'solvent_density': 850 * unit.kilogram / unit.meter**3, + 'box_shape': 'cube', + 'box_vectors': 20 * np.identity(3) * unit.nanometer}, +]) +def test_from_dict(args): + s1 = SolventComponent(**args) assert SolventComponent.from_dict(s1.to_dict()) == s1 @@ -57,7 +88,8 @@ def test_conc(): 1.5 * unit.kg, # probably a tad much salt -0.1 * unit.molar]) # negative conc def test_bad_conc(conc): - with pytest.raises(ValueError): + errmsg = "ion_concentration must be" + with pytest.raises(ValueError, match=errmsg): _ = SolventComponent(positive_ion='Na', negative_ion='Cl', ion_concentration=conc) @@ -74,15 +106,113 @@ def test_solvent_charge(): ('F', 'I'), ]) def test_bad_inputs(pos, neg): - with pytest.raises(ValueError): + errmsg = "Invalid" + with pytest.raises(ValueError, match=errmsg): _ = SolventComponent(positive_ion=pos, negative_ion=neg) +@pytest.mark.parametrize('nsolv', [0, -42]) +def test_negative_num_solvent(nsolv): + errmsg = "num_solvent must be greater than zero" + with pytest.raises(ValueError, match=errmsg): + _ = SolventComponent(num_solvent=nsolv) + + +@pytest.mark.parametrize('box_vectors,solvent_padding', [ + [None, 1.2*unit.nanometer], + [20 * np.identity(3) * unit.angstrom, None], +]) +def test_defined_solvent_with_defined_box_error(box_vectors, solvent_padding): + errmsg = "Cannot define the number of solvent molecules" + with pytest.raises(ValueError, match=errmsg): + _ = SolventComponent( + num_solvent=42, + solvent_density=850*unit.kilogram/unit.meter**3, + solvent_padding=solvent_padding, + box_vectors=box_vectors, + ) + + +def test_solvent_padding_box_vectors_error(): + errmsg = "cannot be defined alongside box_vectors" + with pytest.raises(ValueError, match=errmsg): + _ = SolventComponent( + solvent_padding=1.2*unit.nanometer, + box_vectors=20*np.identity(3)*unit.angstrom, + ) + + +def test_solvent_padding_not_distance_error(): + errmsg = "solvent_padding must be given in units of" + with pytest.raises(ValueError, match=errmsg): + _ = SolventComponent( + solvent_padding=1.2*unit.molar, + ) + + +def test_negative_solvent_padding_error(): + errmsg = "solvent_padding must be positive" + with pytest.raises(ValueError, match=errmsg): + _ = SolventComponent( + solvent_padding=-1*unit.nanometer + ) + + +def test_incompatible_density_error(): + errmsg = "solvent_density must be given in units compatible with g/mL" + with pytest.raises(ValueError, match=errmsg): + _ = SolventComponent( + solvent_density=1*unit.nanometer + ) + + +def test_negative_density_error(): + errmsg = "solvent_density cannot be negative" + with pytest.raises(ValueError, match=errmsg): + _ = SolventComponent( + solvent_density=-850*unit.kilogram/unit.meter**3, + ) + + +def test_unknown_box_type_error(): + errmsg = "Unknown box_shape passed" + with pytest.raises(ValueError, match=errmsg): + _ = SolventComponent( + box_shape='rhombic' + ) + + +def test_no_units_box_vectors(): + errmsg = "box_vectors must be defined as a unit.Quantity" + with pytest.raises(ValueError, match=errmsg): + _ = SolventComponent( + solvent_padding=None, + box_vectors=20*np.identity(3), + ) + + +@pytest.mark.parametrize('box_vectors', [ + 20 * np.identity(2) * unit.angstrom, + np.asarray([ + [0.5, 0.5, np.sqrt(2.0) / 2.0], + [0.0, 1.0, 0.0], + [1.0, 0.0, 0.0], + ]) * unit.angstrom, +]) +def test_box_vectors_not_reduced_form_error(box_vectors): + errmsg = "box_vectors are not in reduced form" + with pytest.raises(ValueError, match=errmsg): + _ = SolventComponent( + solvent_padding=None, + box_vectors=box_vectors, + ) + + class TestSolventComponent(GufeTokenizableTestsMixin): cls = SolventComponent key = "SolventComponent-26b4034ad9dbd9f908dfc298ea8d449f" - repr = "SolventComponent(name=O, Na+, Cl-)" + repr = "SolventComponent(name=O, Na+, Cl-, None, 0.15 molar, 1.2 nanometer, 1.2 nanometer, cube, None)" @pytest.fixture def instance(self): diff --git a/gufe/tests/test_transformation.py b/gufe/tests/test_transformation.py index cfb871bc..83aa08d8 100644 --- a/gufe/tests/test_transformation.py +++ b/gufe/tests/test_transformation.py @@ -32,7 +32,7 @@ class TestTransformation(GufeTokenizableTestsMixin): cls = Transformation key = "Transformation-3166a168ef6ea2a7b2f036415ca52a61" - repr = "Transformation(stateA=ChemicalSystem(name=, components={'ligand': SmallMoleculeComponent(name=toluene), 'solvent': SolventComponent(name=O, K+, Cl-)}), stateB=ChemicalSystem(name=, components={'protein': ProteinComponent(name=), 'solvent': SolventComponent(name=O, K+, Cl-), 'ligand': SmallMoleculeComponent(name=toluene)}), protocol=)" + repr = "Transformation(stateA=ChemicalSystem(name=, components={'ligand': SmallMoleculeComponent(name=toluene), 'solvent': SolventComponent(name=O, K+, Cl-, None, 0.0 molar, 1.2 nanometer, 1.2 nanometer, cube, None)}), stateB=ChemicalSystem(name=, components={'protein': ProteinComponent(name=), 'solvent': SolventComponent(name=O, K+, Cl-, None, 0.0 molar, 1.2 nanometer, 1.2 nanometer, cube, None), 'ligand': SmallMoleculeComponent(name=toluene)}), protocol=)" @pytest.fixture def instance(self, absolute_transformation): @@ -127,7 +127,7 @@ class TestNonTransformation(GufeTokenizableTestsMixin): cls = NonTransformation key = "NonTransformation-8c81ca1e263572bc3235a15a11bad376" - repr = "NonTransformation(stateA=ChemicalSystem(name=, components={'protein': ProteinComponent(name=), 'solvent': SolventComponent(name=O, K+, Cl-), 'ligand': SmallMoleculeComponent(name=toluene)}), stateB=ChemicalSystem(name=, components={'protein': ProteinComponent(name=), 'solvent': SolventComponent(name=O, K+, Cl-), 'ligand': SmallMoleculeComponent(name=toluene)}), protocol=)" + repr = "NonTransformation(stateA=ChemicalSystem(name=, components={'protein': ProteinComponent(name=), 'solvent': SolventComponent(name=O, K+, Cl-, None, 0.0 molar, 1.2 nanometer, 1.2 nanometer, cube, None), 'ligand': SmallMoleculeComponent(name=toluene)}), stateB=ChemicalSystem(name=, components={'protein': ProteinComponent(name=), 'solvent': SolventComponent(name=O, K+, Cl-, None, 0.0 molar, 1.2 nanometer, 1.2 nanometer, cube, None), 'ligand': SmallMoleculeComponent(name=toluene)}), protocol=)" @pytest.fixture def instance(self, complex_equilibrium): diff --git a/gufe/tokenization.py b/gufe/tokenization.py index eb7de2a5..1508efe9 100644 --- a/gufe/tokenization.py +++ b/gufe/tokenization.py @@ -14,6 +14,8 @@ import warnings from typing import Any, Dict, List, Tuple, Union +import numpy as np + from gufe.custom_codecs import ( BYTES_CODEC, DATETIME_CODEC, @@ -535,7 +537,7 @@ def to_keyed_dict(self, include_defaults=True) -> Dict: if not include_defaults: for key, value in self.defaults().items(): - if dct.get(key) == value: + if np.all(dct.get(key) == value): dct.pop(key) return dct