mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-09-18 16:23:20 +08:00
206 lines
9.2 KiB
Python
206 lines
9.2 KiB
Python
"""经典基线：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
|
||
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)
|
||
|
||
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):
|
||
return {"n_samples": int(len(p)), "accuracy": float(accuracy_score(t, p)),
|
||
"precision": float(precision_score(t, p, average="macro", zero_division=0)),
|
||
"recall": float(recall_score(t, p, average="macro", zero_division=0)),
|
||
"f1": float(f1_score(t, p, average="macro", zero_division=0))}
|
||
|
||
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, "clf": clf_metrics, "dist": dist_metrics}
|
||
|
||
def score(task_type, t, p):
|
||
if task_type == "reg": return r2_score(t, p)
|
||
if task_type == "clf": return f1_score(t, p, 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)
|
||
|
||
# ---------- 各模型预测 ----------
|
||
RF_GRID = [dict(n_estimators=300, max_depth=None),
|
||
dict(n_estimators=300, max_depth=12),
|
||
dict(n_estimators=600, max_depth=None)]
|
||
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, **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):
|
||
if model == "rf": return [{"rf": g} for g in RF_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)
|
||
if len(cand) == 1: return cand[0]
|
||
n_splits = max(2, min(n_inner, len(tr) // 2))
|
||
kf = KFold(n_splits=n_splits, shuffle=True, random_state=SEED)
|
||
best, best_s = cand[0], -1e18
|
||
for p in cand:
|
||
scs = []
|
||
for itr, iva in kf.split(tr):
|
||
a, b = tr[itr], tr[iva]
|
||
a, b = a[valid[a]], b[valid[b]]
|
||
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)))
|
||
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)
|
||
m = METRICS[ttype](y[tev], predict(model, trv, tev, y, ttype, nc, params))
|
||
tm[name] = m
|
||
for mk, mv in m.items():
|
||
if mk == "n_samples": 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() |