"""经典基线：RandomForest / Tanimoto-kNN / Tanimoto+tabular-kNN。 用法: python scripts/run_classical_baselines.py \ --input data/interim/internal.csv \ --ref-run models/abl_full/s1_baseline/seed42 \ --seed 42 --n-outer 5 --n-inner 3 --out-root models/abl_full """ from __future__ import annotations import argparse, json from pathlib import Path import numpy as np import pandas as pd from scipy.special import rel_entr from sklearn.ensemble import RandomForestRegressor, RandomForestClassifier from sklearn.model_selection import KFold, StratifiedKFold from sklearn.preprocessing import StandardScaler from sklearn.metrics import (mean_squared_error, mean_absolute_error, r2_score, accuracy_score, precision_score, recall_score, f1_score, confusion_matrix) from lnp_ml.dataset import LNPDataset from lnp_ml.featurization.smiles import RDKitFeaturizer SEED = 42 X_RF = MORGAN = TAB = None # 全局特征，main 里填充 # ---------- 指标（对齐 evaluate_on_test）---------- def reg_metrics(t, p): t, p = np.asarray(t), np.asarray(p) return {"n_samples": int(len(p)), "mse": float(mean_squared_error(t, p)), "rmse": float(np.sqrt(mean_squared_error(t, p))), "mae": float(mean_absolute_error(t, p)), "r2": float(r2_score(t, p))} def clf_metrics(t, p, n_classes): """Classification metrics over the declared label space. Supplying ``labels`` prevents a class absent from one test fold from being silently dropped from the macro average. Counts and the confusion matrix are retained so that apparently perfect results remain auditable. """ t, p = np.asarray(t, dtype=int), np.asarray(p, dtype=int) labels = np.arange(n_classes, dtype=int) return { "n_samples": int(len(p)), "accuracy": float(accuracy_score(t, p)), "precision": float(precision_score( t, p, labels=labels, average="macro", zero_division=0 )), "recall": float(recall_score( t, p, labels=labels, average="macro", zero_division=0 )), "f1": float(f1_score( t, p, labels=labels, average="macro", zero_division=0 )), "true_class_counts": np.bincount(t, minlength=n_classes).astype(int).tolist(), "pred_class_counts": np.bincount(p, minlength=n_classes).astype(int).tolist(), "confusion_matrix": confusion_matrix(t, p, labels=labels).astype(int).tolist(), } def dist_metrics(t, p, eps=1e-10): t = np.clip(np.asarray(t), eps, 1.0); p = np.clip(np.asarray(p), eps, 1.0) kl = float(np.sum(rel_entr(t, p), axis=-1).mean()) m = 0.5 * (t + p) js = float((0.5*np.sum(rel_entr(t, m), axis=-1) + 0.5*np.sum(rel_entr(p, m), axis=-1)).mean()) return {"n_samples": int(len(p)), "kl_divergence": kl, "js_divergence": js} METRICS = {"reg": reg_metrics, "dist": dist_metrics} def score(task_type, t, p, n_classes=None): if task_type == "reg": return r2_score(t, p) if task_type == "clf": labels = np.arange(n_classes, dtype=int) return f1_score(t, p, labels=labels, average="macro", zero_division=0) return -dist_metrics(t, p)["js_divergence"] # ---------- 相似度 / 距离 ---------- def tanimoto_sim(A, B): inter = A @ B.T a = A.sum(1)[:, None]; b = B.sum(1)[None, :] return inter / np.clip(a + b - inter, 1e-10, None) def combined_dist(fp_te, fp_tr, tab_te, tab_tr, alpha): d_fp = 1.0 - tanimoto_sim(fp_te, fp_tr) sc = StandardScaler().fit(tab_tr) Ztr, Zte = sc.transform(tab_tr), sc.transform(tab_te) d_tab = np.sqrt(np.clip(((Zte[:, None, :] - Ztr[None, :, :]) ** 2).sum(-1), 0, None)) def mm(d): lo, hi = float(d.min()), float(d.max()) return (d - lo) / (hi - lo) if hi > lo else np.zeros_like(d) return alpha * mm(d_fp) + (1 - alpha) * mm(d_tab) def knn_from_matrix(S, is_sim, ytr, k, task_type, n_classes): k = min(k, S.shape[1]) order = np.argsort(-S if is_sim else S, axis=1)[:, :k] out = [] for i, idx in enumerate(order): w = S[i, idx] if is_sim else 1.0 / (S[i, idx] + 1e-6) w = np.clip(w, 1e-12, None) yy = ytr[idx] if task_type == "reg": out.append(float(np.average(yy, weights=w))) elif task_type == "clf": sc = np.zeros(n_classes) for c, wt in zip(yy, w): sc[int(c)] += wt out.append(int(sc.argmax())) else: v = np.average(yy, axis=0, weights=w); s = v.sum() out.append(v / s if s > 0 else v) return np.array(out) # ---------- 各模型预测 ---------- # Use separate RF search spaces by task type. Classification endpoints share # one small-sample regularized configuration, while regression and # biodistribution recover the original inner-CV grid. Features and outer-fold # partitions are unchanged. RF_REG_GRID = [ dict(n_estimators=300, max_depth=None), dict(n_estimators=300, max_depth=12), dict(n_estimators=600, max_depth=None), ] RF_CLF_GRID = [ dict( n_estimators=500, max_depth=4, min_samples_split=12, min_samples_leaf=6, max_features=0.25, bootstrap=True, max_samples=0.70, ) ] K_GRID = [3, 5, 10, 15] ALPHA_GRID = [0.3, 0.5, 0.7] def rf_predict(tr, te, y, task_type, nc, params): Xtr, Xte, ytr = X_RF[tr], X_RF[te], y[tr] if task_type == "reg": m = RandomForestRegressor(random_state=SEED, n_jobs=-1, **params); m.fit(Xtr, ytr); return m.predict(Xte) if task_type == "clf": m = RandomForestClassifier( random_state=SEED, n_jobs=-1, class_weight="balanced_subsample", **params, ) m.fit(Xtr, ytr) return m.predict(Xte) m = RandomForestRegressor(random_state=SEED, n_jobs=-1, **params); m.fit(Xtr, ytr) v = np.clip(m.predict(Xte), 0, None); s = v.sum(1, keepdims=True); return np.where(s > 0, v / s, v) def knn_predict(variant, tr, te, y, task_type, nc, params): if variant == "tanimoto_knn": S, is_sim = tanimoto_sim(MORGAN[te], MORGAN[tr]), True else: S, is_sim = combined_dist(MORGAN[te], MORGAN[tr], TAB[te], TAB[tr], params.get("alpha", 0.5)), False return knn_from_matrix(S, is_sim, y[tr], params["k"], task_type, nc) def grid_for(model, task_type): if model == "rf": grid = RF_CLF_GRID if task_type == "clf" else RF_REG_GRID return [{"rf": g} for g in grid] if model == "tanimoto_knn": return [{"k": k} for k in K_GRID] return [{"k": k, "alpha": a} for k in K_GRID for a in ALPHA_GRID] def predict(model, tr, te, y, task_type, nc, params): return rf_predict(tr, te, y, task_type, nc, params["rf"]) if model == "rf" \ else knn_predict(model, tr, te, y, task_type, nc, params) def select_params(model, tr, y, valid, task_type, nc, n_inner): cand = grid_for(model, task_type) if len(cand) == 1: return cand[0] base = tr[valid[tr]] if len(base) < 4: return cand[0] if task_type == "clf": class_counts = np.bincount(y[base].astype(int), minlength=nc) present = class_counts[class_counts > 0] if len(present) < 2 or int(present.min()) < 2: return cand[0] n_splits = min(n_inner, int(present.min())) else: n_splits = max(2, min(n_inner, len(base) // 2)) best, best_s = cand[0], -1e18 for p in cand: scs = [] # Recreate the deterministic split iterator for every candidate. if task_type == "clf": split_iter = StratifiedKFold( n_splits=n_splits, shuffle=True, random_state=SEED ).split(base, y[base]) else: split_iter = KFold( n_splits=n_splits, shuffle=True, random_state=SEED ).split(base) for itr, iva in split_iter: a, b = base[itr], base[iva] if len(a) < 2 or len(b) < 1: continue try: scs.append(score( task_type, y[b], predict(model, a, b, y, task_type, nc, p), nc, )) except Exception: pass if scs and float(np.mean(scs)) > best_s: best_s, best = float(np.mean(scs)), p return best def main(): global SEED, X_RF, MORGAN, TAB ap = argparse.ArgumentParser() ap.add_argument("--input", default="data/interim/internal.csv") ap.add_argument("--ref-run", default="models/abl_full/s1_baseline/seed42") ap.add_argument("--seed", type=int, default=42) ap.add_argument("--n-outer", type=int, default=5) ap.add_argument("--n-inner", type=int, default=3) ap.add_argument("--out-root", default="models/abl_full") ap.add_argument("--models", nargs="+", default=["rf", "tanimoto_knn", "tanimoto_knn_combined"]) args = ap.parse_args() SEED = args.seed ds = LNPDataset(pd.read_csv(args.input)) MORGAN = RDKitFeaturizer().transform(ds.smiles)["morgan"].astype(np.float32) TAB = np.concatenate([ds.comp, ds.phys, ds.help, ds.exp], axis=1).astype(np.float32) X_RF = np.concatenate([MORGAN, TAB], axis=1).astype(np.float32) tasks = {} def add(name, ttype, y, valid, nc=None): if y is not None: tasks[name] = (ttype, np.asarray(y), np.asarray(valid), nc) add("size", "reg", ds.size, ~np.isnan(ds.size) if ds.size is not None else None) add("delivery", "reg", ds.delivery, ~np.isnan(ds.delivery) if ds.delivery is not None else None) if ds.toxic is not None: add("toxic", "clf", ds.toxic, ds.toxic >= 0, 2) if ds.pdi is not None: add("pdi", "clf", ds.pdi, ds.pdi_valid, int(ds.pdi[ds.pdi_valid].max()) + 1) if ds.ee is not None: add("ee", "clf", ds.ee, ds.ee_valid, int(ds.ee[ds.ee_valid].max()) + 1) if ds.biodist is not None: add("biodist", "dist", ds.biodist, ds.biodist_valid) folds = [] for k in range(args.n_outer): d = json.loads((Path(args.ref_run) / f"outer_fold_{k}" / "splits.json").read_text()) folds.append((np.array(d["outer_train_idx"]), np.array(d["outer_test_idx"]))) for model in args.models: fold_results, agg = [], {} for k, (tr, te) in enumerate(folds): tm = {} for name, (ttype, y, valid, nc) in tasks.items(): trv, tev = tr[valid[tr]], te[valid[te]] if len(trv) < 2 or len(tev) < 1: continue params = select_params(model, tr, y, valid, ttype, nc, args.n_inner) pred = predict(model, trv, tev, y, ttype, nc, params) if ttype == "clf": m = clf_metrics(y[tev], pred, nc) else: m = METRICS[ttype](y[tev], pred) tm[name] = m for mk, mv in m.items(): if mk == "n_samples" or not np.isscalar(mv): continue agg.setdefault(name, {}).setdefault(mk, []).append(mv) fold_results.append({"fold": k, "test_metrics": tm}) fdir = Path(args.out_root) / model / f"seed{args.seed}" / f"outer_fold_{k}" fdir.mkdir(parents=True, exist_ok=True) (fdir / "test_metrics.json").write_text(json.dumps(tm, indent=2)) summary_stats = {} for t, md in agg.items(): summary_stats[t] = {} for mk, v in md.items(): summary_stats[t][f"{mk}_mean"] = float(np.mean(v)) summary_stats[t][f"{mk}_std"] = float(np.std(v)) run_dir = Path(args.out_root) / model / f"seed{args.seed}" run_dir.mkdir(parents=True, exist_ok=True) (run_dir / "summary.json").write_text( json.dumps({"fold_results": fold_results, "summary_stats": summary_stats}, indent=2)) print(f"[{model}] saved -> {run_dir / 'summary.json'}") if __name__ == "__main__": main()