Env handler smirks

tests/common.py

@@ -24,7 +24,6 @@
 from timemachine.potentials import Nonbonded
 from timemachine.potentials.potential import GpuImplWrapper
 from timemachine.potentials.types import PotentialFxn
-from timemachine.testsystems.data.ildn_params import get_amber99ildn_patterns
 
 HILBERT_GRID_DIM = 128
 
@@ -394,12 +393,12 @@ def load_split_forcefields() -> SplitForcefield:
     SIG_SCALE = 0.5  # smaller change to prevent overflow
     EPS_SCALE = 2.0
 
-    patterns = get_amber99ildn_patterns()
-    protein_smirks = [x[0] for x in patterns]
-    protein_params = np.ones((len(protein_smirks),)) * 0.5
-
     ff_ref = Forcefield.load_default()
 
+    assert ff_ref.q_handle is not None
+    protein_smirks = ff_ref.q_handle.smirks
+    protein_params = np.ones((len(protein_smirks),)) * 0.5
+
     ff_intra = Forcefield.load_default()
     assert ff_intra.q_handle_intra is not None
     assert ff_intra.lj_handle_intra is not None

tests/test_handlers.py

@@ -1,4 +1,5 @@
 import functools
+from collections import defaultdict
 from copy import deepcopy
 from importlib import resources
 
@@ -17,7 +18,6 @@
 from timemachine.ff.charges import AM1CCC_CHARGES
 from timemachine.ff.handlers import bonded, nonbonded
 from timemachine.md import builders
-from timemachine.testsystems.data.ildn_params import get_amber99ildn_patterns
 
 pytestmark = [pytest.mark.nocuda]
 
@@ -1145,19 +1145,13 @@ def test_env_bcc_peptide_symmetries(protein_path_and_symmetries):
     Test that we can compute BCCs to generate per atom charge offsets and that they can be differentiated
     """
     protein_path, expected_symmetries = protein_path_and_symmetries
-    patterns = get_amber99ildn_patterns()
-
-    smirks = [x[0] for x in patterns]
+    smirks = [x[0] for x in AM1CCC_CHARGES["patterns"]]
     params = np.random.rand(len(smirks)) - 0.5
 
     with resources.path("timemachine.testsystems.data", protein_path) as path_to_pdb:
         host_pdb = app.PDBFile(str(path_to_pdb))
         topology = host_pdb.topology
-
-    pbcc = nonbonded.EnvironmentBCCHandler(smirks, params, DEFAULT_PROTEIN_FF, DEFAULT_WATER_FF, topology)
-
-    assert len(pbcc.bond_atomic_numbers) == len(pbcc.bond_idxs)
-    assert (6, 6) in pbcc.bond_atomic_numbers
+        pbcc = nonbonded.EnvironmentBCCHandler(smirks, params, DEFAULT_PROTEIN_FF, DEFAULT_WATER_FF, topology)
 
     # raw charges are correct are in the order of atoms in the topology
     raw_charges = np.array(pbcc.parameterize(np.zeros_like(params)))
@@ -1193,19 +1187,18 @@ def test_environment_bcc_full_protein(protein_path):
     """
     Test that we can compute BCCs to generate per atom charge offsets and that they can be differentiated
     """
-    patterns = get_amber99ildn_patterns()
-    smirks = [x[0] for x in patterns]
-    params = [x[1] for x in patterns]
-    params = np.random.rand(len(params)) - 0.5
+    patterns = [smirks for (smirks, param) in AM1CCC_CHARGES["patterns"]]
+    np.random.seed(2024)
+    params = np.random.rand(len(patterns)) - 0.5
 
     with resources.path("timemachine.testsystems.data", protein_path) as path_to_pdb:
         host_pdb = app.PDBFile(str(path_to_pdb))
         _, _, _, topology, _ = builders.build_protein_system(host_pdb, DEFAULT_PROTEIN_FF, DEFAULT_WATER_FF)
 
-    pbcc = nonbonded.EnvironmentBCCHandler(smirks, params, DEFAULT_PROTEIN_FF, DEFAULT_WATER_FF, topology)
+    pbcc = nonbonded.EnvironmentBCCHandler(patterns, params, DEFAULT_PROTEIN_FF, DEFAULT_WATER_FF, topology)
 
     # test that we can mechanically parameterize everything
-    pbcc.parameterize(params)
+    final_q_params = pbcc.parameterize(params)
 
     def loss_fn(bcc_params):
         res = pbcc.parameterize(bcc_params)
@@ -1219,7 +1212,7 @@ def loss_fn(bcc_params):
 
     # test that the partial handler gives the same results
     ff = Forcefield.load_default()
-    partial_cc = nonbonded.EnvironmentBCCPartialHandler(smirks, params, None)
+    partial_cc = nonbonded.EnvironmentBCCPartialHandler(patterns, params, None)
     pbcc2 = partial_cc.get_env_handle(topology, ff)
     np.testing.assert_array_equal(pbcc.parameterize(params), pbcc2.parameterize(params))
 
@@ -1231,3 +1224,27 @@ def loss_fn2(bcc_params):
 
     assert loss_fn(params) == loss_fn2(params)
     np.testing.assert_array_equal(grad_fn(params), grad_fn2(params))
+
+    # now check for correctness on the full protein set
+    # if the initial charges are the same, then the final charges should
+    # also be the same by symmetry
+    atom_idx_to_res_name = {}
+    cur_atom = 0
+    for i in range(len(pbcc.template_for_residue)):
+        tfr = pbcc.template_for_residue[i]
+        n_atoms = len(tfr.atoms)
+        for j in range(n_atoms):
+            atom_idx_to_res_name[cur_atom + j] = f"{tfr.name}"
+        cur_atom += n_atoms
+
+    # charges that are the same and belong to the same residue should
+    # be the same after applying the corrections
+    init_q_params = pbcc.initial_charges
+    sym_charge_idxs = defaultdict(list)
+    for i, q in enumerate(init_q_params):
+        sym_charge_idxs[(q, atom_idx_to_res_name[i])].append(i)
+
+    for k, group in sym_charge_idxs.items():
+        first = group[0]
+        for other in group[1:]:
+            np.testing.assert_almost_equal(final_q_params[other], final_q_params[first])

timemachine/ff/__init__.py

@@ -54,6 +54,7 @@ class Forcefield:
             nonbonded.SimpleChargeHandler,
             nonbonded.AM1BCCHandler,
             nonbonded.AM1CCCHandler,
+            nonbonded.AM1BCCCCCHandler,
             nonbonded.PrecomputedChargeHandler,
         ]
     ]
@@ -62,6 +63,7 @@ class Forcefield:
             nonbonded.SimpleChargeIntraHandler,
             nonbonded.AM1BCCIntraHandler,
             nonbonded.AM1CCCIntraHandler,
+            nonbonded.AM1BCCCCCIntraHandler,
             nonbonded.PrecomputedChargeHandler,
         ]
     ]
@@ -184,6 +186,7 @@ def from_handlers(
             elif isinstance(
                 handle,
                 (
+                    nonbonded.AM1BCCCCCIntraHandler,
                     nonbonded.AM1CCCIntraHandler,
                     nonbonded.AM1BCCIntraHandler,
                     nonbonded.SimpleChargeIntraHandler,
@@ -203,6 +206,7 @@ def from_handlers(
             elif isinstance(
                 handle,
                 (
+                    nonbonded.AM1BCCCCCHandler,
                     nonbonded.AM1CCCHandler,
                     nonbonded.AM1BCCHandler,
                     nonbonded.SimpleChargeHandler,
@@ -230,7 +234,9 @@ def from_handlers(
         if q_handle_intra is None:
             # Copy the forcefield parameters to the intramolecular term if not
             # already handled.
-            if isinstance(q_handle, nonbonded.AM1CCCHandler):
+            if isinstance(q_handle, nonbonded.AM1BCCCCCHandler):
+                q_handle_intra = nonbonded.AM1BCCCCCIntraHandler(q_handle.smirks, q_handle.params, q_handle.props)
+            elif isinstance(q_handle, nonbonded.AM1CCCHandler):
                 q_handle_intra = nonbonded.AM1CCCIntraHandler(q_handle.smirks, q_handle.params, q_handle.props)
             elif isinstance(q_handle, nonbonded.AM1BCCHandler):
                 q_handle_intra = nonbonded.AM1BCCIntraHandler(q_handle.smirks, q_handle.params, q_handle.props)