Add checks for partial charges

gufe/components/explicitmoleculecomponent.py

@@ -1,4 +1,5 @@
 import json
+import numpy as np
 import warnings
 from rdkit import Chem
 from typing import Optional
@@ -36,6 +37,48 @@ def _ensure_ofe_name(mol: RDKitMol, name: str) -> str:
     return name
 
 
+def _check_partial_charges(mol: RDKitMol) -> None:
+    """
+    Checks for the presence of partial charges.
+
+    Raises
+    ------
+    ValueError
+      * If the partial charges are not of length atoms.
+      * If the sum of partial charges is not equal to the
+        formal charge.
+    UserWarning
+      * If partial charges are found.
+      * If the partial charges are near 0 for all atoms.
+    """
+    if 'atom.dprop.PartialCharge' not in mol.GetPropNames():
+        return
+
+    p_chgs = np.array(
+        mol.GetProp('atom.dprop.PartialCharge').split(), dtype=float
+    )
+
+    if len(p_chgs) != mol.GetNumAtoms():
+        errmsg = (f"Incorrect number of partial charges: {len(p_chgs)} "
+                  f" were provided for {mol.GetNumAtoms()} atoms")
+        raise ValueError(errmsg)
+
+    if (sum(p_chgs) - Chem.GetFormalCharge(mol)) > 0.01:
+        errmsg = (f"Sum of partial charges {sum(p_chgs)} differs from "
+                  f"RDKit formal charge {Chem.GetFormalCharge(mol)}")
+        raise ValueError(errmsg)
+
+    if np.all(np.isclose(p_chgs, 0.0)):
+        wmsg = (f"Partial charges provided all equal to "
+                "zero. These may be ignored by some Protocols.")
+        warnings.warn(wmsg)
+    else:
+        wmsg = ("Partial charges have been provided, these will "
+                "preferentially be used instead of generating new "
+                "partial charges")
+        warnings.warn(wmsg)
+
+
 class ExplicitMoleculeComponent(Component):
     """Base class for explicit molecules.
 
@@ -48,6 +91,7 @@ class ExplicitMoleculeComponent(Component):
 
     def __init__(self, rdkit: RDKitMol, name: str = ""):
         name = _ensure_ofe_name(rdkit, name)
+        _check_partial_charges(rdkit)
         conformers = list(rdkit.GetConformers())
         if not conformers:
             raise ValueError("Molecule was provided with no conformers.")

gufe/tests/test_smallmoleculecomponent.py

@@ -5,6 +5,7 @@
 import importlib.resources
 try:
     import openff.toolkit.topology
+    from openff.units import unit
 except ImportError:
     HAS_OFFTK = False
 else:
@@ -198,6 +199,39 @@ def test_to_off_name(self, named_ethane):
         assert off_ethane.name == 'ethane'
 
 
+@pytest.mark.skipif(not HAS_OFFTK, reason="no openff tookit available")
+class TestSmallMoleculeComponentPartialCharges:
+    @pytest.fixture(scope='function')
+    def charged_off_ethane(self, ethane):
+        off_ethane = ethane.to_openff()
+        off_ethane.assign_partial_charges(partial_charge_method='am1bcc')
+        return off_ethane
+
+    def test_partial_charges_warning(self, charged_off_ethane):
+        matchmsg = "Partial charges have been provided"
+        with pytest.warns(UserWarning, match=matchmsg):
+            SmallMoleculeComponent.from_openff(charged_off_ethane)
+
+    def test_partial_charges_zero_warning(self, charged_off_ethane):
+        charged_off_ethane.partial_charges[:] = 0 * unit.elementary_charge
+        matchmsg = "Partial charges provided all equal to zero"
+        with pytest.warns(UserWarning, match=matchmsg):
+            SmallMoleculeComponent.from_openff(charged_off_ethane)
+
+    def test_partial_charges_not_formal_error(self, charged_off_ethane):
+        charged_off_ethane.partial_charges[:] = 1 * unit.elementary_charge
+        with pytest.raises(ValueError, match="Sum of partial charges"):
+            SmallMoleculeComponent.from_openff(charged_off_ethane)
+
+    def test_partial_charges_too_few_atoms(self):
+        mol = Chem.AddHs(Chem.MolFromSmiles("CC"))
+        Chem.AllChem.Compute2DCoords(mol)
+        mol.SetProp('atom.dprop.PartialCharge', '1')
+
+        with pytest.raises(ValueError, match="Incorrect number of"):
+            SmallMoleculeComponent.from_rdkit(mol)
+
+
 @pytest.mark.parametrize('mol, charge', [
     ('CC', 0), ('CC[O-]', -1),
 ])