Make charge symmetrization stereo-aware

tests/test_handlers.py

@@ -10,7 +10,7 @@
 from rdkit import Chem
 from rdkit.Chem import AllChem, rdmolops
 
-from timemachine.constants import ONE_4PI_EPS0
+from timemachine.constants import DEFAULT_FF, ONE_4PI_EPS0
 from timemachine.ff import Forcefield
 from timemachine.ff.charges import AM1CCC_CHARGES
 from timemachine.ff.handlers import bonded, nonbonded
@@ -404,10 +404,10 @@ def test_am1_ccc():
 
     # fmt: off
     ligand_params = np.array([
-        -5.36348, 6.98008, -1.68885, -4.04439, 1.93568, -2.25534,
-        -0.87637, -1.8679, -0.66892, -1.6588, 7.67174, -6.1465,
-        -7.26945, 0.93978, 0.93978, 0.93978, 1.83702, 1.69675,
-        1.71667, 1.90762, 5.27508,
+        -5.36358, 6.98002, -1.68898, -4.04447, 1.9357, -2.2555,
+        -0.87654, -1.86804, -0.66777, -1.65897, 7.67174, -6.14658,
+        -7.26958, 0.93984, 0.93984, 0.93984, 1.83699, 1.6967,
+        1.71663, 1.90759, 5.27512,
         ]
     )
     # fmt: on
@@ -691,23 +691,23 @@ def mol_with_precomputed_charges():
     # fmt: off
     precomputed_charges = dict()
     precomputed_charges['ubuntu-18.04.4'] = np.array([
-        -3.385959, 0.26308782, 2.30802986, 1.11718023, -2.11460365,
-        0.9725527, -9.18744615, 0.04042971, -2.49344069, 1.87709392,
-        -1.64877796, 0.57591713, -1.68873618, 1.19414994, -1.04834368,
-        -1.52407053, 0.30670004, -1.26003975, -1.3233364, 33.79971321,
-        -10.86357005, -10.86357005, -1.72362592, -0.10184044, -1.47951542,
-        -1.72362592, -11.32255943, 2.74250182, 1.64583123, 1.55000221,
-        1.69816584, 1.83477824, 2.11436794, 1.99956175, 1.98447416,
-        2.86426252, 2.86426252])
+        -3.38603383, 0.26304593, 2.30806183, 1.11718269, -2.11466618,
+        0.97255779, -9.18759228, 0.04222331, -2.49355872, 1.87712273,
+        -1.64885893, 0.57591133, -1.68888899, 1.19419578, -1.04842043,
+        -1.52418996, 0.3067501, -1.2601894, -1.32341895, 33.79966262,
+        -10.86363258, -10.86363258, -1.72379533, -0.10198921, -1.47957496,
+        -1.72379533, -11.32267114, 2.74244351, 1.64587022, 1.54997692,
+        1.69815056, 1.83481513, 2.11438216, 1.99952399, 1.98443746,
+        2.86429782, 2.86429782])
     precomputed_charges['ubuntu-20.04.2'] = np.array([
-       -3.38772701, 0.26721331, 2.30272567, 1.11670881, -2.10741359,
-       0.96123709, -9.15420638, 0.03088217, -2.48495405, 1.872379,
-       -1.65467143, 0.58181069, -1.6896792, 1.19155677, -1.04775431,
-       -1.51994506, 0.3044605, -1.25933253, -1.32628322, 33.75586471,
-       -10.85178312, -10.85178312, -1.72386163, -0.10278342, -1.4785724,
-       -1.72386163, -11.31914145, 2.74450571, 1.64253075, 1.55165236,
-       1.69875511, 1.83418879, 2.1147215, 1.99413971, 1.98565289,
-       2.86638427, 2.86638427])
+        -3.3878352035023025, 0.26727080686425936, 2.3026757907523128, 1.1166965996609888, -2.107513358517747,
+        0.9612083769496058, -9.154301226767831, 0.03288091151318266, -2.4851042282886255, 1.8723247315824088,
+        -1.6547659251138624, 0.5817522843277193, -1.6897608745990291, 1.191577230801982, -1.0478643525397062,
+        -1.5200599330326403, 0.3044723600341638, -1.2594928059462964, -1.326439016721788, 33.755887188134906,
+        -10.851852320219985, -10.851852320219985, -1.724017252391697, -0.10285399008405359, -1.4786711130785013,
+        -1.724017252391697, -11.319224351341875, 2.7444770823135793, 1.6425628940842003, 1.5516172308731973,
+        1.6987242017361495, 1.8342000296623708, 2.114672322516452, 1.9941244248440126, 1.9856416523206082,
+        2.866430198256134, 2.866430198256134])
     # fmt: on
 
     return dict(mol=test_mol, precomputed_charges=precomputed_charges)
@@ -901,6 +901,35 @@ def test_lennard_jones_handler():
     assert np.all(adjoints[mask] == 0.0)
 
 
+def test_charge_symmetry():
+    """Expect only 2 distinct partial charges applied to benzene"""
+
+    # molecule with 2 atomic symmetry classes
+    benzene = "C1=CC=CC=C1"
+    mol = Chem.MolFromSmiles(benzene)
+    mol = Chem.AddHs(mol)
+
+    # expect 2 distinct charges, {+q, -q}
+    ff = Forcefield.load_from_file(DEFAULT_FF)
+    assigned_charges = np.array(ff.q_handle.parameterize(mol))
+    assert len(set(assigned_charges)) == 2
+
+
+def test_charge_stereo_symmetry():
+    """Expect distinct charges on all atoms, in a case where symmetry is broken by stereo"""
+
+    # molecule where # stereo-oblivious symmetry classes = 8
+    # but # stereo-sensitive symmetry classes = n_atoms = 16
+    # https://github.com/proteneer/timemachine/pull/738#issuecomment-1126179492
+    smi = r"O/C(=C(/O)\C(\Cl)=C(\F)Cl)/C(/Cl)=C(\F)Cl"
+    mol = Chem.AddHs(Chem.MolFromSmiles(smi))
+
+    # expect distinct charges on all atoms
+    ff = Forcefield.load_from_file(DEFAULT_FF)
+    assigned_charges = np.array(ff.q_handle.parameterize(mol))
+    assert len(set(assigned_charges)) == mol.GetNumAtoms()
+
+
 def test_symmetric_am1ccc():
     """Assert that (symmetric_bond_smarts, +1.0) has same behavior as (symmetric_bond_smarts, 0.0) on one test mol"""
 

timemachine/ff/handlers/nonbonded.py

@@ -13,6 +13,7 @@
 from timemachine.ff.handlers.serialize import SerializableMixIn
 from timemachine.ff.handlers.utils import canonicalize_bond
 from timemachine.ff.handlers.utils import match_smirks as rd_match_smirks
+from timemachine.ff.handlers.utils import symmetrize
 from timemachine.graph_utils import convert_to_nx
 
 AM1_CHARGE_CACHE = "AM1Cache"
@@ -79,16 +80,27 @@ def oe_generate_conformations(oemol, sample_hydrogens=True):
 
 
 def oe_assign_charges(mol, charge_model=AM1BCCELF10):
-    """assign partial charges, then premultiply by sqrt(ONE_4PI_EPS0)
-    as an optimization"""
+    """Convert RD mol to OE mol, compute charges, convert to TM units
+
+    Notes
+    -----
+    * Passes `symmetrize=False` to toolkit functions where possible,
+        since `symmetrize=True` symmetrizes using stereo-insensitive equivalence classes
+        See: https://github.com/proteneer/timemachine/pull/738#issuecomment-1124280659
+        Exception: AM1BCCELF10 not configurable
+
+    See Also
+    --------
+    * ff/handlers/utils.py::symmetrize -- symmetrize using stereo-aware symmetry classes
+    """
 
     # imported here for optional dependency
     from openeye import oequacpac
 
     charge_engines = {
-        AM1: oequacpac.OEAM1Charges(symmetrize=True),
-        AM1ELF10: oequacpac.OEELFCharges(oequacpac.OEAM1Charges(symmetrize=True), 10),
-        AM1BCC: oequacpac.OEAM1BCCCharges(symmetrize=True),
+        AM1: oequacpac.OEAM1Charges(symmetrize=False),
+        AM1ELF10: oequacpac.OEELFCharges(oequacpac.OEAM1Charges(symmetrize=False), 10),
+        AM1BCC: oequacpac.OEAM1BCCCharges(symmetrize=False),
         AM1BCCELF10: oequacpac.OEAM1BCCELF10Charges(),
     }
     charge_engine = charge_engines[charge_model]
@@ -110,6 +122,7 @@ def oe_assign_charges(mol, charge_model=AM1BCCELF10):
             raise Exception(f"Unable to assign charges for '{oemol.GetTitle()}'")
 
     partial_charges = np.array([atom.GetPartialCharge() for atom in oemol.GetAtoms()])
+    partial_charges = symmetrize(partial_charges, mol)
 
     # Verify that the charges sum up to an integer
     net_charge = np.sum(partial_charges)
@@ -532,11 +545,13 @@ def static_parameterize(params, smirks, mol):
 
         """
         am1_charges = compute_or_load_am1_charges(mol)
+        symmetrized_charges = symmetrize(am1_charges, mol)
+
         bond_idxs, type_idxs = compute_or_load_bond_smirks_matches(mol, smirks)
 
         deltas = params[type_idxs]
-        q_params = apply_bond_charge_corrections(am1_charges, bond_idxs, deltas)
+        q_params = apply_bond_charge_corrections(symmetrized_charges, bond_idxs, deltas)
 
-        assert q_params.shape[0] == mol.GetNumAtoms()  # check that return shape is consistent with input mol
+        assert q_params.shape == (mol.GetNumAtoms(),)  # check that return shape is consistent with input mol
 
         return q_params

timemachine/ff/handlers/utils.py

@@ -1,3 +1,9 @@
+import jax
+
+jax.config.update("jax_enable_x64", True)
+
+import numpy as np
+from numpy.typing import NDArray
 from rdkit import Chem
 
 
@@ -67,3 +73,34 @@ def match_smirks(mol, smirks):
         matches.append(tuple(mas))
 
     return matches
+
+
+def get_symmetry_classes(mol: Chem.Mol) -> NDArray:
+    """Assign arbitrary integer indices to atoms in mol, such that:
+    * symmetric atoms are assigned distinct indices
+    * indices are contiguous from 0 to n_classes
+    """
+    symmetry_classes = list(Chem.CanonicalRankAtoms(mol, breakTies=False, includeChirality=True))
+    n_classes = len(set(symmetry_classes))
+
+    # make contiguous from 0 to n_classes
+    idx_map = {old_idx: new_idx for (new_idx, old_idx) in enumerate(set(symmetry_classes))}
+    symmetry_classes = [idx_map[old_idx] for old_idx in symmetry_classes]
+    assert set(symmetry_classes) == set(range(n_classes))
+
+    return np.array(symmetry_classes)
+
+
+def symmetrize(per_particle_params: NDArray, mol: Chem.Mol) -> NDArray:
+    """Replace per_particle_params[atom] with mean([per_particle_params[a] for a in symmetry_class(atom)])"""
+
+    # compute the mean per symmetry class
+    symmetry_classes = get_symmetry_classes(mol)
+    class_sums = jax.ops.segment_sum(per_particle_params, symmetry_classes)
+    class_means = class_sums / np.bincount(symmetry_classes)
+
+    # assign all members of a symmetry class the class's mean value
+    symmetrized = class_means[symmetry_classes]
+    assert symmetrized.shape == per_particle_params.shape
+
+    return symmetrized