feat: merge_these_atom_params

Author: diogomart

meeko/molsetup.py

@@ -1176,7 +1176,7 @@ def get_neighbors(self, atom_index: int):
 
     # endregion
 
-    def merge_terminal_atoms(self, indices, merge_rmin_half=False) -> None:
+    def merge_terminal_atoms(self, indices, merge_rmin_half=False, merge_these_atom_params=()) -> None:
         """
         Primarily for merging hydrogens, but will merge the data for any atom or pseudoatom that is bonded to only one
         other atom.
@@ -1194,6 +1194,8 @@ def merge_terminal_atoms(self, indices, merge_rmin_half=False) -> None:
         """
         if merge_rmin_half and "rmin_half" not in self.atom_params:
             raise ValueError("can't merge rmin_half because it's not in atom_params")
+        if merge_rmin_half and "rmin_half" in merge_these_atom_params:
+            raise ValueError("merge_rmin_half=True wants to merge by volume, rmin_half in merge_these_atom_params wants to merge by radius")
         for index in indices:
             if len(self.get_neighbors(index)) != 1:
                 msg = "Atempted to merge atom %d with %d neighbors. "
@@ -1204,6 +1206,10 @@ def merge_terminal_atoms(self, indices, merge_rmin_half=False) -> None:
             self.atoms[neighbor_index].charge += self.get_charge(index)
             self.atoms[index].charge = 0.0
             self.atoms[index].is_ignore = True
+            for param_key in merge_these_atom_params:
+                value = self.atom_params[param_key][index]
+                self.atom_params[param_key][neighbor_index] += value
+                self.atom_params[param_key][index] = 0.0
             if not merge_rmin_half:
                 continue
             r_neigh = self.atom_params["rmin_half"][neighbor_index]

meeko/preparation.py

@@ -84,6 +84,7 @@ class MoleculePreparation:
     def __init__(
         self,
         merge_these_atom_types=("H",),
+        merge_these_atom_params=(),
         merge_rmin_half=False,
         hydrate=False,
         flexible_amides=False,
@@ -120,6 +121,7 @@ def __init__(
         Parameters
         ----------
         merge_these_atom_types
+        merge_these_atom_params
         merge_rmin_half
         hydrate
         flexible_amides
@@ -162,6 +164,7 @@ def __init__(
 
         self.deprecated_setup_access = None
         self.merge_these_atom_types = merge_these_atom_types
+        self.merge_these_atom_params = merge_these_atom_params
         self.merge_rmin_half = merge_rmin_half
         self.hydrate = hydrate
         self.flexible_amides = flexible_amides
@@ -605,6 +608,17 @@ def prepare(
             self.dihedral_params,
         )
 
+        if self.crippen or self.crippen_as_solpar:
+            add_crippen_to_molsetup(setup)
+            if not self.crippen:
+                setup.atom_params["ad4_sol_par"] = setup.atom_params.pop("crippen")
+            elif self.crippen_as_solpar:
+                setup.atom_params["ad4_sol_par"] = [value for value in setup.atom_params["crippen"]]
+
+        if self.override_ad4sol_par_including_q:
+            qasp = self.override_ad4sol_par_including_q_qasp
+            set_ad4sol_par_including_q(setup, qasp)
+
         # Convert molecule to graph and apply trained Espaloma model
         # skip if charges are read from template
         if self.dihedral_model == "espaloma" or (self.charge_model == "espaloma" and self.compute_charges):
@@ -634,7 +648,7 @@ def prepare(
             for atom in setup.atoms:
                 if atom.atom_type == atype_to_merge:
                     indices.add(atom.index)
-        setup.merge_terminal_atoms(indices, self.merge_rmin_half)
+        setup.merge_terminal_atoms(indices, self.merge_rmin_half, self.merge_these_atom_params)
 
         # 3.  assign bond types
         #     - all single bonds rotatable except some amides and SMARTS rigidification
@@ -672,17 +686,6 @@ def prepare(
                     new_atom_info = orig_pdbinfo._replace(name=new_name)
                     atom.pdbinfo = new_atom_info
 
-        if self.crippen or self.crippen_as_solpar:
-            add_crippen_to_molsetup(setup)
-            if not self.crippen:
-                setup.atom_params["ad4_sol_par"] = setup.atom_params.pop("crippen")
-            elif self.crippen_as_solpar:
-                setup.atom_params["ad4_sol_par"] = [value for value in setup.atom_params["crippen"]]
-
-        if self.override_ad4sol_par_including_q:
-            qasp = self.override_ad4sol_par_including_q_qasp
-            set_ad4sol_par_including_q(setup, qasp)
-
         if self.reactive_smarts is None:
             setups = [setup]
         else:

test/parameterization_test.py

@@ -114,13 +114,30 @@ def test_mbondi():
     rdDistGeom.EmbedMolecule(mol, etkdg_v3)
     mk_prep = MoleculePreparation(load_atom_params="gb_mbondi")
     molsetup = mk_prep(mol)[0]
-    print(molsetup.atom_params)
     assert 1.3 in molsetup.atom_params["mbondi_radius"]  # H bound to N or C
     assert 0.8 in molsetup.atom_params["mbondi_radius"]  # H bound to O or S
     assert 1.2 not in molsetup.atom_params["mbondi_radius"]  # other H
     assert molsetup.atom_params["gb_screen"].count(0.72) == 4  # C
     assert molsetup.atom_params["gb_screen"].count(0.79) == 2  # N
 
+def test_merge_crippen():
+    mol = Chem.AddHs(Chem.MolFromSmiles("OC(F)F"))
+    etkdg_v3 = rdDistGeom.ETKDGv3()
+    rdDistGeom.EmbedMolecule(mol, etkdg_v3)
+    mk_prep = MoleculePreparation(crippen=True, load_atom_params=["ad4_types"])
+    molsetup = mk_prep(mol)[0]
+    mk_prep = MoleculePreparation(crippen=True, load_atom_params=["ad4_types"], merge_these_atom_params=["crippen"])
+    molsetup_merge = mk_prep(mol)[0]
+    h_mask = np.array([atom.atom_type == "H" for atom in molsetup.atoms])
+    crippen_orig = np.array(molsetup.atom_params["crippen"])
+    crippen_merge = np.array(molsetup_merge.atom_params["crippen"])
+    abs_diff = np.abs(crippen_merge - crippen_orig)
+    assert abs(sum(crippen_orig) - sum(crippen_merge)) < 1e-8
+    assert np.max(crippen_orig[h_mask]) > 1e-2
+    assert np.max(crippen_merge[h_mask]) < 1e-8
+    assert np.min(abs_diff[h_mask]) > 1e-2
+    assert np.max(abs_diff[~h_mask]) > 1e-2
+
 def test_override_ad4sol_par_including_q():
     qasp = 0.01097
     mol = Chem.AddHs(Chem.MolFromSmiles("c1nc[nH]c1CO"))