S02: Graph Homomorphisms in Reaction Informatics

import warnings

import rdkit
import networkx as nx
import pandas as pd
from pathlib import Path
from synedu.Utils import draw_molecular_graph, mol_to_graph

# Matplotlib warns about layouts we set deliberately; not a lesson result.
warnings.filterwarnings("ignore", category=UserWarning)

print("RDKit version:", rdkit.__version__)
print("NetworkX version:", nx.__version__)

DATA_DIR = Path("data")
CSV_PATH = DATA_DIR / "molecules.csv"
df = pd.read_csv(CSV_PATH)
display(df.head())

from rdkit import Chem
from networkx.algorithms import isomorphism as iso

pairs = {
    "benzene": ("c1ccccc1", "C1=CC=CC=C1"),
    "aniline": ("c1ccccc1N", "c1ccccc1[NH3+]"),
}

graphs = {}
for name, (sa, sb) in pairs.items():
    graphs[f"{name}_a"] = mol_to_graph(Chem.MolFromSmiles(sa))
    graphs[f"{name}_b"] = mol_to_graph(Chem.MolFromSmiles(sb))


def node_match(n1, n2):
    return n1.get("element") == n2.get("element")


def edge_match(e1, e2):
    return (
        abs(float(e1.get("order", 1.0)) - float(e2.get("order", 1.0))) < 1e-9
        and bool(e1.get("aromatic", False)) == bool(e2.get("aromatic", False))
    )


def iso_and_count(G1, G2, nm, em):
    gm = iso.GraphMatcher(G1, G2, node_match=nm, edge_match=em)
    return gm.is_isomorphic(), sum(1 for _ in gm.isomorphisms_iter())


print("=== simple matcher (element + bond labels) ===")
for name in pairs:
    G1 = graphs[f"{name}_a"]
    G2 = graphs[f"{name}_b"]
    iso_flag, n_maps = iso_and_count(G1, G2, node_match, edge_match)
    print(f"{name:8} | isomorphic: {int(iso_flag):1d} | mappings: {n_maps}")

def node_match(n1, n2):
    keys = ("element", "formal_charge", "aromatic", "hcount")
    return all(n1.get(key) == n2.get(key) for key in keys)


def edge_match(e1, e2):
    return (
        abs(float(e1.get("order", 1.0)) - float(e2.get("order", 1.0))) < 1e-9
        and bool(e1.get("aromatic", False)) == bool(e2.get("aromatic", False))
    )


def enumerate_automorphisms(G: nx.Graph):
    GM_self = iso.GraphMatcher(G, G, node_match=node_match, edge_match=edge_match)
    return list(GM_self.isomorphisms_iter())


smiles = 'c1ccccc1'
mol = Chem.MolFromSmiles(smiles)
graph = mol_to_graph(mol)
draw_molecular_graph(graph, show_indices=True)
enumerate_automorphisms(graph)

import networkx as nx
from typing import Dict, Iterable, List, Set, Hashable


def compute_orbits_from_automorphisms(
    G: nx.Graph,
    automorphisms: Iterable[Dict[Hashable, Hashable]] | None = None,
) -> List[Set[Hashable]]:
    """
    Compute vertex orbits induced by graph automorphisms.

    Vertices belong to the same orbit if an automorphism maps one to the
    other. In molecular graphs, each orbit represents symmetry-equivalent
    atoms.

    :param G:
        Input graph.
    :type G: nx.Graph
    :param automorphisms:
        Iterable of automorphisms ``{v: φ(v)}``. If ``None``,
        automorphisms are computed internally.
    :type automorphisms: iterable of dict, optional
    :returns:
        List of vertex orbits (sets of nodes), deterministically ordered.
    :rtype: list[set]
    """
    if automorphisms is None:
        automorphisms = enumerate_automorphisms(G)

    # --- Union–find (disjoint set) ---
    parent: Dict[Hashable, Hashable] = {v: v for v in G.nodes()}

    def find(x: Hashable) -> Hashable:
        """Find set representative with path compression."""
        if parent[x] != x:
            parent[x] = find(parent[x])
        return parent[x]

    def union(a: Hashable, b: Hashable) -> None:
        """Merge the sets containing a and b."""
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[rb] = ra

    # --- Apply automorphism action ---
    for auto in automorphisms:
        for v, fv in auto.items():
            union(v, fv)

    # --- Collect orbits ---
    orbits: Dict[Hashable, Set[Hashable]] = {}
    for v in G.nodes():
        r = find(v)
        orbits.setdefault(r, set()).add(v)

    return sorted(orbits.values(), key=lambda s: min(s))

autos = enumerate_automorphisms(graph)
orbits = compute_orbits_from_automorphisms(graph, autos)

print("Number of automorphisms (benzene):", len(autos))
print("Orbits:", orbits)

from __future__ import annotations
from typing import Dict, List
import networkx as nx
from networkx.algorithms import isomorphism as iso
from rdkit import Chem


def substructure_node_match(host_attrs, pattern_attrs):
    """Match the query labels used for an aromatic ring embedding."""
    return (
        host_attrs.get("element") == pattern_attrs.get("element")
        and bool(host_attrs.get("aromatic", False))
        == bool(pattern_attrs.get("aromatic", False))
    )


def substructure_edge_match(host_attrs, pattern_attrs):
    return (
        abs(
            float(host_attrs.get("order", 1.0))
            - float(pattern_attrs.get("order", 1.0))
        )
        < 1e-9
        and bool(host_attrs.get("aromatic", False))
        == bool(pattern_attrs.get("aromatic", False))
    )


def nx_subgraph_matches(
    host_G: nx.Graph,
    pattern_G: nx.Graph,
    *,
    invert: bool = True,
    node_match_fn=substructure_node_match,
    edge_match_fn=substructure_edge_match,
) -> List[Dict]:
    """
    Enumerate subgraph isomorphisms of `pattern_G` inside `host_G`.

    Notes
    -----
    NetworkX's GraphMatcher(host, pattern).subgraph_monomorphisms_iter()
    yields mappings of the form: host_node -> pattern_node (G1 -> G2).

    If you want the more intuitive direction (pattern -> host), set
    `invert=True`. The default query matcher intentionally omits hydrogen
    count: a fused aromatic carbon can satisfy a ring query even though its
    hydrogen count differs from the corresponding atom in isolated benzene.
    """
    GM = iso.GraphMatcher(
        host_G,
        pattern_G,
        node_match=node_match_fn,
        edge_match=edge_match_fn,
    )

    out: List[Dict] = []
    for m_host_to_pat in GM.subgraph_monomorphisms_iter():
        if invert:
            out.append({p: h for h, p in m_host_to_pat.items()})
        else:
            out.append(m_host_to_pat)
    return out


pattern_mol = Chem.MolFromSmiles("c1ccccc1")  # benzene
host_mol = Chem.MolFromSmiles("c1ccc2ccccc2c1")  # naphthalene

pattern_G = mol_to_graph(pattern_mol)
host_G = mol_to_graph(host_mol)

matches = nx_subgraph_matches(host_G, pattern_G, invert=True)  # pattern_idx -> host_idx

print("Raw subgraph isomorphisms (pattern -> host):", len(matches))

from collections import defaultdict
from typing import Dict, Iterable, List, Mapping, Tuple, Union, Optional


def _build_node_to_rep(
    orbits: Union[Iterable[Iterable[int]], Mapping[int, int], None],
) -> Dict[int, int]:
    """Convert `orbits` into a mapping node -> representative."""
    if orbits is None:
        return {}
    if isinstance(orbits, Mapping):
        return dict(orbits)
    node_to_rep: Dict[int, int] = {}
    for orbit in orbits:
        orbit = set(orbit)
        if not orbit:
            continue
        rep = min(orbit)
        for v in orbit:
            node_to_rep[v] = rep
    return node_to_rep


def dedup_by_host_image_with_orbits(
    matches: Iterable[Dict[int, int]],
    *,
    orbits: Union[Iterable[Iterable[int]], Mapping[int, int], None] = None,
    host_autos: Optional[Iterable[Mapping[int, int]]] = None,
    pattern_node_order: Optional[Iterable[int]] = None,
    return_groups: bool = False,
) -> Union[List[Dict[int, int]], Dict[Tuple[int, ...], List[Dict[int, int]]]]:
    """
    Deduplicate pattern->host mappings by canonical host-image (with optional orbit reps).

    Parameters
    ----------
    matches : iterable of dict
        Each mapping is pattern_node -> host_node.
    host_autos : iterable of host-node mappings, optional
        Automorphisms used to canonicalize the entire mapped host image.
    orbits : deprecated
        Vertex orbits alone cannot canonicalize multi-node host images. They
        are accepted only together with ``host_autos`` for API compatibility.
    pattern_node_order : iterable of pattern node ids, optional
        Order used to form the per-mapping tuple for selecting the representative.
        If None, inferred deterministically from the first mapping (sorted keys).
    return_groups : bool (default False)
        If False (default) return a list of representative mappings (one per canonical class).
        If True return the full dict: canonical_key -> list[mappings].

    Returns
    -------
    List[dict] or Dict[tuple, list]
        See `return_groups` description.
    """
    matches = list(matches)  # materialize to allow multiple passes
    if not matches:
        return {} if return_groups else []

    if orbits is not None and host_autos is None:
        raise ValueError(
            "Vertex orbits alone cannot canonicalize a multi-node host image; "
            "pass host_autos"
        )
    autos = [dict(auto) for auto in host_autos] if host_autos is not None else []

    # infer or validate pattern node order
    if pattern_node_order is None:
        # Use sorted keys of the first mapping (deterministic)
        pattern_node_order = tuple(sorted(matches[0].keys()))
    else:
        pattern_node_order = tuple(pattern_node_order)

    def canonical_host_image(m: Dict[int, int]) -> Tuple[int, ...]:
        image = tuple(sorted(m.values()))
        if not autos:
            return image
        return min(tuple(sorted(auto[node] for node in image)) for auto in autos)

    # Bucket by the host image, canonicalized under one whole-graph symmetry.
    buckets: Dict[Tuple[int, ...], List[Dict[int, int]]] = defaultdict(list)
    for m in matches:
        key = canonical_host_image(m)
        buckets[key].append(m)

    # if user requested full groups, return them (deterministic ordering by key)
    if return_groups:
        return {k: buckets[k] for k in sorted(buckets.keys())}

    # otherwise pick a deterministic representative per bucket:
    # representative = mapping with smallest tuple (m[p] for p in pattern_node_order)
    representatives: List[Dict[int, int]] = []
    for key in sorted(buckets.keys()):
        group = buckets[key]

        # compute lexicographic tuple for each mapping according to pattern_node_order
        def ordering_tuple(m: Dict[int, int]) -> Tuple[int, ...]:
            return tuple(m[p] for p in pattern_node_order)

        rep = min(group, key=ordering_tuple)
        representatives.append(rep)

    return representatives

# 1. compute host orbits (optional - useful if you later want to collapse host symmetry)
host_autos = enumerate_automorphisms(host_G)  # φ : host_node -> host_node
host_orbits = compute_orbits_from_automorphisms(host_G, host_autos)  # list of sets
print("Host orbits:", host_orbits)
print()

# 2. dedup by host-image (default: returns representatives list)
representatives = dedup_by_host_image_with_orbits(matches)  # list of rep mappings
print("Number of representatives:", len(representatives))
for i, rep in enumerate(representatives, start=1):
    print(f" representative #{i}:", rep)
print()

# 3. (optional) show full groups and multiplicities to verify symmetry inflation
groups = dedup_by_host_image_with_orbits(
    matches, return_groups=True
)  # dict: key -> list[mappings]
print("Distinct placements by host-image (groups):", len(groups))
for key, group in groups.items():
    print(" placement key:", key, " multiplicity:", len(group))
    print("  example mapping:", group[0])
print()

# 4. (optional) collapse placements modulo host automorphisms (use host_orbits)
groups_mod_host = dedup_by_host_image_with_orbits(
    matches, orbits=host_orbits, host_autos=host_autos, return_groups=True
)
print("Unique classes modulo host automorphisms:", len(groups_mod_host))
for key, group in groups_mod_host.items():
    print(" canonical host-image key:", key, " multiplicity:", len(group))
    print("  example mapping:", group[0])

from rdkit import Chem
from typing import List, Dict, Tuple


def rdkit_subgraph_matches(
    host_mol: Chem.Mol,
    pattern_mol: Chem.Mol,
    *,
    uniquify: bool = False,
) -> List[Dict[int, int]]:
    """
    Enumerate substructure matches of `pattern_mol` inside `host_mol`.

    Returns
    -------
    List[Dict[int, int]]
        Each mapping is: pattern_atom_idx -> host_atom_idx

    Notes
    -----
    - RDKit returns tuples of host indices ordered by pattern atom order.
    - If `uniquify=False`, symmetry-equivalent embeddings are *not* removed.
    """
    matches: Tuple[Tuple[int, ...], ...] = host_mol.GetSubstructMatches(
        pattern_mol,
        uniquify=uniquify,
    )

    out: List[Dict[int, int]] = []
    for m in matches:
        out.append({p_idx: h_idx for p_idx, h_idx in enumerate(m)})
    return out


# Compare raw vs heuristic-deduplicated matches
matches = rdkit_subgraph_matches(host_mol, pattern_mol, uniquify=False)
print("Raw subgraph isomorphisms (pattern -> host):", len(matches))

matches = rdkit_subgraph_matches(host_mol, pattern_mol, uniquify=True)
print("Heuristically unique subgraph isomorphisms (pattern -> host):", len(matches))

from rdkit import Chem
import pandas as pd

_pattern_smi = "CCO"  # ethanol
_pattern_mol_rdkit = Chem.MolFromSmiles(_pattern_smi)
_pattern_G_nx = mol_to_graph(_pattern_mol_rdkit)

_test_cases = [
    ("naphthalene", "c1ccc2ccccc2c1"),
    ("aspirin", "CC(=O)Oc1ccccc1C(=O)O"),
    ("diethyl ether", "CCOCC"),
    ("ethanol", "CCO"),
    ("propanol", "CCCO"),
    ("glycol", "OCCO"),
]

rows = []
for name, smi in _test_cases:
    hmol = Chem.MolFromSmiles(smi)
    hG = mol_to_graph(hmol)

    nx_raw = nx_subgraph_matches(hG, _pattern_G_nx, invert=True)
    nx_deduped = dedup_by_host_image_with_orbits(nx_raw)

    rdkit_raw = rdkit_subgraph_matches(hmol, _pattern_mol_rdkit, uniquify=False)
    rdkit_uniq = rdkit_subgraph_matches(hmol, _pattern_mol_rdkit, uniquify=True)

    rows.append(
        {
            "molecule": name,
            "NX raw": len(nx_raw),
            "NX deduplicated": len(nx_deduped),
            "RDKit raw": len(rdkit_raw),
            "RDKit uniquify": len(rdkit_uniq),
            "NX==RDKit unique": len(nx_deduped) == len(rdkit_uniq),
        }
    )

df_cmp = pd.DataFrame(rows).set_index("molecule")
display(
    df_cmp.style.map(
        lambda v: (
            "color:green;font-weight:bold"
            if v is True
            else "color:red;font-weight:bold" if v is False else ""
        ),
        subset=["NX==RDKit unique"],
    ).set_caption("Ethanol subgraph - NX vs RDKit match counts")
)

# Diagnostic: aspirin differs because of aromaticity in the node labels.

_aspirin = Chem.MolFromSmiles("CC(=O)Oc1ccccc1C(=O)O")
_ethanol = Chem.MolFromSmiles("CCO")
_aspirin_G = mol_to_graph(_aspirin)
_ethanol_G = mol_to_graph(_ethanol)


def _element_charge_match(a, b):
    return a.get("element") == b.get("element") and int(
        a.get("formal_charge", 0)
    ) == int(b.get("formal_charge", 0))


def _bond_order_match(a, b):
    return float(a.get("order", 1.0)) == float(b.get("order", 1.0))


_nx_strict = nx_subgraph_matches(_aspirin_G, _ethanol_G, invert=True)
_nx_relaxed = nx_subgraph_matches(
    _aspirin_G,
    _ethanol_G,
    invert=True,
    node_match_fn=_element_charge_match,
    edge_match_fn=_bond_order_match,
)
_rdkit_unique = rdkit_subgraph_matches(_aspirin, _ethanol, uniquify=True)

print("NX strict:", _nx_strict)
print("NX relaxed:", _nx_relaxed)
print("RDKit uniquify:", _rdkit_unique)

# rdkit_subgraph_matches returns 0-based RDKit atom indices; graph nodes use
# a different ID scheme - bridge with the rdkit_idx attribute stored on each node
_eth_rdkit_to_node = {d["rdkit_idx"]: n for n, d in _ethanol_G.nodes(data=True)}
_asp_rdkit_to_node = {d["rdkit_idx"]: n for n, d in _aspirin_G.nodes(data=True)}

for mapping in _rdkit_unique:
    print("\nRDKit match:", mapping)
    for p_rdkit, h_rdkit in mapping.items():
        p_node = _eth_rdkit_to_node[p_rdkit]
        h_node = _asp_rdkit_to_node[h_rdkit]
        print(
            f"  pattern atom {p_rdkit} {_ethanol_G.nodes[p_node]} -> host atom {h_rdkit} {_aspirin_G.nodes[h_node]}"
        )