Merge pull request #326 from jthorton/charge_checking

Improve partial charge check

gufe/components/explicitmoleculecomponent.py

@@ -41,6 +41,11 @@ def _check_partial_charges(mol: RDKitMol) -> None:
     """
     Checks for the presence of partial charges.
 
+    Note
+    ----
+    We ensure the charges are set as atom properties
+    to ensure they are detected by OpenFF
+
     Raises
     ------
     ValueError
@@ -68,6 +73,18 @@ def _check_partial_charges(mol: RDKitMol) -> None:
                   f"RDKit formal charge {Chem.GetFormalCharge(mol)}")
         raise ValueError(errmsg)
 
+    # set the charges on the atoms if not already set
+    for i, charge in enumerate(p_chgs):
+        atom = mol.GetAtomWithIdx(i)
+        if not atom.HasProp("PartialCharge"):
+            atom.SetDoubleProp("PartialCharge", charge)
+        else:
+            atom_charge = atom.GetDoubleProp("PartialCharge")
+            if not np.isclose(atom_charge, charge):
+                errmsg = (f"non-equivalent partial charges between atom and "
+                          f"molecule properties: {atom_charge} {charge}")
+                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.")

gufe/tests/test_smallmoleculecomponent.py

@@ -264,6 +264,47 @@ def test_partial_charges_too_few_atoms(self):
         with pytest.raises(ValueError, match="Incorrect number of"):
             SmallMoleculeComponent.from_rdkit(mol)
 
+    def test_partial_charges_applied_to_atoms(self):
+        """
+        Make sure that charges set at the molecule level
+        are transferred to atoms and picked up by openFF.
+        """
+        mol = Chem.AddHs(Chem.MolFromSmiles("C"))
+        Chem.AllChem.Compute2DCoords(mol)
+        # add some fake charges at the molecule level
+        mol.SetProp('atom.dprop.PartialCharge', '-1 0.25 0.25 0.25 0.25')
+        matchmsg = "Partial charges have been provided"
+        with pytest.warns(UserWarning, match=matchmsg):
+            ofe = SmallMoleculeComponent.from_rdkit(mol)
+        # convert to openff and make sure the charges are set
+        off_mol = ofe.to_openff()
+        assert off_mol.partial_charges is not None
+        # check ordering is the same
+        rdkit_mol_with_charges = ofe.to_rdkit()
+        for i, charge in enumerate(off_mol.partial_charges.m):
+            rdkit_atom = rdkit_mol_with_charges.GetAtomWithIdx(i)
+            assert rdkit_atom.GetDoubleProp("PartialCharge") == charge
+
+    def test_inconsistent_charges(self, charged_off_ethane):
+        """
+        An error should be raised if the atom and molecule level
+        charges do not match.
+        """
+        mol = Chem.AddHs(Chem.MolFromSmiles("C"))
+        Chem.AllChem.Compute2DCoords(mol)
+        # add some fake charges at the molecule level
+        mol.SetProp('atom.dprop.PartialCharge', '-1 0.25 0.25 0.25 0.25')
+        # set different charges to the atoms
+        for atom in mol.GetAtoms():
+            atom.SetDoubleProp("PartialCharge", 0)
+
+        # make sure the correct error is raised
+        msg = ("non-equivalent partial charges between "
+               "atom and molecule properties")
+        with pytest.raises(ValueError, match=msg):
+            SmallMoleculeComponent.from_rdkit(mol)
+
+
 
 @pytest.mark.parametrize('mol, charge', [
     ('CC', 0), ('CC[O-]', -1),