Upload folder using huggingface_hub

src/admet.py

@@ -0,0 +1,179 @@
+"""
+ADMET & In Silico Drug-Likeness Profiling Engine.
+Computes physicochemical descriptors, Lipinski Rule of 5, Veber, Ghose, Egan,
+QED, PAINS filter alerts, Brenk toxicophores, and Bioavailability Radar coordinates.
+"""
+from typing import Dict, Any, List, Tuple, Optional
+import numpy as np
+from rdkit import Chem
+from rdkit.Chem import Descriptors, Lipinski, Crippen, QED
+from rdkit.Chem.FilterCatalog import FilterCatalog, FilterCatalogParams
+
+# Initialize FilterCatalogs for PAINS and BRENK structural alerts
+_pains_params = FilterCatalogParams()
+_pains_params.AddCatalog(FilterCatalogParams.FilterCatalogs.PAINS)
+PAINS_CATALOG = FilterCatalog(_pains_params)
+
+_brenk_params = FilterCatalogParams()
+_brenk_params.AddCatalog(FilterCatalogParams.FilterCatalogs.BRENK)
+BRENK_CATALOG = FilterCatalog(_brenk_params)
+
+
+def compute_admet_profile(mol_or_smiles: Any) -> Dict[str, Any]:
+    """
+    Compute comprehensive in silico ADMET and drug-likeness profile.
+    
+    Args:
+        mol_or_smiles: RDKit Mol object or SMILES string.
+        
+    Returns:
+        Dictionary containing descriptors, rule evaluations, alerts, and scores.
+    """
+    if isinstance(mol_or_smiles, str):
+        mol = Chem.MolFromSmiles(mol_or_smiles)
+    else:
+        mol = mol_or_smiles
+
+    if mol is None:
+        raise ValueError("Invalid molecule provided for ADMET calculation.")
+
+    # 1. Fundamental Physicochemical Descriptors
+    mw = float(Descriptors.MolWt(mol))
+    logp = float(Crippen.MolLogP(mol))
+    tpsa = float(Descriptors.TPSA(mol))
+    hbd = int(Lipinski.NumHDonors(mol))
+    hba = int(Lipinski.NumHAcceptors(mol))
+    rot_bonds = int(Lipinski.NumRotatableBonds(mol))
+    heavy_atoms = int(mol.GetNumHeavyAtoms())
+    aromatic_rings = int(Lipinski.NumAromaticRings(mol))
+    fraction_csp3 = float(Descriptors.FractionCSP3(mol))
+    molar_refractivity = float(Crippen.MolMR(mol))
+    qed_score = float(QED.qed(mol))
+
+    # Synthetic Accessibility Approximation (heuristic based on complexity, rings, chiral centers)
+    num_chiral = len(Chem.FindMolChiralCenters(mol, includeUnassigned=True))
+    num_rings = Descriptors.RingCount(mol)
+    bertz = Descriptors.BertzCT(mol)
+    # Normalized complexity heuristic between 1.0 (very easy) and 10.0 (very hard)
+    sascore_est = min(10.0, max(1.0, 1.0 + (bertz / 250.0) + (num_chiral * 0.4) + (max(0, num_rings - 2) * 0.5)))
+
+    # 2. Rule of 5 (Lipinski)
+    lipinski_violations = []
+    if mw > 500:
+        lipinski_violations.append(f"MW > 500 ({mw:.1f} Da)")
+    if logp > 5.0:
+        lipinski_violations.append(f"LogP > 5.0 ({logp:.2f})")
+    if hbd > 5:
+        lipinski_violations.append(f"HBD > 5 ({hbd})")
+    if hba > 10:
+        lipinski_violations.append(f"HBA > 10 ({hba})")
+
+    lipinski_pass = len(lipinski_violations) <= 1
+
+    # 3. Veber Rule (Oral Bioavailability)
+    veber_violations = []
+    if rot_bonds > 10:
+        veber_violations.append(f"Rotatable Bonds > 10 ({rot_bonds})")
+    if tpsa > 140.0:
+        veber_violations.append(f"TPSA > 140 Ų ({tpsa:.1f})")
+    veber_pass = len(veber_violations) == 0
+
+    # 4. Ghose Filter
+    ghose_violations = []
+    if not (160 <= mw <= 480):
+        ghose_violations.append(f"MW not in [160, 480] ({mw:.1f})")
+    if not (-0.4 <= logp <= 5.6):
+        ghose_violations.append(f"LogP not in [-0.4, 5.6] ({logp:.2f})")
+    if not (40 <= molar_refractivity <= 130):
+        ghose_violations.append(f"MR not in [40, 130] ({molar_refractivity:.1f})")
+    if not (20 <= heavy_atoms <= 70):
+        ghose_violations.append(f"Heavy atoms not in [20, 70] ({heavy_atoms})")
+    ghose_pass = len(ghose_violations) == 0
+
+    # 5. Egan Filter (Absorption)
+    egan_violations = []
+    if logp > 5.88:
+        egan_violations.append(f"LogP > 5.88 ({logp:.2f})")
+    if tpsa > 131.6:
+        egan_violations.append(f"TPSA > 131.6 ({tpsa:.1f})")
+    egan_pass = len(egan_violations) == 0
+
+    # 6. PAINS Filter (Pan-Assay Interference Compounds)
+    pains_matches = []
+    for entry in PAINS_CATALOG.GetMatches(mol):
+        pains_matches.append(entry.GetDescription())
+
+    # 7. Brenk Filter (Toxicophores & Reactive Groups)
+    brenk_matches = []
+    for entry in BRENK_CATALOG.GetMatches(mol):
+        brenk_matches.append(entry.GetDescription())
+
+    # 8. Bioavailability Radar Coordinates (Normalized to [0, 1] relative to drug-like box)
+    # Axes: LIPO (LogP), SIZE (MW), POLAR (TPSA), INSOLU (LogS proxy/LogP), INSATU (FractionCsp3), FLEX (RotBonds)
+    # Optimal ranges:
+    # LIPO: -0.7 to 5.0
+    # SIZE: 150 to 500
+    # POLAR: 20 to 130
+    # INSOLU: 0 to 6
+    # INSATU: 0.25 to 1.0
+    # FLEX: 0 to 9
+    radar_coords = {
+        "LIPO": float(np.clip((logp - (-0.7)) / (5.0 - (-0.7)), 0.05, 1.2)),
+        "SIZE": float(np.clip((mw - 150) / (500 - 150), 0.05, 1.2)),
+        "POLAR": float(np.clip((tpsa - 20) / (130 - 20), 0.05, 1.2)),
+        "INSOLU": float(np.clip(max(0, logp) / 6.0, 0.05, 1.2)),
+        "INSATU": float(np.clip(fraction_csp3, 0.05, 1.2)),
+        "FLEX": float(np.clip(rot_bonds / 9.0, 0.05, 1.2))
+    }
+
+    return {
+        "descriptors": {
+            "Molecular Weight": mw,
+            "LogP": logp,
+            "TPSA": tpsa,
+            "H-Bond Donors": hbd,
+            "H-Bond Acceptors": hba,
+            "Rotatable Bonds": rot_bonds,
+            "Heavy Atoms": heavy_atoms,
+            "Aromatic Rings": aromatic_rings,
+            "Fraction Csp3": fraction_csp3,
+            "Molar Refractivity": molar_refractivity,
+            "QED": qed_score,
+            "SA Score (est)": sascore_est
+        },
+        "filters": {
+            "Lipinski": {
+                "passed": lipinski_pass,
+                "num_violations": len(lipinski_violations),
+                "violations": lipinski_violations
+            },
+            "Veber": {
+                "passed": veber_pass,
+                "num_violations": len(veber_violations),
+                "violations": veber_violations
+            },
+            "Ghose": {
+                "passed": ghose_pass,
+                "num_violations": len(ghose_violations),
+                "violations": ghose_violations
+            },
+            "Egan": {
+                "passed": egan_pass,
+                "num_violations": len(egan_violations),
+                "violations": egan_violations
+            }
+        },
+        "alerts": {
+            "PAINS": {
+                "has_alerts": len(pains_matches) > 0,
+                "count": len(pains_matches),
+                "descriptions": pains_matches
+            },
+            "Brenk": {
+                "has_alerts": len(brenk_matches) > 0,
+                "count": len(brenk_matches),
+                "descriptions": brenk_matches
+            }
+        },
+        "radar": radar_coords
+    }

src/analogs.py

@@ -0,0 +1,132 @@
+"""
+In Silico Lead Optimization & Molecular Analog Generator.
+Systematically generates bioisosteric and R-group substituted chemical analogs,
+evaluates property shifts (Delta logS, Delta BBBP, Delta Tox), and discovers Pareto-optimal leads.
+"""
+from typing import List, Dict, Any, Optional
+from rdkit import Chem, RDLogger
+from rdkit.Chem import AllChem, Descriptors, QED
+
+from src.admet import compute_admet_profile
+from src.ensemble import predict_ensemble
+
+RDLogger.DisableLog('rdApp.*')
+
+# Common medicinal chemistry transformation rules
+MEDCHEM_RULES = [
+    # Aromatic functional group replacements
+    {"name": "+Fluorine (-F)", "smarts": "[c:1][H]", "repl": "[c:1]F", "rationale": "Improves metabolic stability & lipophilicity"},
+    {"name": "+Chlorine (-Cl)", "smarts": "[c:1][H]", "repl": "[c:1]Cl", "rationale": "Enhances binding cavity packing"},
+    {"name": "+Methyl (-CH3)", "smarts": "[c:1][H]", "repl": "[c:1]C", "rationale": "Fills hydrophobic subpocket"},
+    {"name": "+Trifluoromethyl (-CF3)", "smarts": "[c:1][H]", "repl": "[c:1]C(F)(F)F", "rationale": "Strongly enhances lipophilic binding"},
+    {"name": "+Hydroxyl (-OH)", "smarts": "[c:1][H]", "repl": "[c:1]O", "rationale": "Improves aqueous solubility (H-bond donor)"},
+    {"name": "+Methoxy (-OCH3)", "smarts": "[c:1][H]", "repl": "[c:1]OC", "rationale": "H-bond acceptor without acidic proton"},
+    {"name": "+Amino (-NH2)", "smarts": "[c:1][H]", "repl": "[c:1]N", "rationale": "Basic center, enhances ionic interactions"},
+    {"name": "+Cyano (-C#N)", "smarts": "[c:1][H]", "repl": "[c:1]C#N", "rationale": "Linear polar group, bioisosteric to halide"},
+    {"name": "+Methylsulfonyl (-SO2CH3)", "smarts": "[c:1][H]", "repl": "[c:1]S(=O)(=O)C", "rationale": "High polarity, decreases lipophilicity"},
+    
+    # Functional group bioisosteres
+    {"name": "Ester -> Amide", "smarts": "[C:1](=O)[O:2][C:3]", "repl": "[C:1](=O)[N:2][C:3]", "rationale": "Improves hydrolytic stability in plasma"},
+    {"name": "Acid -> Tetrazole (Bioisostere)", "smarts": "[C:1](=O)[OH]", "repl": "[C:1]c1nnnn1", "rationale": "Non-classical carboxylic acid bioisostere"},
+    {"name": "Phenol -> Fluoro-benzene", "smarts": "[c:1][OH]", "repl": "[c:1]F", "rationale": "Blocks Phase-II glucuronidation metabolism"}
+]
+
+
+def generate_analogs(parent_smiles: str, max_analogs: int = 15) -> List[Dict[str, Any]]:
+    """
+    Generate chemically plausible analogs of the parent molecule.
+    
+    Args:
+        parent_smiles: SMILES string of the lead compound.
+        max_analogs: Maximum number of distinct valid analogs to generate.
+        
+    Returns:
+        List of generated analog dictionaries with SMILES and transformation details.
+    """
+    mol = Chem.MolFromSmiles(parent_smiles)
+    if mol is None:
+        return []
+
+    # Explicit hydrogens for aromatic replacement rules
+    mol_with_h = Chem.AddHs(mol)
+    generated = {}
+
+    for rule in MEDCHEM_RULES:
+        rxn_smarts = f"({rule['smarts']})>>({rule['repl']})"
+        try:
+            rxn = AllChem.ReactionFromSmarts(rxn_smarts)
+            # Try on mol or mol_with_h
+            reactants = (mol_with_h,) if "[H]" in rule['smarts'] else (mol,)
+            products = rxn.RunReactants(reactants)
+
+            for prod_tuple in products:
+                prod = prod_tuple[0]
+                try:
+                    prod = Chem.RemoveHs(prod)
+                    Chem.SanitizeMol(prod)
+                    smi = Chem.MolToSmiles(prod)
+                    # Filter out identity or duplicates
+                    if smi and smi != parent_smiles and smi not in generated:
+                        generated[smi] = {
+                            "smiles": smi,
+                            "transformation": rule["name"],
+                            "rationale": rule["rationale"]
+                        }
+                        if len(generated) >= max_analogs:
+                            break
+                except Exception:
+                    continue
+        except Exception:
+            continue
+
+        if len(generated) >= max_analogs:
+            break
+
+    return list(generated.values())
+
+
+def evaluate_lead_optimization(
+    parent_smiles: str,
+    target_dataset: str = "esol",
+    max_analogs: int = 12
+) -> Dict[str, Any]:
+    """
+    Generate analogs and evaluate property shifts against the parent compound.
+    """
+    parent_admet = compute_admet_profile(parent_smiles)
+    parent_pred = predict_ensemble(parent_smiles, target_dataset)
+    parent_val = float(parent_pred["mean"][0])
+
+    analogs = generate_analogs(parent_smiles, max_analogs=max_analogs)
+    analog_results = []
+
+    for item in analogs:
+        smi = item["smiles"]
+        try:
+            admet = compute_admet_profile(smi)
+            pred = predict_ensemble(smi, target_dataset)
+            val = float(pred["mean"][0])
+            delta = val - parent_val
+
+            analog_results.append({
+                "SMILES": smi,
+                "Transformation": item["transformation"],
+                "Rationale": item["rationale"],
+                "Predicted Value": val,
+                "Delta": delta,
+                "QED": admet["descriptors"]["QED"],
+                "MW": admet["descriptors"]["Molecular Weight"],
+                "LogP": admet["descriptors"]["LogP"],
+                "TPSA": admet["descriptors"]["TPSA"],
+                "Lipinski Pass": admet["filters"]["Lipinski"]["passed"]
+            })
+        except Exception:
+            continue
+
+    return {
+        "parent_smiles": parent_smiles,
+        "parent_val": parent_val,
+        "parent_qed": parent_admet["descriptors"]["QED"],
+        "target_dataset": target_dataset,
+        "analogs": analog_results
+    }

src/batch_screen.py

@@ -0,0 +1,134 @@
+"""
+High-Throughput Virtual Screening & Batch Profiling Engine.
+Screens libraries of candidate SMILES across ESOL, BBBP, Tox21, and ADMET filters,
+and ranks compounds using a Multi-Objective Drug Desirability Score.
+"""
+from typing import List, Dict, Any, Optional
+import pandas as pd
+import numpy as np
+from rdkit import Chem, RDLogger
+from rdkit.Chem import Descriptors, QED
+
+from src.admet import compute_admet_profile
+from src.ensemble import predict_all_endpoints
+
+RDLogger.DisableLog('rdApp.*')
+
+
+def calculate_desirability_score(
+    qed: float,
+    lipinski_pass: bool,
+    log_s: float,
+    max_tox_prob: float,
+    has_pains: bool,
+    prefer_bbb: bool = False,
+    bbb_prob: float = 0.5
+) -> float:
+    """
+    Compute a composite Drug Desirability Score between 0.0 and 100.0.
+    """
+    # 1. QED component (0 - 30 pts)
+    score_qed = float(np.clip(qed, 0.0, 1.0)) * 30.0
+
+    # 2. Lipinski compliance (0 or 20 pts)
+    score_lipinski = 20.0 if lipinski_pass else 0.0
+
+    # 3. Solubility (0 - 20 pts) - optimal between -1 and -3
+    if log_s >= -2.0:
+        score_sol = 20.0
+    elif log_s >= -4.0:
+        score_sol = 15.0
+    elif log_s >= -6.0:
+        score_sol = 8.0
+    else:
+        score_sol = 2.0
+
+    # 4. Toxicity Safety (0 - 20 pts)
+    score_tox = max(0.0, (1.0 - max_tox_prob) * 20.0)
+
+    # 5. Assay Interference penalty (-10 pts if PAINS alert)
+    penalty_pains = -10.0 if has_pains else 0.0
+
+    # 6. BBB adjustment (0 - 10 pts)
+    if prefer_bbb:
+        score_bbb = bbb_prob * 10.0
+    else:
+        # Non-CNS drug preference (lower BBB is better)
+        score_bbb = (1.0 - bbb_prob) * 10.0
+
+    total = max(0.0, min(100.0, score_qed + score_lipinski + score_sol + score_tox + score_bbb + penalty_pains))
+    return round(total, 1)
+
+
+def screen_batch_molecules(
+    smiles_list: List[str],
+    compound_names: Optional[List[str]] = None,
+    prefer_bbb: bool = False,
+    progress_callback: Optional[Any] = None
+) -> pd.DataFrame:
+    """
+    Screen a collection of SMILES strings and return a ranked DataFrame.
+    """
+    records = []
+    total = len(smiles_list)
+
+    for idx, smi in enumerate(smiles_list):
+        smi = str(smi).strip()
+        if not smi:
+            continue
+
+        name = compound_names[idx] if compound_names and idx < len(compound_names) else f"Compound_{idx+1}"
+        mol = Chem.MolFromSmiles(smi)
+        if mol is None:
+            continue
+
+        try:
+            admet = compute_admet_profile(mol)
+            preds = predict_all_endpoints(smi)
+
+            esol_val = float(preds['esol']['mean'][0])
+            bbbp_prob = float(preds['bbbp']['mean'][0])
+            tox_probs = preds['tox21']['mean']
+            max_tox = float(np.max(tox_probs))
+            mean_tox = float(np.mean(tox_probs))
+            high_risk_tox_count = int(np.sum(tox_probs >= 0.5))
+
+            qed_score = admet['descriptors']['QED']
+            lip_pass = admet['filters']['Lipinski']['passed']
+            pains_alert = admet['alerts']['PAINS']['has_alerts']
+
+            desirability = calculate_desirability_score(
+                qed=qed_score,
+                lipinski_pass=lip_pass,
+                log_s=esol_val,
+                max_tox_prob=max_tox,
+                has_pains=pains_alert,
+                prefer_bbb=prefer_bbb,
+                bbb_prob=bbbp_prob
+            )
+
+            records.append({
+                "Compound ID": name,
+                "SMILES": smi,
+                "Desirability Score": desirability,
+                "QED": round(qed_score, 3),
+                "Solubility logS": round(esol_val, 2),
+                "BBB Probability": f"{bbbp_prob*100:.1f}%",
+                "Max Tox Risk": f"{max_tox*100:.1f}%",
+                "High Risk Assays": high_risk_tox_count,
+                "MW (Da)": round(admet['descriptors']['Molecular Weight'], 1),
+                "LogP": round(admet['descriptors']['LogP'], 2),
+                "TPSA (Ų)": round(admet['descriptors']['TPSA'], 1),
+                "Lipinski": "PASS" if lip_pass else f"FAIL ({admet['filters']['Lipinski']['num_violations']})",
+                "PAINS Alert": "YES" if pains_alert else "Clean"
+            })
+        except Exception:
+            continue
+
+        if progress_callback:
+            progress_callback((idx + 1) / total)
+
+    df = pd.DataFrame(records)
+    if not df.empty:
+        df = df.sort_values(by="Desirability Score", ascending=False).reset_index(drop=True)
+    return df

src/chemical_space.py

@@ -0,0 +1,148 @@
+"""
+Chemical Space Explorer & Structural Similarity Search Engine.
+Projects molecules into 2D chemical space (PCA over Morgan fingerprints & descriptors)
+and searches for nearest reference neighbors via Tanimoto structural similarity.
+"""
+from typing import Dict, List, Tuple, Any, Optional
+import os
+import numpy as np
+import pandas as pd
+from rdkit import Chem, DataStructs, RDLogger
+from rdkit.Chem import AllChem, Descriptors
+
+# Silence RDKit deprecation noise
+RDLogger.DisableLog('rdApp.*')
+
+
+class ChemicalSpaceExplorer:
+    """Manages chemical space projections and nearest neighbor retrieval."""
+    def __init__(self, dataset_csv: str = "data/raw/delaney-processed.csv", max_compounds: int = 1200):
+        self.dataset_csv = dataset_csv
+        self.df = None
+        self.fps = []
+        self.valid_indices = []
+        self.pca_mean = None
+        self.pca_components = None
+        self._load_and_index(max_compounds)
+
+    def _load_and_index(self, max_compounds: int):
+        if not os.path.exists(self.dataset_csv):
+            # Fallback if file missing
+            return
+
+        raw_df = pd.read_csv(self.dataset_csv)
+        if len(raw_df) > max_compounds:
+            raw_df = raw_df.sample(max_compounds, random_state=42).reset_index(drop=True)
+
+        smiles_col = 'smiles' if 'smiles' in raw_df.columns else raw_df.columns[-1]
+        name_col = 'Compound ID' if 'Compound ID' in raw_df.columns else ('name' if 'name' in raw_df.columns else None)
+        target_col = [c for c in raw_df.columns if 'solubility' in c.lower() or 'p_np' in c.lower()]
+        target_name = target_col[0] if target_col else None
+
+        valid_mols = []
+        valid_rows = []
+        fp_matrix = []
+
+        for idx, row in raw_df.iterrows():
+            smi = str(row[smiles_col]).strip()
+            mol = Chem.MolFromSmiles(smi)
+            if mol is not None:
+                # ECFP4 fingerprint (2048-bit)
+                fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=512)
+                arr = np.zeros((512,), dtype=np.float32)
+                DataStructs.ConvertToNumpyArray(fp, arr)
+                
+                # Add basic descriptors for richer projection
+                mw = Descriptors.MolWt(mol) / 500.0
+                logp = Descriptors.MolLogP(mol) / 5.0
+                tpsa = Descriptors.TPSA(mol) / 150.0
+                feats = np.concatenate([arr, [mw, logp, tpsa]])
+
+                fp_matrix.append(feats)
+                self.fps.append(fp)
+                valid_mols.append(mol)
+                valid_rows.append({
+                    "Compound ID": str(row[name_col]) if name_col and pd.notna(row[name_col]) else f"CMPD_{idx}",
+                    "SMILES": smi,
+                    "Target Value": float(row[target_name]) if target_name and pd.notna(row[target_name]) else 0.0
+                })
+
+        if not fp_matrix:
+            return
+
+        self.df = pd.DataFrame(valid_rows)
+        X = np.stack(fp_matrix, axis=0)
+
+        # Compute deterministic 2D PCA for projection
+        self.pca_mean = np.mean(X, axis=0)
+        X_centered = X - self.pca_mean
+        cov = np.cov(X_centered, rowvar=False)
+        eigenvalues, eigenvectors = np.linalg.eigh(cov)
+        idx = np.argsort(eigenvalues)[::-1]
+        self.pca_components = eigenvectors[:, idx[:2]]
+
+        coords = np.dot(X_centered, self.pca_components)
+        self.df["PC1"] = coords[:, 0]
+        self.df["PC2"] = coords[:, 1]
+
+    def project_molecule(self, smiles: str) -> Optional[Tuple[float, float]]:
+        """Project a single SMILES query into the chemical space PC1-PC2 coordinates."""
+        if self.pca_components is None:
+            return None
+
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return None
+
+        fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=512)
+        arr = np.zeros((512,), dtype=np.float32)
+        DataStructs.ConvertToNumpyArray(fp, arr)
+        mw = Descriptors.MolWt(mol) / 500.0
+        logp = Descriptors.MolLogP(mol) / 5.0
+        tpsa = Descriptors.TPSA(mol) / 150.0
+        feats = np.concatenate([arr, [mw, logp, tpsa]])
+
+        centered = feats - self.pca_mean
+        coords = np.dot(centered, self.pca_components)
+        return float(coords[0]), float(coords[1])
+
+    def find_nearest_neighbors(self, smiles: str, top_k: int = 5) -> List[Dict[str, Any]]:
+        """Search reference compounds by Tanimoto similarity over Morgan fingerprints."""
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None or not self.fps:
+            return []
+
+        query_fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=512)
+        sims = DataStructs.BulkTanimotoSimilarity(query_fp, self.fps)
+        top_indices = np.argsort(sims)[::-1][:top_k]
+
+        results = []
+        for rank, i in enumerate(top_indices):
+            row = self.df.iloc[i]
+            results.append({
+                "Rank": rank + 1,
+                "Compound ID": row["Compound ID"],
+                "SMILES": row["SMILES"],
+                "Tanimoto Similarity": float(sims[i]),
+                "Target Value": row["Target Value"]
+            })
+        return results
+
+
+# Global singleton cache
+_explorer_instances = {}
+
+
+def get_chemical_space_explorer(dataset_name: str = "esol") -> ChemicalSpaceExplorer:
+    """Get or instantiate cached chemical space explorer."""
+    name = dataset_name.lower()
+    if name not in _explorer_instances:
+        if name == "esol":
+            path = "data/raw/delaney-processed.csv"
+        elif name == "bbbp":
+            path = "data/raw/BBBP.csv"
+        else:
+            path = "data/raw/delaney-processed.csv"
+
+        _explorer_instances[name] = ChemicalSpaceExplorer(dataset_csv=path)
+    return _explorer_instances[name]

src/conformers.py

@@ -0,0 +1,121 @@
+"""
+3D Conformer Generation and Interactive Molecular Modeling.
+Generates energetic 3D molecular conformations using ETKDG and MMFF94 force field,
+and constructs py3Dmol visualizer HTML blocks for Streamlit components.
+"""
+from typing import Tuple, Optional
+from rdkit import Chem
+from rdkit.Chem import AllChem
+import py3Dmol
+
+
+def generate_3d_conformer(smiles_or_mol: any, seed: int = 42) -> Tuple[Optional[Chem.Mol], Optional[str], Optional[float], str]:
+    """
+    Generate an energy-minimized 3D conformer for a molecule.
+    
+    Args:
+        smiles_or_mol: SMILES string or RDKit Mol object.
+        seed: Random seed for deterministic coordinate embedding.
+        
+    Returns:
+        Tuple of (mol_3d, mol_block_str, energy_kcal_per_mol, status_message)
+    """
+    if isinstance(smiles_or_mol, str):
+        mol = Chem.MolFromSmiles(smiles_or_mol)
+    else:
+        mol = Chem.Mol(smiles_or_mol)
+
+    if mol is None:
+        return None, None, None, "Invalid chemical structure"
+
+    mol_h = Chem.AddHs(mol)
+
+    # Use ETKDG embedding
+    params = AllChem.ETKDGv3() if hasattr(AllChem, 'ETKDGv3') else AllChem.ETKDG()
+    params.randomSeed = seed
+    embed_code = AllChem.EmbedMolecule(mol_h, params)
+
+    if embed_code != 0:
+        # Fallback with relaxed distance geometry
+        embed_code = AllChem.EmbedMolecule(mol_h, randomSeed=seed, useRandomCoords=True)
+        if embed_code != 0:
+            return None, None, None, "Failed to embed 3D coordinates"
+
+    # MMFF94 Force Field Optimization
+    energy = None
+    status = "MMFF94 optimized"
+    if AllChem.MMFFHasAllMoleculeParams(mol_h):
+        mmff_props = AllChem.MMFFGetMoleculeProperties(mol_h)
+        ff = AllChem.MMFFGetMoleculeForceField(mol_h, mmff_props)
+        if ff:
+            ff.Initialize()
+            ff.Minimize(maxIts=500)
+            energy = float(ff.CalcEnergy())
+    else:
+        # Fallback to Universal Force Field (UFF)
+        if AllChem.UFFHasAllMoleculeParams(mol_h):
+            ff = AllChem.UFFGetMoleculeForceField(mol_h)
+            if ff:
+                ff.Initialize()
+                ff.Minimize(maxIts=500)
+                energy = float(ff.CalcEnergy())
+                status = "UFF optimized (MMFF params unavailable for this element)"
+
+    mol_block = Chem.MolToMolBlock(mol_h)
+    return mol_h, mol_block, energy, status
+
+
+def render_3dmol_html(
+    mol_block: str,
+    style: str = "stick",
+    surface_type: str = "none",
+    spin: bool = False,
+    width: int = 500,
+    height: int = 400
+) -> str:
+    """
+    Generate self-contained HTML for interactive 3D molecular visualization using py3Dmol.
+    
+    Args:
+        mol_block: MDL MolBlock representation of the 3D conformer.
+        style: Visual representation ('stick', 'ball_and_stick', 'sphere', 'wireframe').
+        surface_type: Molecular surface representation ('none', 'vdw', 'sas', 'ses').
+        spin: Whether to enable auto-rotation.
+        width: Viewer width in pixels.
+        height: Viewer height in pixels.
+        
+    Returns:
+        HTML string suitable for st.components.v1.html.
+    """
+    view = py3Dmol.view(width=width, height=height)
+    view.addModel(mol_block, "mol")
+
+    if style == "stick":
+        view.setStyle({'stick': {'radius': 0.18, 'colorscheme': 'cyanCarbon'}})
+    elif style == "ball_and_stick":
+        view.setStyle({
+            'stick': {'radius': 0.14, 'colorscheme': 'cyanCarbon'},
+            'sphere': {'scale': 0.28, 'colorscheme': 'cyanCarbon'}
+        })
+    elif style == "sphere":
+        view.setStyle({'sphere': {'scale': 0.85, 'colorscheme': 'cyanCarbon'}})
+    elif style == "wireframe":
+        view.setStyle({'line': {'linewidth': 2.0, 'colorscheme': 'cyanCarbon'}})
+    else:
+        view.setStyle({'stick': {}})
+
+    # Optional molecular surface
+    surf_map = {
+        'vdw': py3Dmol.VDW,
+        'sas': py3Dmol.SAS,
+        'ses': py3Dmol.SES
+    }
+    if surface_type.lower() in surf_map:
+        view.addSurface(surf_map[surface_type.lower()], {'opacity': 0.65, 'color': 'lightblue'})
+
+    if spin:
+        view.spin(True)
+
+    view.zoomTo()
+    view.setBackgroundColor('#0f172a')  # Dark slate background for premium biotech feel
+    return view._make_html()

src/data.py

@@ -0,0 +1,376 @@
+"""
+Molecular Graph Featurization, Dataset Loading, and Scaffold Splitting.
+Supports ESOL, BBBP, and Tox21 datasets with deterministic Bemis-Murcko splitting
+and RDKit SMILES validation.
+"""
+import os
+import gzip
+import logging
+import urllib.request
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Any
+
+import pandas as pd
+import numpy as np
+import torch
+from torch_geometric.data import Data
+from torch_geometric.loader import DataLoader
+import rdkit
+from rdkit import Chem
+from rdkit.Chem import AllChem
+from rdkit.Chem.Scaffolds import MurckoScaffold
+
+# Setup logger
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger(__name__)
+
+# Allowed categorical sets for one-hot encodings
+ALLOWED_ATOMS = [
+    'H', 'B', 'C', 'N', 'O', 'F', 'Na', 'Mg', 'Al', 'Si', 'P', 'S',
+    'Cl', 'K', 'Ca', 'Fe', 'Cu', 'Zn', 'Br', 'I', 'other'
+]
+ALLOWED_DEGREES = [0, 1, 2, 3, 4, 5, 6, 'other']
+ALLOWED_CHARGES = [-2, -1, 0, 1, 2, 'other']
+ALLOWED_HYBRIDIZATIONS = [
+    Chem.rdchem.HybridizationType.SP,
+    Chem.rdchem.HybridizationType.SP2,
+    Chem.rdchem.HybridizationType.SP3,
+    Chem.rdchem.HybridizationType.SP3D,
+    Chem.rdchem.HybridizationType.SP3D2,
+    'other'
+]
+ALLOWED_TOTAL_HS = [0, 1, 2, 3, 4, 'other']
+
+ALLOWED_BONDS = [
+    Chem.rdchem.BondType.SINGLE,
+    Chem.rdchem.BondType.DOUBLE,
+    Chem.rdchem.BondType.TRIPLE,
+    Chem.rdchem.BondType.AROMATIC
+]
+ALLOWED_STEREO = [
+    Chem.rdchem.BondStereo.STEREONONE,
+    Chem.rdchem.BondStereo.STEREOANY,
+    Chem.rdchem.BondStereo.STEREOZ,
+    Chem.rdchem.BondStereo.STEREOE
+]
+
+# MoleculeNet dataset configs
+DATASET_CONFIGS = {
+    'esol': {
+        'url': 'https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/delaney-processed.csv',
+        'raw_filename': 'delaney-processed.csv',
+        'smiles_col': 'smiles',
+        'target_cols': ['measured log solubility in mols per litre'],
+        'task_type': 'regression',
+        'num_tasks': 1,
+        'metric': 'rmse'
+    },
+    'bbbp': {
+        'url': 'https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/BBBP.csv',
+        'raw_filename': 'BBBP.csv',
+        'smiles_col': 'smiles',
+        'target_cols': ['p_np'],
+        'task_type': 'classification',
+        'num_tasks': 1,
+        'metric': 'roc_auc'
+    },
+    'tox21': {
+        'url': 'https://deepchemdata.s3-us-west-1.amazonaws.com/datasets/tox21.csv.gz',
+        'raw_filename': 'tox21.csv.gz',
+        'smiles_col': 'smiles',
+        'target_cols': [
+            'NR-AR', 'NR-AR-LBD', 'NR-AhR', 'NR-Aromatase', 'NR-ER', 'NR-ER-LBD',
+            'NR-PPAR-gamma', 'SR-ARE', 'SR-ATAD5', 'SR-HSE', 'SR-MMP', 'SR-p53'
+        ],
+        'task_type': 'multitask_classification',
+        'num_tasks': 12,
+        'metric': 'roc_auc'
+    }
+}
+
+
+def one_hot_encode(val: Any, allowed: List[Any]) -> List[float]:
+    """Helper to one-hot encode a value given allowed categories with fallback to 'other'."""
+    vec = [0.0] * len(allowed)
+    if val in allowed:
+        vec[allowed.index(val)] = 1.0
+    else:
+        vec[-1] = 1.0
+    return vec
+
+
+def get_atom_features(atom: Chem.Atom) -> List[float]:
+    """Extract node features for an RDKit Atom."""
+    return (
+        one_hot_encode(atom.GetSymbol(), ALLOWED_ATOMS) +
+        one_hot_encode(atom.GetTotalDegree(), ALLOWED_DEGREES) +
+        one_hot_encode(atom.GetFormalCharge(), ALLOWED_CHARGES) +
+        one_hot_encode(atom.GetHybridization(), ALLOWED_HYBRIDIZATIONS) +
+        [1.0 if atom.GetIsAromatic() else 0.0] +
+        one_hot_encode(atom.GetTotalNumHs(), ALLOWED_TOTAL_HS)
+    )
+
+
+def get_bond_features(bond: Chem.Bond) -> List[float]:
+    """Extract edge features for an RDKit Bond."""
+    return (
+        one_hot_encode(bond.GetBondType(), ALLOWED_BONDS) +
+        [1.0 if bond.GetIsConjugated() else 0.0] +
+        [1.0 if bond.IsInRing() else 0.0] +
+        one_hot_encode(bond.GetStereo(), ALLOWED_STEREO)
+    )
+
+
+# Calculate and expose dimensions
+NODE_DIM = len(ALLOWED_ATOMS) + len(ALLOWED_DEGREES) + len(ALLOWED_CHARGES) + \
+           len(ALLOWED_HYBRIDIZATIONS) + 1 + len(ALLOWED_TOTAL_HS)
+EDGE_DIM = len(ALLOWED_BONDS) + 1 + 1 + len(ALLOWED_STEREO)
+
+
+def smiles_to_graph(smiles: str, y: Optional[torch.Tensor] = None) -> Optional[Data]:
+    """
+    Convert a SMILES string to a PyG Data object.
+    Returns None if SMILES is invalid or cannot be parsed by RDKit.
+    """
+    if not isinstance(smiles, str) or len(smiles.strip()) == 0:
+        return None
+
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return None
+
+    # Compute node features
+    atom_feats = [get_atom_features(atom) for atom in mol.GetAtoms()]
+    x = torch.tensor(atom_feats, dtype=torch.float)
+
+    # Compute edges and edge features (undirected)
+    edge_indices = []
+    edge_attrs = []
+    for bond in mol.GetBonds():
+        i = bond.GetBeginAtomIdx()
+        j = bond.GetEndAtomIdx()
+        bond_feats = get_bond_features(bond)
+
+        # Both directions (i -> j) and (j -> i)
+        edge_indices.append([i, j])
+        edge_attrs.append(bond_feats)
+        edge_indices.append([j, i])
+        edge_attrs.append(bond_feats)
+
+    if len(edge_indices) > 0:
+        edge_index = torch.tensor(edge_indices, dtype=torch.long).t().contiguous()
+        edge_attr = torch.tensor(edge_attrs, dtype=torch.float)
+    else:
+        # Isolated atoms (no bonds)
+        edge_index = torch.empty((2, 0), dtype=torch.long)
+        edge_attr = torch.empty((0, EDGE_DIM), dtype=torch.float)
+
+    data = Data(
+        x=x,
+        edge_index=edge_index,
+        edge_attr=edge_attr,
+        smiles=smiles,
+        num_nodes=mol.GetNumAtoms()
+    )
+    if y is not None:
+        data.y = y
+
+    return data
+
+
+def smiles_to_fingerprint(smiles: str, radius: int = 2, n_bits: int = 2048) -> Optional[np.ndarray]:
+    """Generate Morgan fingerprint (ECFP4) bit vector from SMILES."""
+    if not isinstance(smiles, str) or len(smiles.strip()) == 0:
+        return None
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return None
+    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=radius, nBits=n_bits)
+    arr = np.zeros((n_bits,), dtype=np.float32)
+    for bit in fp.GetOnBits():
+        arr[bit] = 1.0
+    return arr
+
+
+def generate_scaffold(smiles: str, include_chirality: bool = False) -> str:
+    """Compute Bemis-Murcko scaffold string for a SMILES."""
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return ""
+    scaffold = MurckoScaffold.MurckoScaffoldSmiles(smiles=smiles, includeChirality=include_chirality)
+    return scaffold
+
+
+def deterministic_scaffold_split(
+    data_list: List[Data],
+    train_frac: float = 0.8,
+    val_frac: float = 0.1,
+    test_frac: float = 0.1
+) -> Tuple[List[Data], List[Data], List[Data]]:
+    """
+    Split molecular dataset into train, val, test using deterministic Bemis-Murcko scaffold split.
+    Sorts scaffold groups by size (descending), then alphabetically by scaffold string.
+    """
+    assert abs(train_frac + val_frac + test_frac - 1.0) < 1e-5, "Fractions must sum to 1.0"
+
+    # Group sample indices by scaffold
+    scaffold_to_indices = defaultdict(list)
+    for idx, data in enumerate(data_list):
+        scaffold = generate_scaffold(data.smiles)
+        scaffold_to_indices[scaffold].append(idx)
+
+    # Sort scaffold groups: primary key = -len(indices) (largest groups first), secondary = scaffold string
+    sorted_scaffold_groups = sorted(
+        scaffold_to_indices.items(),
+        key=lambda item: (-len(item[1]), item[0])
+    )
+
+    n_total = len(data_list)
+    train_cutoff = train_frac * n_total
+    val_cutoff = (train_frac + val_frac) * n_total
+
+    train_indices, val_indices, test_indices = [], [], []
+
+    for _, indices in sorted_scaffold_groups:
+        if len(train_indices) + len(indices) <= train_cutoff:
+            train_indices.extend(indices)
+        elif len(train_indices) + len(val_indices) + len(indices) <= val_cutoff:
+            val_indices.extend(indices)
+        else:
+            test_indices.extend(indices)
+
+    train_data = [data_list[i] for i in train_indices]
+    val_data = [data_list[i] for i in val_indices]
+    test_data = [data_list[i] for i in test_indices]
+
+    return train_data, val_data, test_data
+
+
+def download_raw_dataset(dataset_name: str, raw_dir: str) -> str:
+    """Download raw dataset if not already present."""
+    config = DATASET_CONFIGS[dataset_name.lower()]
+    os.makedirs(raw_dir, exist_ok=True)
+    raw_path = os.path.join(raw_dir, config['raw_filename'])
+
+    if not os.path.exists(raw_path):
+        logger.info(f"Downloading {dataset_name} from {config['url']} -> {raw_path}")
+        urllib.request.urlretrieve(config['url'], raw_path)
+        logger.info(f"Successfully downloaded {dataset_name}.")
+    else:
+        logger.info(f"Found existing raw data file at {raw_path}")
+
+    return raw_path
+
+
+def load_and_preprocess_dataset(
+    dataset_name: str,
+    data_dir: str = 'data',
+    reload: bool = False
+) -> Dict[str, Any]:
+    """
+    Load, parse, validate, and featurize a MoleculeNet dataset.
+    Saves processed splits on disk for fast reload.
+    """
+    dataset_name = dataset_name.lower()
+    if dataset_name not in DATASET_CONFIGS:
+        raise ValueError(f"Unknown dataset '{dataset_name}'. Allowed: {list(DATASET_CONFIGS.keys())}")
+
+    processed_dir = os.path.join(data_dir, 'processed')
+    raw_dir = os.path.join(data_dir, 'raw')
+    os.makedirs(processed_dir, exist_ok=True)
+    cache_path = os.path.join(processed_dir, f"{dataset_name}_processed.pt")
+
+    if os.path.exists(cache_path) and not reload:
+        logger.info(f"Loading cached dataset from {cache_path}")
+        return torch.load(cache_path, weights_only=False)
+
+    config = DATASET_CONFIGS[dataset_name]
+    raw_path = download_raw_dataset(dataset_name, raw_dir)
+
+    # Read dataframe (plain CSV or gzip)
+    if raw_path.endswith('.gz'):
+        with gzip.open(raw_path, 'rt') as f:
+            df = pd.read_csv(f)
+    else:
+        df = pd.read_csv(raw_path)
+
+    smiles_col = config['smiles_col']
+    target_cols = config['target_cols']
+    initial_rows = len(df)
+
+    valid_graphs: List[Data] = []
+    dropped_count = 0
+
+    for idx, row in df.iterrows():
+        smi = row[smiles_col]
+        # Target values
+        targets = [float(row[col]) if pd.notna(row[col]) else float('nan') for col in target_cols]
+        y_tensor = torch.tensor([targets], dtype=torch.float)
+
+        graph = smiles_to_graph(smi, y=y_tensor)
+        if graph is None:
+            dropped_count += 1
+            logger.debug(f"Row {idx}: Failed to parse SMILES '{smi}'. Dropped.")
+        else:
+            valid_graphs.append(graph)
+
+    final_size = len(valid_graphs)
+    logger.info(f"[{dataset_name.upper()}] Initial rows: {initial_rows} | Dropped invalid SMILES: {dropped_count} | Final valid size: {final_size}")
+    print(f"[{dataset_name.upper()}] Dataset processing complete: {initial_rows} rows -> {dropped_count} dropped -> {final_size} valid molecules.")
+
+    # Deterministic Bemis-Murcko scaffold split
+    train_data, val_data, test_data = deterministic_scaffold_split(valid_graphs, 0.8, 0.1, 0.1)
+
+    print(f"[{dataset_name.upper()}] Split sizes -> Train: {len(train_data)} ({len(train_data)/final_size*100:.1f}%), "
+          f"Val: {len(val_data)} ({len(val_data)/final_size*100:.1f}%), "
+          f"Test: {len(test_data)} ({len(test_data)/final_size*100:.1f}%)")
+
+    dataset_dict = {
+        'dataset_name': dataset_name,
+        'config': config,
+        'train': train_data,
+        'val': val_data,
+        'test': test_data,
+        'initial_size': initial_rows,
+        'dropped_count': dropped_count,
+        'final_size': final_size,
+        'node_dim': NODE_DIM,
+        'edge_dim': EDGE_DIM
+    }
+
+    # Save to disk cache
+    torch.save(dataset_dict, cache_path)
+    logger.info(f"Saved processed dataset to {cache_path}")
+
+    return dataset_dict
+
+
+def get_dataloaders(
+    dataset_name: str,
+    batch_size: int = 32,
+    data_dir: str = 'data',
+    reload: bool = False,
+    quick: bool = False,
+    num_workers: int = 0
+) -> Tuple[DataLoader, DataLoader, DataLoader, Dict[str, Any]]:
+    """
+    Get PyG DataLoaders for train, val, test splits.
+    If quick=True, trims dataset sizes for fast smoke testing.
+    """
+    data_bundle = load_and_preprocess_dataset(dataset_name, data_dir=data_dir, reload=reload)
+
+    train_data = data_bundle['train']
+    val_data = data_bundle['val']
+    test_data = data_bundle['test']
+
+    if quick:
+        train_data = train_data[:min(len(train_data), 100)]
+        val_data = val_data[:min(len(val_data), 20)]
+        test_data = test_data[:min(len(test_data), 20)]
+        logger.info(f"[QUICK MODE] Trimmed to: Train={len(train_data)}, Val={len(val_data)}, Test={len(test_data)}")
+
+    train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+    val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+    test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+
+    return train_loader, val_loader, test_loader, data_bundle

src/ensemble.py

@@ -0,0 +1,149 @@
+"""
+Multi-Model Ensemble & Uncertainty Quantification Engine.
+Executes consensus predictions across GCN, GAT, GIN, MPNN, and Baseline models
+and quantifies epistemic uncertainty via standard deviation and confidence intervals.
+"""
+from typing import Dict, List, Any, Tuple, Optional
+import os
+import numpy as np
+import torch
+
+from src.data import smiles_to_graph, DATASET_CONFIGS
+from src.evaluate import load_trained_model
+
+MODEL_NAMES = ['gin', 'mpnn', 'gcn', 'gat', 'baseline']
+
+
+def predict_single_model(model: Any, model_name: str, graph: Any, dataset_name: str) -> np.ndarray:
+    """Run inference for a single model and return output array (shape: [num_tasks])."""
+    if model_name == 'baseline':
+        pred = model.predict([graph])[0]
+        if np.ndim(pred) == 0:
+            pred = np.array([float(pred)])
+        return np.asarray(pred, dtype=np.float32)
+
+    model.eval()
+    with torch.no_grad():
+        out = model(graph.x, graph.edge_index, edge_attr=graph.edge_attr)
+        if dataset_name == 'esol':
+            val = out[0].cpu().numpy()
+        else:
+            # Classification or multi-task: convert logits to probabilities
+            val = torch.sigmoid(out)[0].cpu().numpy()
+    return np.asarray(val, dtype=np.float32)
+
+
+def predict_ensemble(
+    smiles_or_graph: Any,
+    dataset_name: str,
+    loaded_models: Optional[Dict[str, Any]] = None,
+    checkpoint_dir: str = 'results/checkpoints'
+) -> Dict[str, Any]:
+    """
+    Run multi-model ensemble prediction for a specific dataset endpoint.
+    
+    Args:
+        smiles_or_graph: SMILES string or PyG Data graph.
+        dataset_name: One of 'esol', 'bbbp', 'tox21'.
+        loaded_models: Optional preloaded dict of {model_name: model}.
+        checkpoint_dir: Directory containing saved checkpoints.
+        
+    Returns:
+        Dictionary containing individual predictions, ensemble mean, std, CI, and agreement.
+    """
+    ds = dataset_name.lower()
+    if isinstance(smiles_or_graph, str):
+        graph = smiles_to_graph(smiles_or_graph)
+    else:
+        graph = smiles_or_graph
+
+    if loaded_models is None:
+        loaded_models = {}
+        for m in MODEL_NAMES:
+            try:
+                mod, _ = load_trained_model(m, ds, checkpoint_dir=checkpoint_dir)
+                loaded_models[m] = mod
+            except Exception:
+                pass
+
+    if not loaded_models:
+        raise RuntimeError(f"No trained checkpoints found for dataset '{ds}'.")
+
+    predictions = {}
+    vals_list = []
+
+    for m_name, model in loaded_models.items():
+        try:
+            val = predict_single_model(model, m_name, graph, ds)
+            predictions[m_name] = val
+            vals_list.append(val)
+        except Exception as e:
+            continue
+
+    if not vals_list:
+        raise RuntimeError(f"Failed to obtain predictions for '{ds}'.")
+
+    stacked = np.stack(vals_list, axis=0)  # Shape: (num_models, num_tasks)
+    mean_pred = np.mean(stacked, axis=0)
+    std_pred = np.std(stacked, axis=0)
+
+    # 95% Confidence bounds
+    ci_lower = mean_pred - 1.96 * std_pred
+    ci_upper = mean_pred + 1.96 * std_pred
+
+    # Assess model agreement
+    avg_std = float(np.mean(std_pred))
+    if ds == 'esol':
+        # logS standard deviation threshold
+        if avg_std < 0.25:
+            agreement = "High Consensus (Low Uncertainty)"
+        elif avg_std < 0.50:
+            agreement = "Moderate Consensus"
+        else:
+            agreement = "Divergent Predictions (High Uncertainty)"
+    else:
+        # Probability standard deviation threshold
+        if avg_std < 0.10:
+            agreement = "High Consensus (Low Uncertainty)"
+        elif avg_std < 0.20:
+            agreement = "Moderate Consensus"
+        else:
+            agreement = "Divergent Predictions (High Uncertainty)"
+
+    return {
+        "dataset": ds,
+        "task_names": DATASET_CONFIGS[ds]['target_cols'],
+        "individual": predictions,
+        "mean": mean_pred,
+        "std": std_pred,
+        "ci_lower": ci_lower,
+        "ci_upper": ci_upper,
+        "num_models": len(vals_list),
+        "agreement": agreement,
+        "uncertainty_score": avg_std
+    }
+
+
+def predict_all_endpoints(
+    smiles_or_graph: Any,
+    models_cache: Optional[Dict[str, Dict[str, Any]]] = None,
+    checkpoint_dir: str = 'results/checkpoints'
+) -> Dict[str, Any]:
+    """
+    Run full multi-endpoint prediction suite across ESOL, BBBP, and Tox21.
+    """
+    results = {}
+    for ds in ['esol', 'bbbp', 'tox21']:
+        cache_ds = models_cache.get(ds) if models_cache else None
+        try:
+            res = predict_ensemble(
+                smiles_or_graph=smiles_or_graph,
+                dataset_name=ds,
+                loaded_models=cache_ds,
+                checkpoint_dir=checkpoint_dir
+            )
+            results[ds] = res
+        except Exception as e:
+            results[ds] = {"error": str(e)}
+
+    return results

src/evaluate.py

@@ -0,0 +1,342 @@
+"""
+Evaluation and Diagnostics Suite for GNNs and Classical Baselines.
+Generates parity plots, ROC/PR curves, loss trajectories, t-SNE latent embeddings,
+and compiles results/REPORT.md.
+"""
+import os
+import sys
+
+# Ensure project root is in sys.path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import json
+import logging
+from typing import Dict, List, Any, Optional, Tuple
+
+import yaml
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.manifold import TSNE
+from sklearn.metrics import (
+    mean_squared_error, mean_absolute_error, r2_score,
+    roc_curve, precision_recall_curve, roc_auc_score, average_precision_score
+)
+
+from src.data import get_dataloaders, NODE_DIM, EDGE_DIM
+from src.models.gnn import build_gnn_model
+from src.models.baseline import ClassicalFingerprintBaseline, extract_features
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger(__name__)
+
+# Matplotlib styling for publication quality
+plt.style.use('seaborn-v0_8-whitegrid' if 'seaborn-v0_8-whitegrid' in plt.style.available else 'default')
+plt.rcParams['font.sans-serif'] = 'DejaVu Sans'
+plt.rcParams['axes.edgecolor'] = '#cccccc'
+plt.rcParams['axes.linewidth'] = 0.8
+
+
+def load_trained_model(
+    model_name: str,
+    dataset_name: str,
+    checkpoint_dir: str = 'results/checkpoints'
+) -> Tuple[Any, Optional[Dict[str, Any]]]:
+    """Load trained model checkpoint from disk."""
+    model_name = model_name.lower()
+    dataset_name = dataset_name.lower()
+
+    if model_name == 'baseline':
+        if os.environ.get("LOW_MEMORY_MODE") == "1" and dataset_name == "tox21":
+            # Avoid deserializing 650MB Tox21 Random Forest pickle on 512MB cloud free-tier VMs
+            raise RuntimeError("Tox21 RF baseline disabled in LOW_MEMORY_MODE to preserve RAM limit.")
+        ckpt_path = os.path.join(checkpoint_dir, f"{dataset_name}_baseline_best.pkl")
+        if not os.path.exists(ckpt_path):
+            raise FileNotFoundError(f"Baseline checkpoint not found at {ckpt_path}")
+        return ClassicalFingerprintBaseline.load(ckpt_path), None
+    else:
+        ckpt_path = os.path.join(checkpoint_dir, f"{dataset_name}_{model_name}_best.pt")
+        if not os.path.exists(ckpt_path):
+            raise FileNotFoundError(f"GNN checkpoint not found at {ckpt_path}")
+        payload = torch.load(ckpt_path, weights_only=False)
+        kwargs = payload['model_kwargs']
+        model = build_gnn_model(model_name=model_name, **kwargs)
+        model.load_state_dict(payload['model_state_dict'])
+        model.eval()
+        return model, payload
+
+
+def plot_loss_trajectories(dataset_name: str, history_dir: str = 'results/history', save_dir: str = 'results/plots'):
+    """Plot training loss curves across GNN architectures."""
+    os.makedirs(save_dir, exist_ok=True)
+    models = ['gcn', 'gat', 'gin', 'mpnn']
+    colors = {'gcn': '#1f77b4', 'gat': '#ff7f0e', 'gin': '#2ca02c', 'mpnn': '#d62728'}
+
+    plt.figure(figsize=(8, 5))
+    has_data = False
+
+    for m in models:
+        hist_path = os.path.join(history_dir, f"{dataset_name}_{m}_history.json")
+        if os.path.exists(hist_path):
+            with open(hist_path, 'r') as f:
+                hist = json.load(f)
+            losses = hist.get('train_loss', [])
+            if losses:
+                plt.plot(range(1, len(losses) + 1), losses, label=m.upper(), color=colors.get(m, '#333333'), lw=2.0)
+                has_data = True
+
+    if has_data:
+        plt.xlabel("Epoch", fontsize=11, fontweight='bold')
+        plt.ylabel("Training Loss", fontsize=11, fontweight='bold')
+        plt.title(f"Training Loss Trajectory on {dataset_name.upper()} (CPU-Optimized)", fontsize=12, fontweight='bold')
+        plt.legend(frameon=True)
+        plt.tight_layout()
+        plot_path = os.path.join(save_dir, f"{dataset_name}_loss_curves.png")
+        plt.savefig(plot_path, dpi=300)
+        plt.close()
+        logger.info(f"Saved loss trajectory plot -> {plot_path}")
+
+
+def plot_parity(
+    dataset_name: str,
+    model_name: str = 'mpnn',
+    checkpoint_dir: str = 'results/checkpoints',
+    save_dir: str = 'results/plots'
+):
+    """Generate Parity Plot (Predicted vs Actual) for ESOL Regression."""
+    os.makedirs(save_dir, exist_ok=True)
+    _, _, test_loader, meta = get_dataloaders(dataset_name, batch_size=32)
+    test_data = meta['test']
+
+    model, _ = load_trained_model(model_name, dataset_name, checkpoint_dir)
+
+    if model_name == 'baseline':
+        preds = model.predict(test_data).flatten()
+        _, y_true = extract_features(test_data)
+        y_true = y_true.flatten()
+    else:
+        preds_list, targets_list = [], []
+        with torch.no_grad():
+            for batch in test_loader:
+                out = model(batch.x, batch.edge_index, edge_attr=batch.edge_attr, batch=batch.batch)
+                preds_list.append(out.view(-1).numpy())
+                targets_list.append(batch.y.view(-1).numpy())
+        preds = np.concatenate(preds_list)
+        y_true = np.concatenate(targets_list)
+
+    rmse = np.sqrt(mean_squared_error(y_true, preds))
+    r2 = r2_score(y_true, preds)
+
+    plt.figure(figsize=(6.5, 6))
+    plt.scatter(y_true, preds, alpha=0.7, color='#2b5c8f', edgecolors='k', linewidth=0.5, s=45, label='Test Compounds')
+
+    min_val = min(y_true.min(), preds.min()) - 0.5
+    max_val = max(y_true.max(), preds.max()) + 0.5
+    plt.plot([min_val, max_val], [min_val, max_val], 'r--', lw=1.8, label='Ideal 1:1 Parity')
+
+    plt.xlim(min_val, max_val)
+    plt.ylim(min_val, max_val)
+    plt.xlabel("Measured Aqueous Solubility logS (mol/L)", fontsize=11, fontweight='bold')
+    plt.ylabel(f"Predicted logS by {model_name.upper()} (mol/L)", fontsize=11, fontweight='bold')
+    plt.title(f"ESOL Parity Plot ({model_name.upper()} on Scaffold Test Split)\nRMSE: {rmse:.3f} | R²: {r2:.3f}", fontsize=12, fontweight='bold')
+    plt.legend(frameon=True)
+    plt.tight_layout()
+
+    plot_path = os.path.join(save_dir, f"{dataset_name}_parity_{model_name}.png")
+    plt.savefig(plot_path, dpi=300)
+    plt.close()
+    logger.info(f"Saved parity plot -> {plot_path}")
+
+
+def plot_roc_pr_curves(
+    dataset_name: str,
+    checkpoint_dir: str = 'results/checkpoints',
+    save_dir: str = 'results/plots'
+):
+    """Generate ROC and PR Curves for BBBP or Tox21 classification."""
+    os.makedirs(save_dir, exist_ok=True)
+    models = ['baseline', 'gcn', 'gat', 'gin', 'mpnn']
+    colors = {'baseline': '#7f7f7f', 'gcn': '#1f77b4', 'gat': '#ff7f0e', 'gin': '#2ca02c', 'mpnn': '#d62728'}
+
+    _, _, test_loader, meta = get_dataloaders(dataset_name, batch_size=32)
+    test_data = meta['test']
+
+    fig, (ax_roc, ax_pr) = plt.subplots(1, 2, figsize=(13, 5.5))
+
+    for m in models:
+        try:
+            model, _ = load_trained_model(m, dataset_name, checkpoint_dir)
+        except Exception:
+            continue
+
+        if m == 'baseline':
+            preds = model.predict(test_data)[:, 0]
+            _, y_true = extract_features(test_data)
+            y_true = y_true[:, 0]
+        else:
+            preds_list, targets_list = [], []
+            with torch.no_grad():
+                for batch in test_loader:
+                    out = torch.sigmoid(model(batch.x, batch.edge_index, edge_attr=batch.edge_attr, batch=batch.batch))
+                    preds_list.append(out[:, 0].numpy())
+                    targets_list.append(batch.y[:, 0].numpy())
+            preds = np.concatenate(preds_list)
+            y_true = np.concatenate(targets_list)
+
+        mask = ~np.isnan(y_true)
+        y_eval = y_true[mask]
+        p_eval = preds[mask]
+
+        if len(np.unique(y_eval)) > 1:
+            fpr, tpr, _ = roc_curve(y_eval, p_eval)
+            roc_auc = roc_auc_score(y_eval, p_eval)
+            prec, rec, _ = precision_recall_curve(y_eval, p_eval)
+            pr_auc = average_precision_score(y_eval, p_eval)
+
+            c = colors.get(m, '#333333')
+            ax_roc.plot(fpr, tpr, label=f"{m.upper()} (AUC = {roc_auc:.3f})", color=c, lw=2.0)
+            ax_pr.plot(rec, prec, label=f"{m.upper()} (AUC = {pr_auc:.3f})", color=c, lw=2.0)
+
+    ax_roc.plot([0, 1], [0, 1], 'k--', lw=1.2, label='Random Chance')
+    ax_roc.set_xlabel("False Positive Rate", fontsize=11, fontweight='bold')
+    ax_roc.set_ylabel("True Positive Rate", fontsize=11, fontweight='bold')
+    ax_roc.set_title(f"ROC Curves on {dataset_name.upper()} Test Split", fontsize=12, fontweight='bold')
+    ax_roc.legend(loc='lower right', frameon=True)
+
+    ax_pr.set_xlabel("Recall", fontsize=11, fontweight='bold')
+    ax_pr.set_ylabel("Precision", fontsize=11, fontweight='bold')
+    ax_pr.set_title(f"Precision-Recall Curves on {dataset_name.upper()} Test Split", fontsize=12, fontweight='bold')
+    ax_pr.legend(loc='lower left', frameon=True)
+
+    plt.tight_layout()
+    plot_path = os.path.join(save_dir, f"{dataset_name}_roc_pr.png")
+    plt.savefig(plot_path, dpi=300)
+    plt.close()
+    logger.info(f"Saved ROC/PR curves -> {plot_path}")
+
+
+def plot_tsne_embeddings(
+    dataset_name: str,
+    model_name: str = 'mpnn',
+    checkpoint_dir: str = 'results/checkpoints',
+    save_dir: str = 'results/plots'
+):
+    """Extract latent embeddings and project into 2D with t-SNE."""
+    os.makedirs(save_dir, exist_ok=True)
+    _, _, test_loader, meta = get_dataloaders(dataset_name, batch_size=32)
+
+    try:
+        model, _ = load_trained_model(model_name, dataset_name, checkpoint_dir)
+    except Exception as e:
+        logger.warning(f"Could not load {model_name} for t-SNE: {e}")
+        return
+
+    embeddings_list, targets_list = [], []
+    with torch.no_grad():
+        for batch in test_loader:
+            emb = model.get_embedding(batch.x, batch.edge_index, edge_attr=batch.edge_attr, batch=batch.batch)
+            embeddings_list.append(emb.numpy())
+            targets_list.append(batch.y[:, 0].numpy())
+
+    embeddings = np.concatenate(embeddings_list, axis=0)
+    targets = np.concatenate(targets_list, axis=0)
+
+    # Filter out NaNs if any
+    mask = ~np.isnan(targets)
+    embeddings = embeddings[mask]
+    targets = targets[mask]
+
+    # Run t-SNE
+    perplexity = min(30, max(5, len(embeddings) // 4))
+    tsne = TSNE(n_components=2, perplexity=perplexity, random_state=42)
+    embeddings_2d = tsne.fit_transform(embeddings)
+
+    plt.figure(figsize=(7, 6))
+    if dataset_name.lower() == 'esol':
+        sc = plt.scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], c=targets, cmap='viridis', s=50, alpha=0.85, edgecolors='k', linewidth=0.3)
+        cbar = plt.colorbar(sc)
+        cbar.set_label("Solubility logS (mol/L)", fontsize=10, fontweight='bold')
+    else:
+        # Binary classification
+        scatter_0 = plt.scatter(embeddings_2d[targets == 0, 0], embeddings_2d[targets == 0, 1], c='#3b82f6', s=50, alpha=0.8, edgecolors='k', linewidth=0.3, label='Negative (0)')
+        scatter_1 = plt.scatter(embeddings_2d[targets == 1, 0], embeddings_2d[targets == 1, 1], c='#ef4444', s=50, alpha=0.8, edgecolors='k', linewidth=0.3, label='Positive (1)')
+        plt.legend(frameon=True)
+
+    plt.title(f"t-SNE Latent Space ({model_name.upper()} on {dataset_name.upper()} Test Split)", fontsize=12, fontweight='bold')
+    plt.xlabel("t-SNE Dimension 1", fontsize=11, fontweight='bold')
+    plt.ylabel("t-SNE Dimension 2", fontsize=11, fontweight='bold')
+    plt.tight_layout()
+
+    plot_path = os.path.join(save_dir, f"{dataset_name}_tsne_{model_name}.png")
+    plt.savefig(plot_path, dpi=300)
+    plt.close()
+    logger.info(f"Saved t-SNE visualization -> {plot_path}")
+
+
+def generate_benchmark_report(
+    benchmark_results: Dict[str, Dict[str, Any]],
+    output_path: str = 'results/REPORT.md'
+):
+    """Compile comprehensive benchmark report across datasets and models."""
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    lines = [
+        "# Benchmark Report: Graph Neural Networks for Molecular Property Prediction",
+        "",
+        "## Executive Summary",
+        "This benchmark evaluates 5 distinct molecular property prediction architectures across 3 canonical MoleculeNet benchmarks using **strict Bemis-Murcko scaffold splitting (80/10/10)**.",
+        "All models were trained and benchmarked locally on an Intel CPU with 4 physical threads.",
+        "",
+        "---",
+        "",
+        "## Benchmark Results",
+        ""
+    ]
+
+    for dataset_name, models_dict in benchmark_results.items():
+        lines.append(f"### Dataset: `{dataset_name.upper()}`")
+        if dataset_name.lower() == 'esol':
+            lines.append("| Model Architecture | Test RMSE (mol/L) ↓ | Test MAE (mol/L) ↓ | Test R² ↑ | Train Time (s) |")
+            lines.append("| :--- | :--- | :--- | :--- | :--- |")
+            for m, data in models_dict.items():
+                agg = data['aggregated_test_metrics']
+                rmse_str = f"{agg['rmse']['mean']:.3f} ± {agg['rmse']['std']:.3f}"
+                mae_str = f"{agg['mae']['mean']:.3f} ± {agg['mae']['std']:.3f}"
+                r2_str = f"{agg['r2']['mean']:.3f} ± {agg['r2']['std']:.3f}"
+                time_str = f"{data['mean_train_time_sec']:.1f}s"
+                lines.append(f"| **{m.upper()}** | {rmse_str} | {mae_str} | {r2_str} | {time_str} |")
+        else:
+            lines.append("| Model Architecture | Test ROC-AUC ↑ | Test PR-AUC ↑ | Train Time (s) |")
+            lines.append("| :--- | :--- | :--- | :--- |")
+            for m, data in models_dict.items():
+                agg = data['aggregated_test_metrics']
+                roc_str = f"{agg['roc_auc']['mean']:.3f} ± {agg['roc_auc']['std']:.3f}"
+                pr_str = f"{agg['pr_auc']['mean']:.3f} ± {agg['pr_auc']['std']:.3f}"
+                time_str = f"{data['mean_train_time_sec']:.1f}s"
+                lines.append(f"| **{m.upper()}** | {roc_str} | {pr_str} | {time_str} |")
+        lines.append("")
+
+    lines.extend([
+        "---",
+        "",
+        "## Inductive Biases & Architectural Insights",
+        "",
+        "1. **Message Passing (MPNN)**: Leverages edge attributes (bond conjugation, ring membership, stereochemistry). Performs exceptionally well on solubility where ring strain and bond conjugation directly dictate crystal lattice energy and water solvation.",
+        "2. **Attention Mechanism (GAT)**: Dynamically re-weights atomic neighbor interactions through multi-head self-attention. Especially effective for recognizing functional pharmacophores without rigid topology assumptions.",
+        "3. **Graph Isomorphism (GIN)**: Maximally expressive for distinguishing graph topologies via multi-layer perceptron aggregations.",
+        "4. **Classical Fingerprint Baseline (Random Forest + ECFP4)**: Strong competitive baseline that serves as a sanity check. While fast, it ignores global 3D topology and continuous atom-bond interactions.",
+        "",
+        "---",
+        "",
+        "## Diagnostic Plots",
+        "- **Training Trajectories**: `results/plots/{dataset}_loss_curves.png`",
+        "- **Parity Plots**: `results/plots/esol_parity_{model}.png`",
+        "- **ROC & PR Curves**: `results/plots/{dataset}_roc_pr.png`",
+        "- **t-SNE Embeddings**: `results/plots/{dataset}_tsne_{model}.png`"
+    ])
+
+    with open(output_path, 'w') as f:
+        f.write("\n".join(lines))
+
+    logger.info(f"Report compiled successfully -> {output_path}")

src/explain.py

@@ -0,0 +1,207 @@
+"""
+Model Explainability Suite using GNNExplainer and RDKit 2D Substructure Rendering.
+Supports both single-task and multi-task models, with edge-aware and non-edge GNNs.
+"""
+import os
+import sys
+
+# Ensure project root is in sys.path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import logging
+import argparse
+from typing import Dict, List, Optional, Tuple, Any, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch_geometric.explain import Explainer, GNNExplainer, ModelConfig
+from rdkit import Chem
+from rdkit.Chem.Draw import rdMolDraw2D
+import matplotlib.pyplot as plt
+
+from src.data import smiles_to_graph
+from src.models.gnn import build_gnn_model
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger(__name__)
+
+
+class TaskSpecificWrapper(nn.Module):
+    """Wraps a multi-task GNN to isolate a single task output for GNNExplainer."""
+    def __init__(self, base_model: nn.Module, task_idx: int = 0):
+        super().__init__()
+        self.base_model = base_model
+        self.task_idx = task_idx
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        edge_index: torch.Tensor,
+        edge_attr: Optional[torch.Tensor] = None,
+        batch: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        out = self.base_model(x, edge_index, edge_attr=edge_attr, batch=batch)
+        if out.ndim == 1:
+            return out.view(-1, 1)
+        return out[:, self.task_idx:self.task_idx + 1]
+
+
+def explain_molecule(
+    model: nn.Module,
+    smiles: str,
+    task_type: str = 'regression',
+    task_idx: int = 0,
+    num_tasks: int = 1,
+    num_epochs: int = 50,
+    device: str = 'cpu'
+) -> Tuple[np.ndarray, np.ndarray, float]:
+    """
+    Explain a molecular graph prediction using GNNExplainer.
+    Returns:
+        atom_importances: np.ndarray of shape (num_atoms,)
+        edge_importances: np.ndarray of shape (num_edges,)
+        pred_value: float scalar prediction
+    """
+    model.eval()
+    graph = smiles_to_graph(smiles)
+    if graph is None:
+        raise ValueError(f"Invalid SMILES string: '{smiles}'")
+
+    graph = graph.to(device)
+
+    # Determine prediction
+    with torch.no_grad():
+        raw_pred = model(graph.x, graph.edge_index, edge_attr=graph.edge_attr)
+        if task_type != 'regression':
+            raw_pred = torch.sigmoid(raw_pred)
+        pred_scalar = float(raw_pred[0, task_idx].item()) if raw_pred.ndim > 1 else float(raw_pred[0].item())
+
+    # Wrap model if multi-task
+    if num_tasks > 1:
+        wrapped_model = TaskSpecificWrapper(model, task_idx=task_idx)
+    else:
+        wrapped_model = model
+
+    mode = 'regression' if task_type == 'regression' else 'binary_classification'
+    explainer = Explainer(
+        model=wrapped_model,
+        algorithm=GNNExplainer(epochs=num_epochs),
+        explanation_type='model',
+        node_mask_type='attributes',
+        edge_mask_type='object' if graph.edge_index.size(1) > 0 else None,
+        model_config=ModelConfig(
+            mode=mode,
+            task_level='graph',
+            return_type='raw'
+        )
+    )
+
+    explanation = explainer(
+        x=graph.x,
+        edge_index=graph.edge_index,
+        edge_attr=graph.edge_attr
+    )
+
+    # Atom importance: sum across feature attributes
+    if explanation.node_mask is not None:
+        atom_scores = explanation.node_mask.sum(dim=-1).detach().cpu().numpy()
+    else:
+        atom_scores = np.ones(graph.num_nodes, dtype=np.float32)
+
+    # Edge importance
+    if explanation.edge_mask is not None and graph.edge_index.size(1) > 0:
+        edge_scores = explanation.edge_mask.detach().cpu().numpy()
+    else:
+        edge_scores = np.zeros(graph.edge_index.size(1), dtype=np.float32)
+
+    return atom_scores, edge_scores, pred_scalar
+
+
+def render_molecule_explanation(
+    smiles: str,
+    atom_importances: np.ndarray,
+    width: int = 500,
+    height: int = 400,
+    colormap_name: str = 'Reds',
+    output_path: Optional[str] = None
+) -> bytes:
+    """
+    Render 2D molecular structure with atom-level importance heatmaps.
+    Outputs high-resolution PNG bytes and optionally saves to disk.
+    """
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        raise ValueError(f"Invalid SMILES string '{smiles}'")
+
+    num_atoms = mol.GetNumAtoms()
+    if len(atom_importances) != num_atoms:
+        atom_scores = np.resize(atom_importances, num_atoms)
+    else:
+        atom_scores = np.array(atom_importances, dtype=np.float32)
+
+    # Normalize scores between 0.15 and 0.95 for clear visual contrast
+    min_s = float(atom_scores.min())
+    max_s = float(atom_scores.max())
+    range_s = max_s - min_s if max_s > min_s else 1.0
+    norm_scores = 0.15 + 0.8 * ((atom_scores - min_s) / range_s)
+
+    cmap = plt.colormaps[colormap_name]
+    atom_colors = {}
+    radii = {}
+
+    for i in range(num_atoms):
+        rgba = cmap(float(norm_scores[i]))
+        atom_colors[i] = (float(rgba[0]), float(rgba[1]), float(rgba[2]))
+        radii[i] = 0.35 + 0.25 * float(norm_scores[i])
+
+    drawer = rdMolDraw2D.MolDraw2DCairo(width, height)
+    draw_options = drawer.drawOptions()
+    draw_options.useBWAtomPalette()
+    draw_options.scaleBondWidth = True
+    draw_options.bondLineWidth = 2
+
+    drawer.DrawMolecule(
+        mol,
+        highlightAtoms=list(atom_colors.keys()),
+        highlightAtomColors=atom_colors,
+        highlightAtomRadii=radii
+    )
+    drawer.FinishDrawing()
+    png_bytes = drawer.GetDrawingText()
+
+    if output_path is not None:
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        with open(output_path, 'wb') as f:
+            f.write(png_bytes)
+        logger.info(f"Saved explanation rendering -> {output_path}")
+
+    return png_bytes
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Explain GNN molecular property prediction")
+    parser.add_argument("--smiles", type=str, default="CC(=O)OC1=CC=CC=C1C(=O)O", help="SMILES to explain")
+    parser.add_argument("--model", type=str, default="gcn", help="Model name (gcn, gat, gin, mpnn)")
+    parser.add_argument("--dataset", type=str, default="esol", help="Dataset name (esol, bbbp, tox21)")
+    parser.add_argument("--task_idx", type=int, default=0, help="Task index to explain (0-11 for Tox21)")
+    parser.add_argument("--output", type=str, default="results/plots/explanation_demo.png", help="Output PNG path")
+    args = parser.parse_args()
+
+    from src.evaluate import load_trained_model
+    model, payload = load_trained_model(args.model, args.dataset)
+    num_tasks = payload['num_tasks'] if payload else 1
+    task_type = 'regression' if args.dataset == 'esol' else ('multitask_classification' if args.dataset == 'tox21' else 'classification')
+
+    atom_imp, edge_imp, pred = explain_molecule(
+        model=model,
+        smiles=args.smiles,
+        task_type=task_type,
+        task_idx=args.task_idx,
+        num_tasks=num_tasks,
+        num_epochs=30
+    )
+
+    print(f"Prediction for {args.smiles}: {pred:.4f}")
+    print(f"Top 3 important atom indices: {np.argsort(atom_imp)[-3:][::-1]}")
+    render_molecule_explanation(args.smiles, atom_imp, output_path=args.output)

src/models/baseline.py

@@ -0,0 +1,215 @@
+"""
+Classical Machine Learning Baselines for Molecular Property Prediction.
+Uses 2048-bit Morgan Fingerprints (ECFP4) with Random Forest and Gradient Boosting.
+"""
+import os
+import sys
+
+# Ensure project root is in sys.path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
+
+import pickle
+import logging
+from typing import Dict, List, Optional, Tuple, Any, Union
+
+import numpy as np
+import torch
+from torch_geometric.data import Data
+from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier, GradientBoostingRegressor
+from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, roc_auc_score, average_precision_score
+from rdkit import Chem
+from rdkit.Chem import Descriptors, AllChem
+
+from src.data import smiles_to_fingerprint
+
+logger = logging.getLogger(__name__)
+
+
+def extract_features(data_list: List[Data], n_bits: int = 2048, use_descriptors: bool = True) -> Tuple[np.ndarray, np.ndarray]:
+    """
+    Extract Morgan fingerprints (ECFP4) and optional physicochemical descriptors for a list of PyG Data objects.
+    Returns:
+        X: np.ndarray of shape (num_samples, num_features)
+        y: np.ndarray of shape (num_samples, num_tasks)
+    """
+    X_list = []
+    y_list = []
+
+    for data in data_list:
+        smiles = data.smiles
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            continue
+
+        # 2048-bit Morgan fingerprint
+        fp = smiles_to_fingerprint(smiles, radius=2, n_bits=n_bits)
+        if fp is None:
+            continue
+
+        if use_descriptors:
+            # Add MolWt, LogP, TPSA, NumRotatableBonds, HBD, HBA
+            desc = np.array([
+                Descriptors.MolWt(mol),
+                Descriptors.MolLogP(mol),
+                Descriptors.TPSA(mol),
+                Descriptors.NumRotatableBonds(mol),
+                Descriptors.NumHDonors(mol),
+                Descriptors.NumHAcceptors(mol)
+            ], dtype=np.float32)
+            feat = np.concatenate([fp, desc])
+        else:
+            feat = fp
+
+        X_list.append(feat)
+
+        if hasattr(data, 'y') and data.y is not None:
+            y_val = data.y.cpu().numpy()
+            if y_val.ndim == 1:
+                y_val = y_val.reshape(1, -1)
+            y_list.append(y_val.reshape(-1))
+
+    X = np.stack(X_list, axis=0) if len(X_list) > 0 else np.empty((0, n_bits))
+    y = np.stack(y_list, axis=0) if len(y_list) > 0 else np.empty((0, 1))
+
+    return X, y
+
+
+class ClassicalFingerprintBaseline:
+    """
+    Classical Machine Learning Baseline supporting Regression, Binary Classification,
+    and Multi-task Masked Classification.
+    """
+    def __init__(
+        self,
+        task_type: str = 'regression',
+        num_tasks: int = 1,
+        model_type: str = 'random_forest',
+        n_estimators: int = 100,
+        random_state: int = 42,
+        n_bits: int = 2048,
+        use_descriptors: bool = True
+    ):
+        self.task_type = task_type
+        self.num_tasks = num_tasks
+        self.model_type = model_type
+        self.n_estimators = n_estimators
+        self.random_state = random_state
+        self.n_bits = n_bits
+        self.use_descriptors = use_descriptors
+        self.models: List[Any] = []
+        self.is_fitted = False
+
+    def _create_single_model(self):
+        if self.task_type == 'regression':
+            if self.model_type == 'gradient_boosting':
+                return GradientBoostingRegressor(n_estimators=self.n_estimators, random_state=self.random_state)
+            return RandomForestRegressor(n_estimators=self.n_estimators, random_state=self.random_state, n_jobs=-1)
+        else:
+            # Classification
+            return RandomForestClassifier(n_estimators=self.n_estimators, random_state=self.random_state, n_jobs=-1)
+
+    def fit(self, train_data: List[Data]) -> 'ClassicalFingerprintBaseline':
+        """Train baseline model(s) on PyG data list."""
+        X, y = extract_features(train_data, n_bits=self.n_bits, use_descriptors=self.use_descriptors)
+        self.models = []
+
+        if self.task_type == 'regression':
+            for t in range(self.num_tasks):
+                y_t = y[:, t]
+                valid_mask = ~np.isnan(y_t)
+                model = self._create_single_model()
+                model.fit(X[valid_mask], y_t[valid_mask])
+                self.models.append(model)
+        else:
+            # Classification: Fit per task to seamlessly handle missing/NaN targets (e.g. Tox21)
+            for t in range(self.num_tasks):
+                y_t = y[:, t]
+                valid_mask = ~np.isnan(y_t)
+                model = self._create_single_model()
+                # Check if both classes are present
+                unique_classes = np.unique(y_t[valid_mask])
+                if len(unique_classes) > 1:
+                    model.fit(X[valid_mask], y_t[valid_mask].astype(int))
+                else:
+                    # Trivial case: single class in split
+                    model = None
+                self.models.append(model)
+
+        self.is_fitted = True
+        return self
+
+    def predict(self, test_data: List[Data]) -> np.ndarray:
+        """
+        Generate predictions for PyG Data objects.
+        Returns:
+            preds: np.ndarray of shape (num_samples, num_tasks)
+        """
+        if not self.is_fitted:
+            raise RuntimeError("Model is not fitted yet.")
+
+        X, _ = extract_features(test_data, n_bits=self.n_bits, use_descriptors=self.use_descriptors)
+        preds = np.zeros((len(X), self.num_tasks), dtype=np.float32)
+
+        for t in range(self.num_tasks):
+            model = self.models[t]
+            if model is None:
+                preds[:, t] = 0.5
+            elif self.task_type == 'regression':
+                preds[:, t] = model.predict(X)
+            else:
+                # Return probability of positive class
+                prob = model.predict_proba(X)
+                if prob.shape[1] > 1:
+                    preds[:, t] = prob[:, 1]
+                else:
+                    preds[:, t] = prob[:, 0]
+
+        return preds
+
+    def evaluate(self, test_data: List[Data]) -> Dict[str, float]:
+        """Evaluate fitted model on test dataset and return metric dictionary."""
+        preds = self.predict(test_data)
+        _, y_true = extract_features(test_data, n_bits=self.n_bits, use_descriptors=self.use_descriptors)
+
+        metrics = {}
+        if self.task_type == 'regression':
+            rmse_list, mae_list, r2_list = [], [], []
+            for t in range(self.num_tasks):
+                yt = y_true[:, t]
+                yp = preds[:, t]
+                mask = ~np.isnan(yt)
+                if np.sum(mask) > 1:
+                    rmse_list.append(np.sqrt(mean_squared_error(yt[mask], yp[mask])))
+                    mae_list.append(mean_absolute_error(yt[mask], yp[mask]))
+                    r2_list.append(r2_score(yt[mask], yp[mask]))
+
+            metrics['rmse'] = float(np.mean(rmse_list))
+            metrics['mae'] = float(np.mean(mae_list))
+            metrics['r2'] = float(np.mean(r2_list))
+        else:
+            roc_list, pr_list = [], []
+            for t in range(self.num_tasks):
+                yt = y_true[:, t]
+                yp = preds[:, t]
+                mask = ~np.isnan(yt)
+                # Check that both classes exist
+                if len(np.unique(yt[mask])) > 1:
+                    roc_list.append(roc_auc_score(yt[mask], yp[mask]))
+                    pr_list.append(average_precision_score(yt[mask], yp[mask]))
+
+            metrics['roc_auc'] = float(np.mean(roc_list)) if roc_list else 0.5
+            metrics['pr_auc'] = float(np.mean(pr_list)) if pr_list else 0.0
+
+        return metrics
+
+    def save(self, filepath: str):
+        """Save fitted model to disk."""
+        os.makedirs(os.path.dirname(filepath), exist_ok=True)
+        with open(filepath, 'wb') as f:
+            pickle.dump(self, f)
+
+    @classmethod
+    def load(cls, filepath: str) -> 'ClassicalFingerprintBaseline':
+        """Load fitted model from disk."""
+        with open(filepath, 'rb') as f:
+            return pickle.load(f)

src/models/gnn.py

@@ -0,0 +1,298 @@
+"""
+Modular Graph Neural Network (GNN) Architectures for Molecular Property Prediction.
+Includes GCN, GAT, GIN, and MPNN (NNConv) with dual global pooling and residual connections.
+"""
+from typing import Optional, List, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import (
+    GCNConv,
+    GATConv,
+    GINConv,
+    NNConv,
+    global_mean_pool,
+    global_add_pool
+)
+
+
+class BaseGNN(nn.Module):
+    """
+    Base class for molecular GNNs.
+    Implements multi-layer messaging, batch norm, residual connections,
+    dual pooling ([mean_pool || add_pool]), and task prediction head.
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_dim: int = 64,
+        num_layers: int = 3,
+        num_tasks: int = 1,
+        dropout: float = 0.2,
+        residual: bool = True
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.num_tasks = num_tasks
+        self.dropout_rate = dropout
+        self.residual = residual
+
+        self.convs = nn.ModuleList()
+        self.bns = nn.ModuleList()
+
+        # Dual pooling produces 2 * hidden_dim features
+        self.head = nn.Sequential(
+            nn.Linear(hidden_dim * 2, hidden_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(hidden_dim, num_tasks)
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        edge_index: torch.Tensor,
+        edge_attr: Optional[torch.Tensor] = None,
+        batch: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """Forward pass through GNN backbone, dual pooling, and task head."""
+        h = self.get_node_representations(x, edge_index, edge_attr)
+
+        if batch is None:
+            batch = torch.zeros(h.size(0), dtype=torch.long, device=h.device)
+
+        # Dual pooling: concatenation of mean and sum pooling
+        mean_p = global_mean_pool(h, batch)
+        sum_p = global_add_pool(h, batch)
+        pooled = torch.cat([mean_p, sum_p], dim=1)
+
+        return self.head(pooled)
+
+    def get_embedding(
+        self,
+        x: torch.Tensor,
+        edge_index: torch.Tensor,
+        edge_attr: Optional[torch.Tensor] = None,
+        batch: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """Extract graph-level latent representations (used for t-SNE / analysis)."""
+        h = self.get_node_representations(x, edge_index, edge_attr)
+        if batch is None:
+            batch = torch.zeros(h.size(0), dtype=torch.long, device=h.device)
+
+        mean_p = global_mean_pool(h, batch)
+        sum_p = global_add_pool(h, batch)
+        return torch.cat([mean_p, sum_p], dim=1)
+
+    def get_node_representations(
+        self,
+        x: torch.Tensor,
+        edge_index: torch.Tensor,
+        edge_attr: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """Subclasses implement convolution message passing over nodes."""
+        raise NotImplementedError
+
+
+class GCN(BaseGNN):
+    """Graph Convolutional Network (Kipf & Welling)."""
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_dim: int = 64,
+        num_layers: int = 3,
+        num_tasks: int = 1,
+        dropout: float = 0.2,
+        residual: bool = True,
+        **kwargs
+    ):
+        super().__init__(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual)
+        for i in range(num_layers):
+            dim_in = in_channels if i == 0 else hidden_dim
+            self.convs.append(GCNConv(dim_in, hidden_dim))
+            self.bns.append(nn.BatchNorm1d(hidden_dim))
+
+    def get_node_representations(
+        self,
+        x: torch.Tensor,
+        edge_index: torch.Tensor,
+        edge_attr: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        h = x
+        for i in range(self.num_layers):
+            h_new = self.convs[i](h, edge_index)
+            h_new = self.bns[i](h_new)
+            h_new = F.relu(h_new)
+            h_new = F.dropout(h_new, p=self.dropout_rate, training=self.training)
+            if self.residual and h.shape == h_new.shape:
+                h = h + h_new
+            else:
+                h = h_new
+        return h
+
+
+class GAT(BaseGNN):
+    """Graph Attention Network with Multi-Head Attention (Veličković et al.)."""
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_dim: int = 64,
+        num_layers: int = 3,
+        num_tasks: int = 1,
+        dropout: float = 0.2,
+        residual: bool = True,
+        heads: int = 4,
+        **kwargs
+    ):
+        super().__init__(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual)
+        self.heads = heads
+        assert hidden_dim % heads == 0, f"hidden_dim ({hidden_dim}) must be divisible by heads ({heads})"
+        head_dim = hidden_dim // heads
+
+        for i in range(num_layers):
+            dim_in = in_channels if i == 0 else hidden_dim
+            self.convs.append(GATConv(dim_in, head_dim, heads=heads, concat=True, dropout=dropout))
+            self.bns.append(nn.BatchNorm1d(hidden_dim))
+
+    def get_node_representations(
+        self,
+        x: torch.Tensor,
+        edge_index: torch.Tensor,
+        edge_attr: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        h = x
+        for i in range(self.num_layers):
+            h_new = self.convs[i](h, edge_index)
+            h_new = self.bns[i](h_new)
+            h_new = F.elu(h_new)
+            h_new = F.dropout(h_new, p=self.dropout_rate, training=self.training)
+            if self.residual and h.shape == h_new.shape:
+                h = h + h_new
+            else:
+                h = h_new
+        return h
+
+
+class GIN(BaseGNN):
+    """Graph Isomorphism Network with MLP aggregators (Xu et al.)."""
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_dim: int = 64,
+        num_layers: int = 3,
+        num_tasks: int = 1,
+        dropout: float = 0.2,
+        residual: bool = True,
+        train_eps: bool = True,
+        **kwargs
+    ):
+        super().__init__(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual)
+        for i in range(num_layers):
+            dim_in = in_channels if i == 0 else hidden_dim
+            mlp = nn.Sequential(
+                nn.Linear(dim_in, hidden_dim),
+                nn.BatchNorm1d(hidden_dim),
+                nn.ReLU(),
+                nn.Linear(hidden_dim, hidden_dim)
+            )
+            self.convs.append(GINConv(mlp, train_eps=train_eps))
+            self.bns.append(nn.BatchNorm1d(hidden_dim))
+
+    def get_node_representations(
+        self,
+        x: torch.Tensor,
+        edge_index: torch.Tensor,
+        edge_attr: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        h = x
+        for i in range(self.num_layers):
+            h_new = self.convs[i](h, edge_index)
+            h_new = self.bns[i](h_new)
+            h_new = F.relu(h_new)
+            h_new = F.dropout(h_new, p=self.dropout_rate, training=self.training)
+            if self.residual and h.shape == h_new.shape:
+                h = h + h_new
+            else:
+                h = h_new
+        return h
+
+
+class MPNN(BaseGNN):
+    """
+    Message Passing Neural Network using Edge Attributes (Gilmer et al.).
+    Utilizes NNConv where message passing weights are conditioned on edge features.
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        edge_dim: int,
+        hidden_dim: int = 64,
+        num_layers: int = 3,
+        num_tasks: int = 1,
+        dropout: float = 0.2,
+        residual: bool = True,
+        **kwargs
+    ):
+        super().__init__(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual)
+        self.edge_dim = edge_dim
+        self.edge_nns = nn.ModuleList()
+
+        for i in range(num_layers):
+            dim_in = in_channels if i == 0 else hidden_dim
+            edge_nn = nn.Sequential(
+                nn.Linear(edge_dim, 32),
+                nn.ReLU(),
+                nn.Linear(32, dim_in * hidden_dim)
+            )
+            self.edge_nns.append(edge_nn)
+            self.convs.append(NNConv(dim_in, hidden_dim, edge_nn, aggr='mean'))
+            self.bns.append(nn.BatchNorm1d(hidden_dim))
+
+    def get_node_representations(
+        self,
+        x: torch.Tensor,
+        edge_index: torch.Tensor,
+        edge_attr: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        # Fallback for empty edges or missing edge attributes
+        if edge_attr is None or edge_attr.size(0) == 0:
+            edge_attr = torch.zeros((edge_index.size(1), self.edge_dim), dtype=torch.float, device=x.device)
+
+        h = x
+        for i in range(self.num_layers):
+            h_new = self.convs[i](h, edge_index, edge_attr)
+            h_new = self.bns[i](h_new)
+            h_new = F.relu(h_new)
+            h_new = F.dropout(h_new, p=self.dropout_rate, training=self.training)
+            if self.residual and h.shape == h_new.shape:
+                h = h + h_new
+            else:
+                h = h_new
+        return h
+
+
+def build_gnn_model(
+    model_name: str,
+    in_channels: int,
+    edge_dim: int = 10,
+    num_tasks: int = 1,
+    hidden_dim: int = 64,
+    num_layers: int = 3,
+    dropout: float = 0.2,
+    residual: bool = True,
+    **kwargs
+) -> BaseGNN:
+    """Factory function for creating GNN models."""
+    name = model_name.lower()
+    if name == 'gcn':
+        return GCN(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual, **kwargs)
+    elif name == 'gat':
+        return GAT(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual, **kwargs)
+    elif name == 'gin':
+        return GIN(in_channels, hidden_dim, num_layers, num_tasks, dropout, residual, **kwargs)
+    elif name == 'mpnn':
+        return MPNN(in_channels, edge_dim, hidden_dim, num_layers, num_tasks, dropout, residual, **kwargs)
+    else:
+        raise ValueError(f"Unknown GNN architecture '{model_name}'. Choose from: gcn, gat, gin, mpnn")

src/train.py

@@ -0,0 +1,398 @@
+"""
+Training and Validation Pipeline for GNNs and Classical Baselines.
+Includes CPU thread limits, masked loss, early stopping, LR scheduling,
+metrics computation, and multi-seed benchmarking.
+"""
+import os
+import sys
+
+# Ensure project root is in sys.path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import time
+import json
+import random
+import logging
+import argparse
+from typing import Dict, List, Tuple, Any, Optional
+
+import yaml
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim import Adam
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from torch_geometric.loader import DataLoader
+from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score, roc_auc_score, average_precision_score
+
+from src.data import get_dataloaders, NODE_DIM, EDGE_DIM
+from src.models.gnn import build_gnn_model, BaseGNN
+from src.models.baseline import ClassicalFingerprintBaseline
+
+# Enforce 4 CPU threads matching hardware physical cores
+torch.set_num_threads(4)
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger(__name__)
+
+
+def set_seed(seed: int = 42):
+    """Set random seed for reproducibility."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def compute_loss(preds: torch.Tensor, targets: torch.Tensor, task_type: str) -> torch.Tensor:
+    """Compute task loss with NaN masking for missing multi-task targets."""
+    mask = ~torch.isnan(targets)
+    if mask.sum() == 0:
+        return torch.tensor(0.0, requires_grad=True, device=preds.device)
+
+    if task_type == 'regression':
+        return F.mse_loss(preds[mask], targets[mask])
+    else:
+        # Binary or multi-task classification
+        return F.binary_cross_entropy_with_logits(preds[mask], targets[mask])
+
+
+def evaluate_gnn(
+    model: nn.Module,
+    loader: DataLoader,
+    task_type: str,
+    device: torch.device
+) -> Tuple[Dict[str, float], np.ndarray, np.ndarray]:
+    """
+    Evaluate GNN model on a DataLoader.
+    Returns:
+        metrics: dictionary of evaluation metrics
+        all_preds: numpy array of predictions
+        all_targets: numpy array of targets
+    """
+    model.eval()
+    preds_list = []
+    targets_list = []
+
+    with torch.no_grad():
+        for batch in loader:
+            batch = batch.to(device)
+            out = model(batch.x, batch.edge_index, edge_attr=batch.edge_attr, batch=batch.batch)
+            if task_type != 'regression':
+                out = torch.sigmoid(out)
+            preds_list.append(out.cpu().numpy())
+            targets_list.append(batch.y.cpu().numpy())
+
+    preds = np.concatenate(preds_list, axis=0) if preds_list else np.empty((0, 1))
+    targets = np.concatenate(targets_list, axis=0) if targets_list else np.empty((0, 1))
+
+    num_tasks = targets.shape[1] if targets.ndim > 1 else 1
+    metrics: Dict[str, float] = {}
+
+    if task_type == 'regression':
+        rmse_list, mae_list, r2_list = [], [], []
+        for t in range(num_tasks):
+            yt = targets[:, t]
+            yp = preds[:, t]
+            mask = ~np.isnan(yt)
+            if np.sum(mask) > 1:
+                rmse_list.append(np.sqrt(mean_squared_error(yt[mask], yp[mask])))
+                mae_list.append(mean_absolute_error(yt[mask], yp[mask]))
+                r2_list.append(r2_score(yt[mask], yp[mask]))
+
+        metrics['rmse'] = float(np.mean(rmse_list)) if rmse_list else 0.0
+        metrics['mae'] = float(np.mean(mae_list)) if mae_list else 0.0
+        metrics['r2'] = float(np.mean(r2_list)) if r2_list else 0.0
+    else:
+        roc_list, pr_list = [], []
+        for t in range(num_tasks):
+            yt = targets[:, t]
+            yp = preds[:, t]
+            mask = ~np.isnan(yt)
+            if len(np.unique(yt[mask])) > 1:
+                roc_list.append(roc_auc_score(yt[mask], yp[mask]))
+                pr_list.append(average_precision_score(yt[mask], yp[mask]))
+
+        metrics['roc_auc'] = float(np.mean(roc_list)) if roc_list else 0.5
+        metrics['pr_auc'] = float(np.mean(pr_list)) if pr_list else 0.0
+
+    return metrics, preds, targets
+
+
+def train_single_seed(
+    model_name: str,
+    config: Dict[str, Any],
+    seed: int = 42,
+    quick: bool = False,
+    override_epochs: Optional[int] = None,
+    save_dir: str = 'results/checkpoints'
+) -> Dict[str, Any]:
+    """Train a single model architecture on a specific random seed."""
+    set_seed(seed)
+    device = torch.device('cpu')
+    os.makedirs(save_dir, exist_ok=True)
+    os.makedirs('results/history', exist_ok=True)
+
+    dataset_name = config['dataset']
+    task_type = config['task_type']
+    num_tasks = config['num_tasks']
+    primary_metric = config['metric']
+    higher_is_better = config.get('higher_is_better', False)
+
+    train_cfg = config.get('training', {})
+    model_cfg = config.get('model', {})
+
+    batch_size = train_cfg.get('batch_size', 32)
+    max_epochs = override_epochs or (2 if quick else train_cfg.get('epochs', 50))
+    lr = train_cfg.get('lr', 0.001)
+    weight_decay = train_cfg.get('weight_decay', 1e-4)
+    patience = train_cfg.get('patience', 15)
+    grad_clip = train_cfg.get('grad_clip', 1.0)
+
+    # 1. Classical Baseline Branch
+    if model_name.lower() == 'baseline':
+        logger.info(f"Training Classical Baseline (Random Forest) on {dataset_name.upper()} (Seed={seed})...")
+        _, _, _, meta = get_dataloaders(dataset_name, batch_size=batch_size, quick=quick)
+        baseline = ClassicalFingerprintBaseline(
+            task_type=task_type,
+            num_tasks=num_tasks,
+            n_estimators=20 if quick else 100,
+            random_state=seed
+        )
+        t0 = time.time()
+        baseline.fit(meta['train'])
+        train_time = time.time() - t0
+
+        val_metrics = baseline.evaluate(meta['val'])
+        test_metrics = baseline.evaluate(meta['test'])
+
+        ckpt_path = os.path.join(save_dir, f"{dataset_name}_baseline_best.pkl")
+        baseline.save(ckpt_path)
+
+        result = {
+            'model_name': 'baseline',
+            'dataset_name': dataset_name,
+            'seed': seed,
+            'train_time': train_time,
+            'val_metrics': val_metrics,
+            'test_metrics': test_metrics,
+            'checkpoint_path': ckpt_path
+        }
+        return result
+
+    # 2. GNN Training Branch
+    train_loader, val_loader, test_loader, meta = get_dataloaders(dataset_name, batch_size=batch_size, quick=quick)
+
+    model = build_gnn_model(
+        model_name=model_name,
+        in_channels=NODE_DIM,
+        edge_dim=EDGE_DIM,
+        num_tasks=num_tasks,
+        hidden_dim=model_cfg.get('hidden_dim', 64),
+        num_layers=model_cfg.get('num_layers', 3),
+        dropout=model_cfg.get('dropout', 0.2),
+        residual=model_cfg.get('residual', True)
+    ).to(device)
+
+    optimizer = Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
+    scheduler = ReduceLROnPlateau(
+        optimizer,
+        mode='max' if higher_is_better else 'min',
+        factor=train_cfg.get('lr_factor', 0.5),
+        patience=train_cfg.get('lr_patience', 5),
+        min_lr=1e-5
+    )
+
+    best_val_score = -float('inf') if higher_is_better else float('inf')
+    best_epoch = 0
+    patience_counter = 0
+    best_model_state = None
+    history = {'train_loss': [], 'val_metric': [], 'epoch_times': []}
+
+    t0 = time.time()
+    for epoch in range(1, max_epochs + 1):
+        ep_t0 = time.time()
+        model.train()
+        total_loss = 0.0
+        num_batches = 0
+
+        for batch in train_loader:
+            batch = batch.to(device)
+            optimizer.zero_grad()
+            out = model(batch.x, batch.edge_index, edge_attr=batch.edge_attr, batch=batch.batch)
+            loss = compute_loss(out, batch.y, task_type)
+            loss.backward()
+
+            if grad_clip > 0:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip)
+
+            optimizer.step()
+            total_loss += loss.item()
+            num_batches += 1
+
+        avg_loss = total_loss / max(num_batches, 1)
+        val_metrics, _, _ = evaluate_gnn(model, val_loader, task_type, device)
+        val_score = val_metrics[primary_metric]
+
+        scheduler.step(val_score)
+        ep_time = time.time() - ep_t0
+
+        history['train_loss'].append(avg_loss)
+        history['val_metric'].append(val_score)
+        history['epoch_times'].append(ep_time)
+
+        # Check improvement
+        is_improved = (val_score > best_val_score) if higher_is_better else (val_score < best_val_score)
+        if is_improved:
+            best_val_score = val_score
+            best_epoch = epoch
+            patience_counter = 0
+            best_model_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
+        else:
+            patience_counter += 1
+
+        if epoch % 5 == 0 or epoch == 1 or epoch == max_epochs:
+            logger.info(f"[{model_name.upper()}|Epoch {epoch:02d}/{max_epochs:02d}] Loss: {avg_loss:.4f} | Val {primary_metric.upper()}: {val_score:.4f} (Best: {best_val_score:.4f}) | Time: {ep_time:.2f}s")
+
+        if patience_counter >= patience and not quick:
+            logger.info(f"Early stopping triggered at epoch {epoch}. Best epoch was {best_epoch}.")
+            break
+
+    total_train_time = time.time() - t0
+
+    # Load best model for test evaluation
+    if best_model_state is not None:
+        model.load_state_dict(best_model_state)
+
+    val_metrics, _, _ = evaluate_gnn(model, val_loader, task_type, device)
+    test_metrics, _, _ = evaluate_gnn(model, test_loader, task_type, device)
+
+    # Save checkpoint
+    best_ckpt_path = os.path.join(save_dir, f"{dataset_name}_{model_name.lower()}_best.pt")
+    ckpt_payload = {
+        'model_name': model_name.lower(),
+        'dataset_name': dataset_name,
+        'model_state_dict': model.state_dict(),
+        'config': config,
+        'val_metrics': val_metrics,
+        'test_metrics': test_metrics,
+        'node_dim': NODE_DIM,
+        'edge_dim': EDGE_DIM,
+        'num_tasks': num_tasks,
+        'model_kwargs': {
+            'in_channels': NODE_DIM,
+            'edge_dim': EDGE_DIM,
+            'num_tasks': num_tasks,
+            'hidden_dim': model_cfg.get('hidden_dim', 64),
+            'num_layers': model_cfg.get('num_layers', 3),
+            'dropout': model_cfg.get('dropout', 0.2),
+            'residual': model_cfg.get('residual', True)
+        }
+    }
+    torch.save(ckpt_payload, best_ckpt_path)
+    logger.info(f"Saved best checkpoint to {best_ckpt_path} (Test {primary_metric.upper()}: {test_metrics[primary_metric]:.4f})")
+
+    # Save history
+    history_path = f"results/history/{dataset_name}_{model_name.lower()}_history.json"
+    with open(history_path, 'w') as f:
+        json.dump(history, f)
+
+    return {
+        'model_name': model_name.lower(),
+        'dataset_name': dataset_name,
+        'seed': seed,
+        'train_time': total_train_time,
+        'best_epoch': best_epoch,
+        'val_metrics': val_metrics,
+        'test_metrics': test_metrics,
+        'checkpoint_path': best_ckpt_path
+    }
+
+
+def train_dataset(
+    dataset_name: str,
+    models: List[str],
+    config_path: Optional[str] = None,
+    seeds: Optional[List[int]] = None,
+    quick: bool = False,
+    override_epochs: Optional[int] = None
+) -> Dict[str, Any]:
+    """Train multiple models on a given dataset across seeds, aggregating mean and std."""
+    if config_path is None:
+        config_path = f"configs/{dataset_name.lower()}.yaml"
+
+    with open(config_path, 'r') as f:
+        config = yaml.safe_load(f)
+
+    active_seeds = seeds or ([42] if quick else config.get('training', {}).get('seeds', [42, 123, 999]))
+    results_summary: Dict[str, Any] = {}
+
+    for model_name in models:
+        seed_results = []
+        for seed in active_seeds:
+            logger.info(f"\n==================================================")
+            logger.info(f"Training {model_name.upper()} on {dataset_name.upper()} | Seed {seed}")
+            logger.info(f"==================================================")
+            res = train_single_seed(
+                model_name=model_name,
+                config=config,
+                seed=seed,
+                quick=quick,
+                override_epochs=override_epochs
+            )
+            seed_results.append(res)
+
+        # Aggregate metrics across seeds
+        metric_keys = seed_results[0]['test_metrics'].keys()
+        agg_metrics = {}
+        for k in metric_keys:
+            vals = [r['test_metrics'][k] for r in seed_results]
+            agg_metrics[k] = {
+                'mean': float(np.mean(vals)),
+                'std': float(np.std(vals)),
+                'values': vals
+            }
+
+        avg_time = float(np.mean([r['train_time'] for r in seed_results]))
+
+        results_summary[model_name] = {
+            'aggregated_test_metrics': agg_metrics,
+            'mean_train_time_sec': avg_time,
+            'runs': seed_results
+        }
+
+    return results_summary
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Train GNNs and Baselines for Molecular Property Prediction")
+    parser.add_argument("--config", type=str, default=None, help="Path to YAML configuration")
+    parser.add_argument("--dataset", type=str, default="esol", help="Dataset name: esol, bbbp, tox21")
+    parser.add_argument("--model", type=str, default="all", help="Model name: baseline, gcn, gat, gin, mpnn, or all")
+    parser.add_argument("--quick", action="store_true", help="Run rapid smoke test (trimmed dataset, 2 epochs)")
+    parser.add_argument("--epochs", type=int, default=None, help="Override number of training epochs")
+    parser.add_argument("--seeds", type=int, nargs="+", default=None, help="Override seeds, e.g., --seeds 42")
+
+    args = parser.parse_args()
+
+    cfg_path = args.config if args.config is not None else f"configs/{args.dataset.lower()}.yaml"
+    all_models = ['baseline', 'gcn', 'gat', 'gin', 'mpnn']
+    models_to_train = all_models if args.model.lower() == 'all' else [args.model.lower()]
+
+    summary = train_dataset(
+        dataset_name=args.dataset,
+        models=models_to_train,
+        config_path=cfg_path,
+        seeds=args.seeds,
+        quick=args.quick,
+        override_epochs=args.epochs
+    )
+
+    print("\n" + "="*50)
+    print(f"BENCHMARK SUMMARY FOR {args.dataset.upper()}")
+    print("="*50)
+    for m, d in summary.items():
+        metrics_str = ", ".join([f"{k.upper()}: {v['mean']:.4f} ± {v['std']:.4f}" for k, v in d['aggregated_test_metrics'].items()])
+        print(f"{m.upper():<10} | {metrics_str} | Avg Time: {d['mean_train_time_sec']:.1f}s")