S09: Context Radius Expansion

%matplotlib inline
import rdkit
import pandas as pd
import networkx as nx
from pathlib import Path
import importlib.metadata as _m
from synedu.Utils import load_database
import logging


def _silence():
    logging.disable(logging.CRITICAL)


_silence()

print("RDKit    :", rdkit.__version__)
print("NetworkX :", nx.__version__)
print("SynKit   :", _m.version('synkit'))

DATA_PATH = Path("data") / "smart.json.gz"
data = pd.DataFrame(load_database(DATA_PATH)[:2000])
print(f"\nDataset : {data.shape[0]} reactions  ·  columns: {list(data.columns)}")
data.head(3)

from sklearn.model_selection import train_test_split

train_df, test_df = train_test_split(data, test_size=0.2, random_state=42)
print(f"Train : {len(train_df)}  |  Test : {len(test_df)}")

from synedu.Utils import cluster_its_graphs, rsmi_to_its, wl_hash
from synkit.Chem.Reaction.standardize import Standardize

train = train_df.to_dict("records")
for v in train:
    v["ITS"] = rsmi_to_its(v["smart"], core=False)  # full ITS (all atoms)
    v["RC"] = rsmi_to_its(v["smart"], core=True)  # reaction-center only
    v["wl"] = wl_hash(v["RC"])

_class_ids = cluster_its_graphs([v["RC"] for v in train])
for v, class_id in zip(train, _class_ids):
    v["class"] = class_id
cluster = train

_df_cls = pd.DataFrame(cluster)
print("Rule family sizes (top 8):")
print(_df_cls["class"].value_counts().head(8).to_string())

_std = Standardize()
test = test_df.to_dict("records")
for v in test:
    _rsmi = _std.fit(v["smart"], remove_aam=True)
    v["r"], v["p"] = _rsmi.split(">>")
    v["rsmi"] = _rsmi
    v["ITS"] = rsmi_to_its(v["smart"], core=False)
    v["RC"] = rsmi_to_its(v["smart"], core=True)
    v["wl"] = wl_hash(v["RC"])

print(
    f"Test set standardised and assigned to reaction-center families: {len(test)} entries"
)
print("Example — reactant :", test[0]["r"][:60], "...")

def reaction_center_nodes(its: nx.Graph) -> set:
    """Return atoms incident to any bond that changes between reactant and product."""
    rc = set()
    for u, v, d in its.edges(data=True):
        br, bp = d.get("order", (0.0, 0.0))
        if abs(float(br) - float(bp)) > 1e-6:
            rc.update([u, v])
    return rc


# Quick check on the first training reaction
_its_ex = cluster[0]["ITS"]
_rc_ex = reaction_center_nodes(_its_ex)
print(
    f"ITS has {_its_ex.number_of_nodes()} nodes, " f"{_its_ex.number_of_edges()} edges"
)
print(f"Reaction-center nodes: {sorted(_rc_ex)}")

def expand_context(its: nx.Graph, r: int) -> nx.Graph:
    """Return the ITS subgraph within r hops of the reaction center."""
    seed = reaction_center_nodes(its)
    frontier = set(seed)
    for _ in range(r):
        next_shell = set()
        for v in frontier:
            next_shell.update(its.neighbors(v))
        frontier = next_shell - seed  # only brand-new nodes
        seed |= frontier
    return its.subgraph(seed).copy()


# At r=0 the node set is exactly the reaction-center endpoints. Because this
# function returns a vertex-induced subgraph, unchanged edges between those
# endpoints may also be present.
_ctx0 = expand_context(_its_ex, 0)
_ctx1 = expand_context(_its_ex, 1)
_ctx2 = expand_context(_its_ex, 2)
print(f"r=0 → {_ctx0.number_of_nodes()} nodes")
print(f"r=1 → {_ctx1.number_of_nodes()} nodes")
print(f"r=2 → {_ctx2.number_of_nodes()} nodes")
print(f"Full ITS → {_its_ex.number_of_nodes()} nodes")

from synedu.Utils.rxn_vis import visualize_reaction
from IPython.display import SVG, display

_samples = family_train[:3]
for i, v in enumerate(_samples):
    display(
        SVG(
            visualize_reaction(
                v["smart"],
                svg=True,
                legend=f"Family {TOP_CLASS} — example {i+1}",
                highlight_changes=True,
            )
        )
    )

from collections import Counter

MAX_RADIUS = 5


def make_radius_templates(records, max_radius=MAX_RADIUS):
    """Build deduplicated r-hop context templates for one reaction family."""
    templates = []
    for r in range(max_radius + 1):
        candidates = []
        for idx, record in enumerate(records):
            K = expand_context(record["ITS"], r)
            h = wl_hash(K)
            candidates.append({"K": K, "hash": h, "source_idx": idx})

        class_ids = cluster_its_graphs([candidate["K"] for candidate in candidates])
        class_frequency = Counter(class_ids)
        seen_classes = set()
        level = []
        ranked = sorted(
            zip(candidates, class_ids),
            key=lambda item: (-class_frequency[item[1]], item[1]),
        )
        for candidate, class_id in ranked:
            if class_id in seen_classes:
                continue
            seen_classes.add(class_id)
            candidate["frequency"] = class_frequency[class_id]
            level.append(candidate)
        templates.append(level)
    return templates


_templates = make_radius_templates(family_train, max_radius=MAX_RADIUS)

for r, tlist in enumerate(_templates):
    print(f"r = {r}: {len(tlist):3d} unique templates")

from synkit.Synthesis.Reactor.syn_reactor import SynReactor

_std2 = Standardize()


def _unique_preserve_order(values):
    seen = set()
    out = []
    for value in values:
        if value in seen:
            continue
        seen.add(value)
        out.append(value)
    return out


def _apply_templates(entries, template_graphs, *, n_entries=25):
    """Apply context-graph templates to up to n_entries test reactions."""
    subset = entries[:n_entries]
    for v in subset:
        fw = []
        for K in template_graphs:
            try:
                reactor = SynReactor(
                    substrate=v["r"],
                    template=K,
                    explicit_h=True,
                    implicit_temp=False,
                    automorphism=True,
                )
                fw.extend(reactor.smarts_list)
            except Exception:
                pass
        clean = []
        for s in fw:
            try:
                c = _std2.fit(s)
                if c:
                    clean.append(c)
            except Exception:
                pass
        v["fw_clean"] = _unique_preserve_order(clean)
    return subset


def f1_tradeoff(recall, enrichment):
    """Harmonic mean between recall and enrichment."""
    if recall + enrichment == 0:
        return 0.0
    return 2 * recall * enrichment / (recall + enrichment)


def recall_enrichment_at_k(entries, *, key_gt="rsmi", key_pred="fw_clean", k=50):
    """Compute standardized recall/enrichment using top-k or all candidates."""
    recalls = []
    enrichments = []

    for entry in entries:
        gt = entry.get(key_gt)
        all_preds = entry.get(key_pred, [])
        preds = all_preds if k is None else all_preds[:k]
        hit = int(gt in preds)
        recalls.append(hit)
        enrichments.append(hit / len(preds) if preds else 0.0)

    if not entries:
        return {"recall_at_k": 0.0, "enrichment_at_k": 0.0}

    return {
        "recall_at_k": sum(recalls) / len(recalls),
        "enrichment_at_k": sum(enrichments) / len(enrichments),
    }


print("Evaluation helpers ready.")

import copy
import time

_eval_results = []
_N_TEST = 25  # same-family test entries used per radius (keep notebook fast)
_EVAL_K = 50
_eval_source = family_test[:_N_TEST]

print(f"Evaluating {len(_eval_source)} test reactions from class {TOP_CLASS}.")

for r, tlist in enumerate(_templates):
    t0 = time.time()
    tmpl_graphs = [t["K"] for t in tlist]

    # Work on fresh copies so we do not accumulate fw_clean across iterations.
    _entries = copy.deepcopy(_eval_source)
    _entries = _apply_templates(_entries, tmpl_graphs, n_entries=len(_entries))

    elapsed = time.time() - t0
    m = recall_enrichment_at_k(_entries, k=_EVAL_K)
    score = f1_tradeoff(m["recall_at_k"], m["enrichment_at_k"])

    _eval_results.append(
        {
            "r": r,
            "n_templates": len(tlist),
            "elapsed": round(elapsed, 2),
            f"Recall@{_EVAL_K}": round(m["recall_at_k"], 3),
            f"Enrichment@{_EVAL_K}": round(m["enrichment_at_k"], 4),
            f"F1@{_EVAL_K}": round(score, 4),
            "tradeoff_score": round(score, 4),
        }
    )
    print(
        f"r={r}  templates={len(tlist):3d}  "
        f"Recall@{_EVAL_K}={_eval_results[-1][f'Recall@{_EVAL_K}']:.3f}  "
        f"Enrichment@{_EVAL_K}={_eval_results[-1][f'Enrichment@{_EVAL_K}']:.4f}  "
        f"F1@{_EVAL_K}={_eval_results[-1][f'F1@{_EVAL_K}']:.4f}  "
        f"({elapsed:.1f}s)"
    )

_eval_df = pd.DataFrame(_eval_results)
print("\nDone.")

# Three-class radius check: r = 0, 1, 2
_RADIUS_CHECK = [0, 1, 2]
_CHECK_K = None
_CHECK_LABEL = "all"

_class_rows = []
for class_id, train_count in _df_cls["class"].value_counts().items():
    class_train = [v for v in cluster if v["class"] == class_id]
    class_wl = class_train[0]["wl"]
    class_test = [v for v in test if v["wl"] == class_wl]
    if class_test:
        _class_rows.append(
            {
                "class": int(class_id),
                "wl": class_wl,
                "train_n": len(class_train),
                "test_n": len(class_test),
            }
        )
    if len(_class_rows) == 3:
        break

_check_classes = pd.DataFrame(_class_rows)
print("Classes selected for the practical radius check:")
print(_check_classes[["class", "train_n", "test_n"]].to_string(index=False))

_multi_results = []

for _, class_row in _check_classes.iterrows():
    class_id = int(class_row["class"])
    class_train = [v for v in cluster if v["class"] == class_id]
    class_test = [v for v in test if v["wl"] == class_row["wl"]]
    class_templates = make_radius_templates(class_train, max_radius=max(_RADIUS_CHECK))

    print(f"\nClass {class_id}: evaluating {len(class_test)} test reactions")
    for r in _RADIUS_CHECK:
        t0 = time.time()
        tmpl_graphs = [t["K"] for t in class_templates[r]]
        entries = copy.deepcopy(class_test)
        entries = _apply_templates(entries, tmpl_graphs, n_entries=len(entries))
        elapsed = time.time() - t0
        metrics = recall_enrichment_at_k(entries, k=_CHECK_K)
        f1 = f1_tradeoff(metrics["recall_at_k"], metrics["enrichment_at_k"])

        _multi_results.append(
            {
                "class": class_id,
                "r": r,
                "train_n": len(class_train),
                "test_n": len(class_test),
                "n_templates": len(class_templates[r]),
                f"Recall({_CHECK_LABEL})": metrics["recall_at_k"],
                f"Enrichment({_CHECK_LABEL})": metrics["enrichment_at_k"],
                f"F1({_CHECK_LABEL})": f1,
                "tradeoff_score": f1,
                "elapsed": elapsed,
            }
        )
        print(
            f"  r={r}: templates={len(class_templates[r]):3d}  "
            f"Recall({_CHECK_LABEL})={metrics['recall_at_k']:.3f}  "
            f"Enrichment({_CHECK_LABEL})={metrics['enrichment_at_k']:.4f}  "
            f"F1({_CHECK_LABEL})={f1:.4f}  ({elapsed:.1f}s)"
        )

_multi_df = pd.DataFrame(_multi_results)
_macro_df = _multi_df.groupby("r", as_index=False).agg(
    n_classes=("class", "nunique"),
    mean_templates=("n_templates", "mean"),
    recall_mean=(f"Recall({_CHECK_LABEL})", "mean"),
    enrichment_mean=(f"Enrichment({_CHECK_LABEL})", "mean"),
    f1_mean=(f"F1({_CHECK_LABEL})", "mean"),
    elapsed_total=("elapsed", "sum"),
)
_macro_df["tradeoff_score"] = _macro_df["f1_mean"]
_best_macro = _macro_df.loc[_macro_df["tradeoff_score"].idxmax()]

print("\nMacro-average across the three classes:")
print(
    _macro_df.assign(
        mean_templates=lambda d: d["mean_templates"].round(1),
        recall_mean=lambda d: d["recall_mean"].round(3),
        enrichment_mean=lambda d: d["enrichment_mean"].round(4),
        f1_mean=lambda d: d["f1_mean"].round(4),
        elapsed_total=lambda d: d["elapsed_total"].round(1),
    ).to_string(index=False)
)