Merge pull request #18 from OpenFreeEnergy/solvent_work

Tightening up `SolventComponent`

gufe/solventcomponent.py

@@ -1,86 +1,128 @@
 # This code is part of OpenFE and is licensed under the MIT license.
 # For details, see https://github.com/OpenFreeEnergy/gufe
-from typing import Tuple
-import math
+from __future__ import annotations
+
 from openff.units import unit
+from typing import Optional, Tuple
 
 from gufe import Component
 
+_CATIONS = {'Cs', 'K', 'Li', 'Na', 'Rb'}
+_ANIONS = {'Cl', 'Br', 'F', 'I'}
+
 
 # really wanted to make this a dataclass but then can't sort & strip ion input
 class SolventComponent(Component):
-
     """Solvent molecules in a chemical state
 
     This component represents the abstract idea of the solvent and ions present
     around the other components, rather than a list of specific water molecules
     and their coordinates.  This abstract representation is later made concrete
-    by specific methods.
+    by specific MD engine methods.
     """
-    def __init__(self, smiles: str = 'O',
-                 ions: Tuple[str, ...] = None,
+    _smiles: str
+    _positive_ion: Optional[str]
+    _negative_ion: Optional[str]
+    _neutralize: bool
+    _ion_concentration: unit.Quantity
+
+    def __init__(self, *,  # force kwarg usage
+                 smiles: str = 'O',
+                 positive_ion: Optional[str] = None,
+                 negative_ion: Optional[str] = None,
                  neutralize: bool = True,
-                 concentration: unit.Quantity = None):
+                 ion_concentration: unit.Quantity = None):
         """
         Parameters
         ----------
         smiles : str, optional
           smiles of the solvent, default 'O' (water)
-        ions : list of str, optional
-          ions in the system, default `None`
+        positive_ion, negative_ion : str, optional
+          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, default `None`
         neutralize : bool, optional
-          if the net charge on the chemical state is neutralized by the ions in this
-          solvent component.  Default `True`
-        concentration : openff-units.unit.Quantity, optional
+          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
           ionic concentration required, default `None`
           this must be supplied with units, e.g. "1.5 * unit.molar"
+
+        Examples
+        --------
+        To create a sodium chloride solution at 0.2 molar concentration::
+
+          >>> s = SolventComponent(position_ion='Na', negative_ion='Cl',
+          ...                      ion_concentration=0.2 * unit.molar)
+
         """
         self._smiles = smiles
-        if ions is not None:
-            # strip and sort so that ('Na', 'Cl-') is equivalent to
-            # ('Cl', 'Na+')
-            self._ions = tuple(sorted(i.strip('+-').capitalize()
-                                      for i in ions))
-        else:
-            self._ions = tuple()
+        if positive_ion is not None:
+            norm = positive_ion.strip('-+').capitalize()
+            if norm not in _CATIONS:
+                raise ValueError(f"Invalid positive ion, got {positive_ion}")
+            positive_ion = norm + '+'
+        self._positive_ion = positive_ion
+        if negative_ion is not None:
+            norm = negative_ion.strip('-+').capitalize()
+            if norm not in _ANIONS:
+                raise ValueError(f"Invalid negative ion, got {negative_ion}")
+            negative_ion = norm + '-'
+        self._negative_ion = negative_ion
+
         self._neutralize = neutralize
-        self._concentration = concentration
+        if ion_concentration is not None:
+            if (not isinstance(ion_concentration, unit.Quantity) or
+               not ion_concentration.is_compatible_with(unit.molar)):
+                raise ValueError(f"ion_concentration must be given in units of"
+                                 f" concentration, got {ion_concentration}")
+            # concentration requires both ions be given
+            if ion_concentration > 0:
+                if self._negative_ion is None or self._positive_ion is None:
+                    raise ValueError("Ions must be given for concentration")
+        self._ion_concentration = ion_concentration
 
     @property
     def smiles(self) -> str:
         """SMILES representation of the solvent molecules"""
         return self._smiles
 
     @property
-    def ions(self) -> Tuple[str, ...]:
-        """The ions in the solvent state"""
-        return self._ions
+    def positive_ion(self) -> Optional[str]:
+        """The cation in the solvent state"""
+        return self._positive_ion
+
+    @property
+    def negative_ion(self) -> Optional[str]:
+        """The anion in the solvent state"""
+        return self._negative_ion
 
     @property
     def neutralize(self) -> bool:
         """If the solvent neutralizes the system overall"""
         return self._neutralize
 
     @property
-    def concentration(self) -> unit.Quantity:
+    def ion_concentration(self) -> unit.Quantity:
         """Concentration of ions in the solvent state"""
-        return self._concentration
+        return self._ion_concentration
 
     @property
     def total_charge(self):
         """Solvents don't have a formal charge defined so this returns None"""
         return None
 
     def __eq__(self, other):
-        try:
-            return (self.smiles == other.smiles and
-                    self.ions == other.ions and
-                    self.concentration == other.concentration)
-        except AttributeError:
-            return False
+        if not isinstance(other, SolventComponent):
+            return NotImplemented
+        return (self.smiles == other.smiles and
+                self.positive_ion == other.positive_ion and
+                self.negative_ion == other.negative_ion and
+                self.ion_concentration == other.ion_concentration)
 
     def __hash__(self):
-        return hash((self.smiles, self.ions, self.concentration))
+        return hash((self.smiles, self.positive_ion, self.negative_ion,
+                     self.ion_concentration))
 
     @classmethod
     def from_dict(cls, d):
@@ -89,6 +131,7 @@ def from_dict(cls, d):
 
     def to_dict(self):
         """For serialization"""
-        return {'smiles': self.smiles, 'ions': self.ions,
-                'concentration': self.concentration,
+        return {'smiles': self.smiles, 'positive_ion': self.positive_ion,
+                'negative_ion': self.negative_ion,
+                'ion_concentration': self.ion_concentration,
                 'neutralize': self._neutralize}

gufe/tests/test_chemicalsystem.py

@@ -14,7 +14,7 @@ def prot_comp(PDB_181L_path):
 
 @pytest.fixture
 def solv_comp():
-    yield gufe.SolventComponent(ions=('K', 'Cl'))
+    yield gufe.SolventComponent(positive_ion='K', negative_ion='Cl')
 
 
 @pytest.fixture

gufe/tests/test_solvents.py

@@ -8,52 +8,86 @@ def test_defaults():
     s = SolventComponent()
 
     assert s.smiles == 'O'
-    assert s.ions == tuple()
-    assert s.concentration is None
+    assert s.positive_ion is None
+    assert s.negative_ion is None
+    assert s.ion_concentration is None
 
 
-@pytest.mark.parametrize('other,', [
-    # test: ordering, charge dropping, case sensitivity
-    ('Cl', 'Na'), ('Na+', 'Cl-'), ('cl', 'na'),
+@pytest.mark.parametrize('pos, neg', [
+    # test: charge dropping, case sensitivity
+    ('Na', 'Cl'), ('Na+', 'Cl-'), ('na', 'cl'),
 ])
-def test_hash(other):
-    s1 = SolventComponent(ions=('Na', 'Cl'))
-    s2 = SolventComponent(ions=other)
+def test_hash(pos, neg):
+    s1 = SolventComponent(positive_ion='Na', negative_ion='Cl')
+    s2 = SolventComponent(positive_ion=pos, negative_ion=neg)
 
     assert s1 == s2
     assert hash(s1) == hash(s2)
-
+    assert s2.positive_ion == 'Na+'
+    assert s2.negative_ion == 'Cl-'
 
 def test_neq():
-    s1 = SolventComponent(ions=('Na', 'Cl'))
-    s2 = SolventComponent(ions=('K', 'Cl'))
+    s1 = SolventComponent(positive_ion='Na', negative_ion='Cl')
+    s2 = SolventComponent(positive_ion='K', negative_ion='Cl')
 
     assert s1 != s2
 
 
 def test_to_dict():
-    s = SolventComponent(ions=('Na', 'Cl'))
+    s = SolventComponent(positive_ion='Na', negative_ion='Cl')
 
-    assert s.to_dict() == {'smiles': 'O', 'ions': ('Cl', 'Na'),
+    assert s.to_dict() == {'smiles': 'O',
+                           'positive_ion': 'Na+',
+                           'negative_ion': 'Cl-',
                            'neutralize': True,
-                           'concentration': None}
+                           'ion_concentration': None}
 
 
 def test_from_dict():
-    s1 = SolventComponent(ions=('Na', 'Cl'),
-                          concentration=1.75 * unit.molar,
+    s1 = SolventComponent(positive_ion='Na', negative_ion='Cl',
+                          ion_concentration=1.75 * unit.molar,
                           neutralize=False)
 
     assert SolventComponent.from_dict(s1.to_dict()) == s1
 
 
 def test_conc():
-    s = SolventComponent(ions=('Na', 'Cl'), concentration=1.75 * unit.molar)
+    s = SolventComponent(positive_ion='Na', negative_ion='Cl',
+                         ion_concentration=1.75 * unit.molar)
+
+    assert s.ion_concentration == unit.Quantity('1.75 M')
+
 
-    assert s.concentration == unit.Quantity('1.75 M')
+@pytest.mark.parametrize('conc,',
+                         [1.22,  # no units, 1.22 what?
+                          1.5 * unit.kg])  # probably a tad much salt
+def test_bad_conc(conc):
+    with pytest.raises(ValueError):
+        _ = SolventComponent(positive_ion='Na', negative_ion='Cl',
+                             ion_concentration=conc)
 
 
 def test_solvent_charge():
-    s = SolventComponent(ions=('Na', 'Cl'), concentration=1.75 * unit.molar)
+    s = SolventComponent(positive_ion='Na', negative_ion='Cl',
+                         ion_concentration=1.75 * unit.molar)
 
     assert s.total_charge is None
+
+
+@pytest.mark.parametrize('pos, neg,', [
+    ('Na', 'C'),
+    ('F', 'I'),
+])
+def test_bad_inputs(pos, neg):
+    with pytest.raises(ValueError):
+        _ = SolventComponent(positive_ion=pos, negative_ion=neg)
+
+
+@pytest.mark.parametrize('pos, neg', [
+    ('Na', None), (None, 'Cl'), (None, None),
+])
+def test_conc_no_ions(pos, neg):
+    # if you specify concentration you must also give ions
+    with pytest.raises(ValueError):
+        _ = SolventComponent(positive_ion=pos, negative_ion=neg,
+                             ion_concentration=1.5 * unit.molar)