Use cache to speed up applying forcefield to connections

gmso/parameterization/molecule_utils.py

@@ -56,3 +56,39 @@ def assert_no_boundary_bonds(top):
         assert site1.molecule == site2.molecule, assertion_msg.format(
             site1.name, site1.molecule, site2.name, site2.molecule
         )
+
+
+def build_molecule_connection_index(top, is_group=False):
+    """Build a dict mapping molecule/group label -> connection lists.
+
+    Returns:
+        dict: {
+            label: {
+                "bonds": [...],
+                "angles": [...],
+                "dihedrals": [...],
+                "impropers": [...]
+            }
+        }
+    """
+    index = {}
+
+    def _label_of(site):
+        return getattr(site, "group") if is_group else site.molecule.name
+
+    def _bucket(connections, key):
+        for conn in connections:
+            members = conn.connection_members
+            labels = {_label_of(s) for s in members}
+            label = labels.pop()
+            entry = index.setdefault(
+                label, {"bonds": [], "angles": [], "dihedrals": [], "impropers": []}
+            )
+            entry[key].append(conn)
+
+    _bucket(top.bonds, "bonds")
+    _bucket(top.angles, "angles")
+    _bucket(top.dihedrals, "dihedrals")
+    _bucket(top.impropers, "impropers")
+
+    return index

gmso/parameterization/topology_parameterizer.py

@@ -20,6 +20,7 @@
 )
 from gmso.parameterization.molecule_utils import (
     assert_no_boundary_bonds,
+    build_molecule_connection_index,
     molecule_angles,
     molecule_bonds,
     molecule_dihedrals,
@@ -132,19 +133,24 @@ def _parameterize_connections(
         ff,
         label_type=None,
         label=None,
+        connection_index=None,
     ):
         """Parameterize connections with appropriate potentials from the forcefield."""
         if label_type and label:
-            bonds = molecule_bonds(top, label, True if label_type == "group" else False)
-            angles = molecule_angles(
-                top, label, True if label_type == "group" else False
-            )
-            dihedrals = molecule_dihedrals(
-                top, label, True if label_type == "group" else False
-            )
-            impropers = molecule_impropers(
-                top, label, True if label_type == "group" else False
-            )
+            if connection_index is not None:
+                entry = connection_index.get(
+                    label, {"bonds": [], "angles": [], "dihedrals": [], "impropers": []}
+                )
+                bonds = entry["bonds"]
+                angles = entry["angles"]
+                dihedrals = entry["dihedrals"]
+                impropers = entry["impropers"]
+            else:
+                is_group = True if label_type == "group" else False
+                bonds = molecule_bonds(top, label, is_group)
+                angles = molecule_angles(top, label, is_group)
+                dihedrals = molecule_dihedrals(top, label, is_group)
+                impropers = molecule_impropers(top, label, is_group)
         else:
             bonds = top.bonds
             angles = top.angles
@@ -194,27 +200,48 @@ def _parameterize_virtual_sites(self, top, sites, bonds, ff):
     def _apply_connection_parameters(self, connections, ff, error_on_missing=True):
         """Find and assign potentials from the forcefield for the provided connections."""
         visited = dict()
+        sig_cache = dict()
         for connection in connections:
-            group, connection_identifiers = self.connection_identifier(connection)
-            match = None
-            for identifier_key in connection_identifiers:
-                if tuple(identifier_key) in visited:
-                    match = visited[tuple(identifier_key)]
-                    break
-
-                match = ff.get_potential(
-                    group=group,
-                    key=identifier_key,
-                    exact_match=True,
+            use_classes = all(
+                [site.atom_type.atomclass for site in connection.connection_members]
+            )
+            if use_classes:
+                sig = tuple(
+                    site.atom_type.atomclass for site in connection.connection_members
                 )
-                if match:
-                    visited[tuple(identifier_key)] = match
-                    break
+            else:
+                sig = tuple(
+                    site.atom_type.name for site in connection.connection_members
+                )
+            if getattr(connection, "bonds", None):
+                sig += tuple(b.bond_order for b in connection.bonds)
+            elif hasattr(connection, "bond_order"):
+                sig += (connection.bond_order,)
+            if sig in sig_cache:
+                match = sig_cache[sig]
+            else:
+                group, connection_identifiers = self.connection_identifier(connection)
+                match = None
+                for identifier_key in connection_identifiers:
+                    id_tuple = tuple(identifier_key)
+                    if id_tuple in visited:
+                        match = visited[id_tuple]
+                        break
+                    match = ff.get_potential(
+                        group=group,
+                        key=identifier_key,
+                        exact_match=True,
+                    )
+                    if match:
+                        visited[id_tuple] = match
+                        break
+                sig_cache[sig] = match
 
             if not match and error_on_missing:
+                group = POTENTIAL_GROUPS[type(connection)]
                 raise ParameterizationError(
                     f"No parameters found for connection {connection}, group: {group}, "
-                    f"identifiers: {connection_identifiers} in the Forcefield."
+                    f"identifiers: {connection_identifiers} in the force field."
                 )
             elif match:
                 setattr(connection, group, match[0].clone(self.config.fast_copy))
@@ -269,7 +296,13 @@ def _apply_virtual_site_parameters(self, virtual_sites, ff, error_on_missing=Tru
                     )
 
     def _parameterize(
-        self, top, typemap, label_type=None, label=None, speedup_by_moltag=False
+        self,
+        top,
+        typemap,
+        label_type=None,
+        label=None,
+        speedup_by_moltag=False,
+        connection_index=None,
     ):
         """Parameterize a topology/subtopology based on an atomtype map."""
         if label and label_type:
@@ -285,10 +318,7 @@ def _parameterize(
             sites, typemap, forcefield, speedup_by_moltag=speedup_by_moltag
         )
         self._parameterize_connections(
-            top,
-            forcefield,
-            label_type,
-            label,
+            top, forcefield, label_type, label, connection_index=connection_index
         )
         if forcefield.virtual_types:
             self._parameterize_virtual_sites(top, sites, bonds, forcefield)
@@ -380,6 +410,10 @@ def run_parameterization(self):
                 )
 
             assert_no_boundary_bonds(self.topology)
+            is_group = self.config.match_ff_by == "group"
+            connection_index = build_molecule_connection_index(
+                self.topology, is_group=is_group
+            )
             for label in labels:
                 if label not in self.forcefields:
                     logger.warning(
@@ -395,12 +429,14 @@ def run_parameterization(self):
                         self.config.speedup_by_moltag,
                         self.config.speedup_by_molgraph,
                     )
+
                     self._parameterize(
-                        self.topology,
-                        typemap,
+                        top=self.topology,
+                        typemap=typemap,
                         label_type=self.config.match_ff_by,
                         label=label,
                         speedup_by_moltag=self.config.speedup_by_moltag,  # This will be removed from the future iterations
+                        connection_index=connection_index,
                     )
         else:
             typemap = self._get_atomtypes(
@@ -439,12 +475,7 @@ def connection_identifier(
     ):  # This can extended to incorporate a pluggable object from the forcefield.
         """Return the group and list of identifiers for a connection to query the forcefield for its potential."""
         group = POTENTIAL_GROUPS[type(connection)]
-        return (
-            group,
-            [
-                *connection.get_connection_identifiers(),  # the viable keys made up of the bond orders
-            ],
-        )
+        return group, connection.get_connection_identifiers()
 
     @staticmethod
     def virtual_site_identifier(