From 5f602dfb3d1c8393a40d61fbe828b9d4896f0ab3 Mon Sep 17 00:00:00 2001
From: avcopan <avcopan@gmail.com>
Date: Thu, 5 Mar 2026 13:17:10 -0500
Subject: [PATCH] Classify a reaction directly from AMChIs

---
 src/automol/graph/__init__.py      |  1 +
 src/automol/graph/base/__init__.py |  1 +
 src/automol/reac/enum.py           | 40 +++++++++++++++++++++++++++++-
 src/automol/smarts/_core.py        |  8 +++---
 4 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/src/automol/graph/__init__.py b/src/automol/graph/__init__.py
index 2223c9a8..c82469f0 100644
--- a/src/automol/graph/__init__.py
+++ b/src/automol/graph/__init__.py
@@ -56,6 +56,7 @@
 # # isomorphisms and equivalence
 # submodules:
 from .base import enum, ts, vmat
+from .base import ReactionSmarts
 
 # # getters
 # # setters
diff --git a/src/automol/graph/base/__init__.py b/src/automol/graph/base/__init__.py
index 75672f3f..6871ef9a 100644
--- a/src/automol/graph/base/__init__.py
+++ b/src/automol/graph/base/__init__.py
@@ -28,6 +28,7 @@
 # # constructors
 # submodules:
 from . import enum, ts, vmat
+from .enum import ReactionSmarts
 
 # # getters
 # # setters
diff --git a/src/automol/reac/enum.py b/src/automol/reac/enum.py
index 4cf8b793..20e5275c 100644
--- a/src/automol/reac/enum.py
+++ b/src/automol/reac/enum.py
@@ -1,6 +1,6 @@
 """Functions for enumerating reactions."""
 
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence, Collection
 
 from .. import amchi, graph, smiles
 from .. import smarts as smarts_
@@ -67,6 +67,44 @@ def from_graphs(
     return rxns
 
 
+def classify_from_amchis(
+    class_smarts: dict[str, str], rct_chis: Sequence[str], prd_chis: Collection[str]
+) -> str | None:
+    """Classify a reaction by matching it to a set of SMARTS templates.
+
+    :param class_smarts: Dictionary mapping class names to SMARTS strings
+    :param rct_chis: Reactant ChIs
+    :param prd_chis: Product ChIs
+    :return: Class name or None if no match is found
+    """
+    if any(map(amchi.has_stereo, rct_chis)) or any(map(amchi.has_stereo, prd_chis)):
+        msg = f"Cannot match SMARTS for AMChIs with stereo:\n{rct_chis} = {prd_chis}"
+        raise ValueError(msg)
+
+    comp_prd_chis = set(prd_chis)
+
+    # Prepare reactants graph
+    rct_gras = [graph.explicit(amchi.graph(c)) for c in rct_chis]
+    rct_gras, _ = graph.standard_keys_for_sequence(rct_gras)
+    rcts_gra = graph.union_from_sequence(rct_gras)
+
+    for class_, smarts in class_smarts.items():
+        nrcts, nprds = smarts_.shape(smarts)
+        if len(rct_chis) != nrcts or len(prd_chis) != nprds:
+            continue
+
+        # Enumerate reactions
+        ts_gras = graph.enum.reactions(smarts, rcts_gra)
+        for ts_gra in ts_gras:
+            prds_gra = graph.ts.products_graph(ts_gra, stereo=False, dummy=False)
+            prd_gras = graph.connected_components(prds_gra)
+            prd_chis = set(map(graph.amchi, prd_gras))
+            if prd_chis == comp_prd_chis:
+                return class_
+
+    return None
+
+
 # helpers
 def reactants(
     smarts: str,
diff --git a/src/automol/smarts/_core.py b/src/automol/smarts/_core.py
index 1e4508af..4cfe7e18 100644
--- a/src/automol/smarts/_core.py
+++ b/src/automol/smarts/_core.py
@@ -13,19 +13,19 @@ def shape(smarts: str) -> tuple[list[int], list[int]]:
     return rd.shape(rd.from_smarts(smarts))
 
 
-def reactant_count(smarts: str) -> tuple[list[int], list[int]]:
+def reactant_count(smarts: str) -> int:
     """Get number of reactants in SMARTS string.
 
     :param smarts: SMARTS string
-    :return: Reaction shape
+    :return: Reactant count
     """
     return rd.reactant_count(rd.from_smarts(smarts))
 
 
-def product_count(smarts: str) -> tuple[list[int], list[int]]:
+def product_count(smarts: str) -> int:
     """Get number of products in SMARTS string.
 
     :param smarts: SMARTS string
-    :return: Reaction shape
+    :return: Product count
     """
     return rd.product_count(rd.from_smarts(smarts))