mirror of
https://github.com/RYDE-WORK/lnp_ml.git
synced 2026-09-18 19:13:21 +08:00
282 lines
12 KiB
Python
282 lines
12 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, 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()
|