Several more atom mapping improvements (#839)

This PR improves several aspects of atom mappings.
two key fixes:
1) Stopping conformer generation in `ARCSpecies._scissors`. This is a
major issue that improves both accuracy and performence:
    1) Performence - less RDKit calls
2) Accuracy - Generating conformers distorts the XYZ of the scissored
products.
Moreover, this is critical for cyclic species, since debugging showed we
did not include the original XYZ in the ARCSpecies, which caused errors.
2) the `identify_superimposable_candidates`, which was being called
$\mathcal{O}(N^2)$ times and reduces to $\mathcal{O}(N)$. More calls are
not required, and does not increase accuracy. (There is some more math
behind it, and it can be discussed offline @alongd).

Other fixes added, and more generalization of AM testing (manually
verified)

Author: kfir4444

arc/common.py

@@ -36,6 +36,12 @@
 logger = logging.getLogger('arc')
 logging.getLogger('matplotlib.font_manager').disabled = True
 
+try:
+    from rdkit import RDLogger
+    RDLogger.DisableLog('rdApp.*')
+except ImportError:
+    pass
+
 # Absolute path to the ARC folder.
 ARC_PATH = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
 ARC_TESTING_PATH = os.path.join(ARC_PATH, 'arc', 'testing')

arc/family/family.py

@@ -24,6 +24,37 @@
 logger = get_logger()
 
 
+REACTION_FAMILY_CACHE: dict[tuple[str, bool], 'ReactionFamily'] = {}
+
+# Pre-compiled regex patterns
+ENTRY_PATTERN = re.compile(r'entry\((.*?)\)', re.DOTALL)
+LABEL_PATTERN = re.compile(r'label\s*=\s*(["\'])(.*?)\1|label\s*=\s*(\w+)')
+GROUP_PATTERN = re.compile(r'group\s*=\s*(?:("""(.*?)"""|"(.*?)"|\'(.*?)\')|(OR\{.*?\}))', re.DOTALL)
+REVERSIBLE_PATTERN = re.compile(r'reversible\s*=\s*(True|False)')
+OWN_REVERSE_PATTERN = re.compile(r'ownReverse\s*=\s*(True|False)')
+RECIPE_PATTERN = re.compile(r'recipe\((.*?)\)', re.DOTALL)
+REACTANTS_PATTERN = re.compile(r'reactants\s*=\s*\[(.*?)\]', re.DOTALL)
+PRODUCTS_PATTERN = re.compile(r'products\s*=\s*\[(.*?)\]', re.DOTALL)
+ACTIONS_PATTERN = re.compile(r'actions\s*=\s*\[(.*?)\]', re.DOTALL)
+
+
+def get_reaction_family(label: str, consider_arc_families: bool = True) -> 'ReactionFamily':
+    """
+    A helper function for getting a cached ReactionFamily object.
+
+    Args:
+        label (str): The reaction family label.
+        consider_arc_families (bool, optional): Whether to consider ARC's custom families.
+
+    Returns:
+        ReactionFamily: The ReactionFamily object.
+    """
+    key = (label, consider_arc_families)
+    if key not in REACTION_FAMILY_CACHE:
+        REACTION_FAMILY_CACHE[key] = ReactionFamily(label=label, consider_arc_families=consider_arc_families)
+    return REACTION_FAMILY_CACHE[key]
+
+
 def get_rmg_db_subpath(*parts: str, must_exist: bool = False) -> str:
     """Return a path under the RMG database, handling both source and packaged layouts."""
     if RMG_DB_PATH is None:
@@ -108,7 +139,17 @@ def __init__(self,
         self.groups_as_lines = read_groups_file_lines(label, consider_arc_families)
         self.reversible = is_reversible(self.groups_as_lines)
         self.own_reverse = is_own_reverse(self.groups_as_lines)
-        self.reactants = get_reactant_groups_from_template(self.groups_as_lines)
+        
+        reactant_labels = get_initial_reactant_labels_from_template(self.groups_as_lines)
+        all_necessary_entries = get_entries(self.groups_as_lines, entry_labels=reactant_labels, recursive=True)
+        self.reactants = get_reactant_groups_from_template(self.groups_as_lines, entries=all_necessary_entries)
+        self.entries = all_necessary_entries
+        
+        self.groups = {}
+        for reactant_group in self.reactants:
+            for label in reactant_group:
+                if label not in self.groups and label in self.entries:
+                    self.groups[label] = Group().from_adjacency_list(self.entries[label])
         self.reactant_num = self.get_reactant_num()
         self.product_num = get_product_num(self.groups_as_lines)
         entry_labels = list()
@@ -156,6 +197,9 @@ def generate_products(self,
                 for group_label in group_labels:
                     group = self.groups_by_label[group_label]
                     for mol in reactant.mol_list or [reactant.mol]:
+                        if not any(a.atomtype for a in mol.atoms):
+                            # Update atomtypes if they are missing (e.g., from SMILES)
+                            mol.update_atomtypes(log_species=False, raise_exception=False)
                         splits = group.split()
                         if mol.is_subgraph_isomorphic(other=group, save_order=True) \
                                 or len(splits) > 1 and any(mol.is_subgraph_isomorphic(other=g, save_order=True) for g in splits):
@@ -297,9 +341,15 @@ def generate_bimolecular_products(self,
                         group_2 = self.groups_by_label[reactant_to_group_map_2['subgroup']]
                         isomorphic_subgraphs_1 = mol_1.find_subgraph_isomorphisms(other=group_1, save_order=True)
                         isomorphic_subgraphs_2 = mol_2.find_subgraph_isomorphisms(other=group_2, save_order=True)
+                        
                         if len(isomorphic_subgraphs_1) and len(isomorphic_subgraphs_2):
                             for isomorphic_subgraph_1 in isomorphic_subgraphs_1:
                                 for isomorphic_subgraph_2 in isomorphic_subgraphs_2:
+                                    # Create the combined isomorphic subgraph.
+                                    # Note: get_isomorphic_subgraph needs to know which subgraph corresponds to which template index.
+                                    # It assumes mol_1 corresponds to the first group match and mol_2 to the second.
+                                    # The labels are already inside the group_atom.label.
+                                    
                                     isomorphic_subgraph_dicts.append(
                                         {'mols': [mol_1, mol_2],
                                          'subgroups': (reactant_to_group_map_1['subgroup'],
@@ -422,7 +472,7 @@ def get_reactant_num(self) -> int:
                 if match:
                     return int(match.group(1))
         if len(self.reactants) == 1:
-            group = Group().from_adjacency_list(get_group_adjlist(self.groups_as_lines, entry_label=self.reactants[0][0]))
+            group = self.groups[self.reactants[0][0]]
             groups = group.split()
             return len(groups)
         else:
@@ -523,7 +573,7 @@ def determine_possible_reaction_products_from_family(rxn: ARCReaction,
                     and whether the family's template also represents its own reverse.
     """
     product_dicts = list()
-    family = ReactionFamily(label=family_label, consider_arc_families=consider_arc_families)
+    family = get_reaction_family(label=family_label, consider_arc_families=consider_arc_families)
     products = family.generate_products(reactants=rxn.get_reactants_and_products(return_copies=True)[0])
     if products:
         for group_labels, product_lists in products.items():
@@ -765,11 +815,10 @@ def is_reversible(groups_as_lines: list[str]) -> bool:
     Returns:
         bool: Whether the reaction family is reversible.
     """
-    for line in groups_as_lines:
-        if 'reversible = True' in line:
-            return True
-        if 'reversible = False' in line:
-            return False
+    groups_str = ''.join(groups_as_lines)
+    match = REVERSIBLE_PATTERN.search(groups_str)
+    if match:
+        return match.group(1) == 'True'
     return True
 
 
@@ -780,36 +829,41 @@ def is_own_reverse(groups_as_lines: list[str]) -> bool:
     Returns:
         bool: Whether the reaction family's template also represents its own reverse.
     """
-    for line in groups_as_lines:
-        if 'ownReverse=True' in line:
-            return True
-        if 'ownReverse=False' in line:
-            return False
+    groups_str = ''.join(groups_as_lines)
+    match = OWN_REVERSE_PATTERN.search(groups_str)
+    if match:
+        return match.group(1) == 'True'
     return False
 
 
-def get_reactant_groups_from_template(groups_as_lines: list[str]) -> list[list[str]]:
+def get_reactant_groups_from_template(groups_as_lines: list[str],
+                                      entries: dict[str, str] | None = None,
+                                      ) -> list[list[str]]:
     """
     Get the reactant groups from a template content string.
     Descends the entries if a group is defined as an OR complex,
     e.g.: group = "OR{Xtrirad_H, Xbirad_H, Xrad_H, X_H}"
 
     Args:
         groups_as_lines (list[str]): The template content string.
+        entries (dict[str, str], optional): Pre-extracted entries.
 
     Returns:
         list[list[str]]: The non-complex reactant groups.
     """
     reactant_labels = get_initial_reactant_labels_from_template(groups_as_lines)
+    if entries is None:
+        entries = get_entries(groups_as_lines, entry_labels=reactant_labels)
     result = list()
     for reactant_label in reactant_labels:
-        if 'OR{' not in get_group_adjlist(groups_as_lines, entry_label=reactant_label):
+        adj = get_group_adjlist(groups_as_lines, entry_label=reactant_label, entries=entries)
+        if 'OR{' not in adj:
             result.append([reactant_label])
         else:
             stack = [reactant_label]
-            while any('OR{' in get_group_adjlist(groups_as_lines, entry_label=label) for label in stack):
+            while any('OR{' in get_group_adjlist(groups_as_lines, entry_label=label, entries=entries) for label in stack):
                 label = stack.pop(0)
-                group_adjlist = get_group_adjlist(groups_as_lines, entry_label=label)
+                group_adjlist = get_group_adjlist(groups_as_lines, entry_label=label, entries=entries)
                 if 'OR{' not in group_adjlist:
                     stack.append(label)
                 else:
@@ -851,7 +905,7 @@ def descent_complex_group(group: str) -> list[str]:
         list[str]: The non-complex reactant group labels, e.g.: ['Xtrirad_H', 'Xbirad_H', 'Xrad_H', 'X_H'].
     """
     if group.startswith('OR{') and group.endswith('}'):
-        group = [g.strip() for g in group[3:-1].split(',')]
+        group = [c.strip() for c in group[3:-1].split(',')]
     if isinstance(group, str):
         group = [group]
     return group
@@ -871,13 +925,15 @@ def get_initial_reactant_labels_from_template(groups_as_lines: list[str],
     Returns:
         list[str]: The reactant groups.
     """
-    labels = list()
-    for line in groups_as_lines:
-        match = re.search(r'products=\[(.*?)\]', line) if products else re.search(r'reactants=\[(.*?)\]', line)
-        if match:
-            labels = match.group(1).replace('"', '').split(', ')
-            break
-    return labels
+    groups_str = ''.join(groups_as_lines)
+    pattern = PRODUCTS_PATTERN if products else REACTANTS_PATTERN
+    match = pattern.search(groups_str)
+    if match:
+        content = match.group(1)
+        # Use regex to find all quoted strings (with backreferences) or unquoted words
+        labels = re.findall(r'(["\'])(.*?)\1|(\w+)', content)
+        return [label[1] or label[2] for label in labels]
+    return list()
 
 
 def get_recipe_actions(groups_as_lines: list[str]) -> list[list[str]]:
@@ -982,43 +1038,77 @@ def split_entries(groups_str: str) -> list[str]:
 
 def get_entries(groups_as_lines: list[str],
                 entry_labels: list[str],
+                recursive: bool = False,
                 ) -> dict[str, str]:
     """
-    Get the requested entries grom a template content string.
+    Get the requested entries from a template content string.
 
     Args:
         groups_as_lines (list[str]): The template content string.
-        entry_labels (list[str]): The entry labels to extract.
+        entry_labels (list[str], optional): The entry labels to extract. If None, all entries are extracted.
+        recursive (bool, optional): Whether to recursively extract child entries for OR complexes.
 
     Returns:
         dict[str, str]: The extracted entries, keys are the labels, values are the groups.
     """
-    groups_str = ''.join(groups_as_lines)
-    entries = split_entries(groups_str)
-    specific_entries = dict()
-    for i, entry in enumerate(entries):
-        label_match = re.search(r'label\s*=\s*"(.*?)"', entry)
-        group_match = re.search(r'group\s*=(.*?)(?=\w+\s*=)', entry, re.DOTALL)
-        if label_match is not None and group_match is not None and label_match.group(1) in entry_labels:
-            specific_entries[label_match.group(1)] = clean_text(group_match.group(1))
-        if i > 2000:
-            break
-    return specific_entries
+    groups_str = "\n" + "".join(groups_as_lines)
+    # Split by `entry(` but keep the delimiter-ish part
+    parts = re.split(r"\nentry\s*\(", groups_str)
+    
+    temp_entries = {}
+    label_pat = re.compile(r"label\s*=\s*(?:([\"'])(.*?)\1|(\w+))")
+    group_pat = re.compile(r"group\s*=\s*(?:\"\"\"(.*?)\"\"\"|([\"'])(.*?)\2|(OR\{.*?\}))", re.DOTALL)
+
+    for part in parts[1:]: # Skip the header
+        label_match = label_pat.search(part)
+        group_match = group_pat.search(part)
+        if label_match and group_match:
+            label = label_match.group(2) or label_match.group(3)
+            # Extract the matched regex group (1 for triple quotes, 3 for single/double quotes, 4 for OR complex)
+            adj = group_match.group(1) or group_match.group(3) or group_match.group(4)
+            temp_entries[label] = clean_text(adj)
+
+    if entry_labels is None:
+        return temp_entries
+
+    all_entries = {}
+    to_process = list(entry_labels)
+    processed = set()
+    while to_process:
+        label = to_process.pop()
+        if label in processed or label not in temp_entries:
+            continue
+        processed.add(label)
+        adj = temp_entries[label]
+        if recursive and 'OR{' in adj:
+            # Match OR{label1, label2, ...}
+            or_match = re.search(r'OR\s*\{\s*(.*?)\s*\}', adj, re.DOTALL)
+            if or_match:
+                children_str = or_match.group(1)
+                children = [c.strip() for c in children_str.split(',')]
+                to_process.extend(children)
+        else:
+            all_entries[label] = adj
+    return all_entries
 
 
 def get_group_adjlist(groups_as_lines: list[str],
                       entry_label: str,
+                      entries: dict[str, str] | None = None,
                       ) -> str:
     """
     Get the corresponding group value for the given entry label.
 
     Args:
         groups_as_lines (list[str]): The template content string.
         entry_label (str): The entry label to extract.
+        entries (dict[str, str], optional): Pre-extracted entries.
 
     Returns:
         str: The extracted group.
     """
+    if entries is not None and entry_label in entries:
+        return entries[entry_label]
     specific_entries = get_entries(groups_as_lines, entry_labels=[entry_label])
     return specific_entries[entry_label]
 

arc/job/adapters/ts/heuristics_test.py

@@ -36,6 +36,10 @@
 from arc.species.species import ARCSpecies
 from arc.species.zmat import _compare_zmats, get_parameter_from_atom_indices
 
+from arc.species.species import check_isomorphism
+from arc.species.zmat import remove_zmat_atom_0
+from arc.species.converter import relocate_zmat_dummy_atoms_to_the_end
+
 
 class TestHeuristicsAdapter(unittest.TestCase):
     """
@@ -1409,11 +1413,32 @@ def test_get_new_zmat2_map(self):
         # expected_new_map = {0: 12, 1: 13, 2: 'X24', 3: 14, 4: 15, 5: 16, 6: 'X25', 7: 17, 8: 'X26', 9: 18, 10: 19,
         #                     11: 20, 12: 21, 13: 22, 14: 'X27', 15: 23, 16: 'X28', 17: 2, 18: 3, 19: 1, 21: 4, 23: 0,
         #                     25: 7, 26: 6, 28: 5, 20: 'X8', 22: 'X9', 24: 'X10', 27: 'X11'}
-        expected_new_map = {0: 12, 1: 13, 2: 'X24', 3: 14, 4: 15, 5: 16, 6: 'X25', 7: 17, 8: 'X26', 9: 18, 10: 19,
-                            11: 20, 12: 21, 13: 22, 14: 'X27', 15: 23, 16: 'X28', 17: 2, 18: 1, 19: 3, 21: 0, 23: 4,
-                            25: 5, 26: 6, 28: 7, 20: 'X8', 22: 'X9', 24: 'X10', 27: 'X11'}
-
-        self.assertEqual(new_map, expected_new_map)
+        
+        # Test isomorphism of the mapped reactant_2 part
+        zmat_2_mod = remove_zmat_atom_0(self.zmat_6)
+        zmat_2_mod['map'] = relocate_zmat_dummy_atoms_to_the_end(zmat_2_mod['map'])
+        spc_from_zmat_2 = ARCSpecies(label='spc_from_zmat_2', xyz=zmat_2_mod, multiplicity=reactant_2.multiplicity,
+                                     number_of_radicals=reactant_2.number_of_radicals, charge=reactant_2.charge)
+        
+        # Verify that all physical atom indices in new_map that came from zmat_2 correctly map to reactant_2
+        # Atom indices in new_map are for the combined species.
+        # Atoms 0-16 in self.zmat_5, atoms 1-12 in self.zmat_6 (13 atoms total, index 0 removed).
+        # In get_new_zmat_2_map, zmat_2 atoms are mapped to indices in new_map.
+        
+        num_atoms_1 = len(self.zmat_5['symbols'])
+        atom_map = dict()
+        for i in range(1, len(self.zmat_6['symbols'])):
+            if not isinstance(self.zmat_6['symbols'][i], str) or self.zmat_6['symbols'][i] != 'X':
+                # This is a physical atom in zmat_2 (at index i)
+                # Its index in the combined Z-Matrix is num_atoms_1 + i - 1
+                combined_idx = num_atoms_1 + i - 1
+                if combined_idx in new_map:
+                    # new_map[combined_idx] is the index in reactant_2
+                    # i-1 is the index in spc_from_zmat_2
+                    atom_map[i-1] = new_map[combined_idx]
+
+        # Verify the atom_map is a valid isomorphism
+        self.assertTrue(check_isomorphism(spc_from_zmat_2.mol, reactant_2.mol, atom_map))
 
     def test_get_new_map_based_on_zmat_1(self):
         """Test the get_new_map_based_on_zmat_1() function."""

arc/job/adapters/ts/linear_test.py

@@ -247,9 +247,9 @@ def _make_rxn_2() -> ARCReaction:
                                                     H       0.27058353   -0.73979548    1.43184405""")])
 
 
-class TestHeuristicsAdapter(unittest.TestCase):
+class TestLinearAdapter(unittest.TestCase):
     """
-    Contains unit tests for the HeuristicsAdapter class.
+    Contains unit tests for the LinearAdapter class.
     """
 
     @classmethod
@@ -1133,7 +1133,7 @@ def test_interpolate_1_plus_2_cycloaddition(self):
             self.assertEqual(len(ts_xyz['symbols']), 10)
             self.assertFalse(colliding_atoms(ts_xyz),
                              msg=f'Collision in 1+2_Cycloaddition TS:\n{xyz_to_str(ts_xyz)}')
-        expected_ts = """C       1.59999925   -0.11618654   -0.14166302
+        expected_ts_1 = """C       1.59999925   -0.11618654   -0.14166302
 C       0.29517860   -0.02143486   -0.02613492
 C      -1.15821797   -1.12490772    0.14486040
 C      -0.81238032    0.84414025    0.04444949
@@ -1143,7 +1143,18 @@ def test_interpolate_1_plus_2_cycloaddition(self):
 H      -1.52801447   -1.64655150   -0.72678867
 H      -0.94547237    1.40230195    0.96062403
 H      -1.11212744    1.33826544   -0.86912905"""
-        self.assertTrue(any(almost_equal_coords(ts, str_to_xyz(expected_ts)) for ts in ts_xyzs))
+        expected_ts_2 = """C       1.59999925   -0.11618654   -0.14166302
+C       0.29517860   -0.02143486   -0.02613492
+C      -0.92013120   -0.71833111    0.10894610
+C      -0.99229728    1.28107087    0.04554500
+H       2.21797993    0.77036923   -0.22897655
+H       2.09015362   -1.08321135   -0.15246324
+H      -1.12327237   -1.17593811    1.06705013
+H      -1.28992770   -1.23997489   -0.76270297
+H      -1.12538933    1.83923257    0.96171954
+H      -1.29204440    1.77519606   -0.86803354"""
+        self.assertTrue(any(almost_equal_coords(ts, str_to_xyz(expected_ts_1)) for ts in ts_xyzs) or
+                        any(almost_equal_coords(ts, str_to_xyz(expected_ts_2)) for ts in ts_xyzs))
         # The TS should have extended forming bonds compared to the product.
         # Get atom map to find the carbene atom in the product.
         atom_map = map_rxn(rxn=rxn)