| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import pickle |
| from collections import Counter |
| from pathlib import Path |
| from types import SimpleNamespace |
|
|
| import numpy as np |
| from scipy.stats import wilcoxon |
|
|
| from scoobdoob.validation import msm_build as mb |
| from scoobdoob.validation.selectors import ( |
| SEED_RELAX, |
| STRIDE, |
| build_selectors, |
| enrich, |
| state_means, |
| ) |
|
|
| HERE = Path(__file__).resolve().parent |
| CACHE = HERE / "data/p53/pocket_sc_feats_s10.pkl" |
| FRAMES = HERE / "data/p53/frame_drug/frames.json" |
| DEFAULT_OUT = HERE / "data/p53/trajectory_cv_tilt.json" |
|
|
|
|
| def assign_folds(meta, nfold=10): |
| """Keep every fold heterogeneous while holding out independent trajectories.""" |
| folds = [[] for _ in range(nfold)] |
| folded = [i for i, m in enumerate(meta) if m["cls"] == "folded"] |
| regular = [ |
| i |
| for i, m in enumerate(meta) |
| if m["cls"] == "seed" and not m["name"].startswith("barrier_") |
| ] |
| barrier = [ |
| i |
| for i, m in enumerate(meta) |
| if m["cls"] == "seed" and m["name"].startswith("barrier_") |
| ] |
| for group, offset in ((folded, 0), (regular, 0), (barrier, 5)): |
| for j, i in enumerate(group): |
| folds[(j + offset) % nfold].append(i) |
| return folds |
|
|
|
|
| def fit_model(data, train, k, lag, ncomp, seed): |
| feats = [data["feats"][i] for i in train] |
| cvs = [data["cvs"][i] for i in train] |
|
|
| proj, tica = mb.tica_reduce(feats, lag=lag, n_comp=ncomp, scaling="kinetic_map") |
| dtrajs, centers, km = mb.cluster_microstates(proj, k=k, seed=seed) |
| state_xy, counts = mb.state_mean_features(dtrajs, cvs, k) |
| model, active, sym2active, _ = mb.build_msm(dtrajs, lag=lag, k=k, reversible=True) |
| msm = mb.MSM( |
| model=model, |
| dtrajs=dtrajs, |
| centers=centers, |
| state_xy=state_xy, |
| counts=counts, |
| active_symbols=active, |
| lag=lag, |
| tica=tica, |
| _sym2active=sym2active, |
| ) |
|
|
| held_dtrajs = {} |
| for i in range(len(data["feats"])): |
| if i in train: |
| continue |
| x = tica.transform(np.asarray(data["feats"][i], dtype=np.float64)) |
| held_dtrajs[i] = km.predict(x).astype(np.int64) |
| return msm, held_dtrajs |
|
|
|
|
| def training_context(msm, data, train, lag): |
| """Quantities needed by ``build_selectors``, calculated on training only.""" |
| active = np.asarray(msm.active_symbols) |
| s2a = {int(s): i for i, s in enumerate(active)} |
| n = len(active) |
| cvsum, cvcnt = np.zeros((n, 2)), np.zeros(n) |
| dsum, dcnt = np.zeros(n), np.zeros(n) |
| C = np.zeros((n, n)) |
| starts = [] |
|
|
| for local, orig in enumerate(train): |
| a = np.array([s2a.get(int(s), -1) for s in msm.dtrajs[local]]) |
| cv = np.asarray(data["cvs"][orig]) |
| feat = np.asarray(data["feats"][orig]) |
| nt = min(len(a), len(cv), len(feat)) |
| a, cv, feat = a[:nt], cv[:nt], feat[:nt] |
| good = a >= 0 |
| np.add.at(cvsum, a[good], cv[good]) |
| np.add.at(cvcnt, a[good], 1) |
| np.add.at(dsum, a[good], feat[good].mean(axis=1)) |
| np.add.at(dcnt, a[good], 1) |
| if data["meta"][orig]["cls"] == "folded" and good[0]: |
| starts.append(int(a[0])) |
| pairs = (a[:-lag] >= 0) & (a[lag:] >= 0) |
| np.add.at(C, (a[:-lag][pairs], a[lag:][pairs]), 1) |
|
|
| if not starts: |
| raise RuntimeError("training split has no active folded start state") |
| cvm = cvsum / np.maximum(cvcnt[:, None], 1) |
| dilation = np.full(n, np.nan) |
| nz = dcnt > 0 |
| dilation[nz] = dsum[nz] / dcnt[nz] |
| a0 = Counter(starts).most_common(1)[0][0] |
| return cvm, C, a0, dilation |
|
|
|
|
| def frame_table(data, held_dtrajs, held, args): |
| """Map held-out fpocket frame records into the held-out model labels.""" |
| by_orig = {int(m.get("orig_traj", i)): i for i, m in enumerate(data["meta"])} |
| held = set(held) |
| rows = json.loads(FRAMES.read_text()) |
| out_state, out_drug, out_traj = [], [], [] |
|
|
| for r in rows: |
| ci = by_orig.get(int(r["traj"])) |
| if ci not in held: |
| continue |
| value = r.get(args.observable) |
| if value is None: |
| continue |
| core = r.get("core_rmsd_nm") |
| zn = r.get("zn_n_coord", -1) |
| if core is None or core > args.max_core_rmsd: |
| continue |
| if zn >= 0 and zn < args.min_zn: |
| continue |
| off = SEED_RELAX if data["meta"][ci]["cls"] == "seed" else 0 |
| pos = int(r["frame"]) // STRIDE - off |
| dt = held_dtrajs[ci] |
| if not 0 <= pos < len(dt): |
| continue |
| out_state.append(int(dt[pos])) |
| out_drug.append(float(value)) |
| out_traj.append(ci) |
| return (np.asarray(out_state), np.asarray(out_drug), np.asarray(out_traj)) |
|
|
|
|
| def percentile_ci(x): |
| if len(x) < 2: |
| return [float("nan"), float("nan")] |
| return np.percentile(x, [2.5, 97.5]).tolist() |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--k", type=int, default=100) |
| ap.add_argument("--lag", type=int, default=1) |
| ap.add_argument("--ncomp", type=int, default=6) |
| ap.add_argument("--seed", type=int, default=0) |
| ap.add_argument("--target-q", type=float, default=0.15) |
| ap.add_argument( |
| "--observable", |
| default="site_drug", |
| choices=["site_drug", "site_volume", "site_apolar", "site_sasa"], |
| ) |
| ap.add_argument("--drug-thresh", type=float, default=0.5) |
| ap.add_argument( |
| "--alpha-rho-sweep", |
| type=float, |
| nargs="*", |
| default=[0.0, 0.1, 0.25, 0.36, 0.5, 1.0, 1.5, 2.0, 3.0, 5.0, 8.0, 12.0, 20.0], |
| help="Algorithm-1 tilt strengths to evaluate (w=exp(-alpha_rho*rho), no " |
| "committor shaping). alpha_rho=0 is the no-density-tilt control. The " |
| "tilt strength is a knob; report the sweep, not one value.", |
| ) |
| ap.add_argument("--max-core-rmsd", type=float, default=0.30) |
| ap.add_argument("--min-zn", type=int, default=3) |
| ap.add_argument("--out", type=Path, default=DEFAULT_OUT) |
| args = ap.parse_args() |
|
|
| data = pickle.load(open(CACHE, "rb")) |
| folds = assign_folds(data["meta"]) |
| all_idx = set(range(len(data["meta"]))) |
| result = {"config": vars(args) | {"out": str(args.out)}, "folds": []} |
|
|
| selector_names = ( |
| "equilibrium", |
| "uniform_states", |
| "tpt_flux", |
| "doob_untilted", |
| "fk_doob", |
| ) + tuple(f"alg1[a={a:g}]" for a in args.alpha_rho_sweep) |
| for fi, held in enumerate(folds): |
| train = sorted(all_idx - set(held)) |
| print(f"[fold {fi}] train={len(train)} held={len(held)}") |
| msm, held_dt = fit_model(data, train, args.k, args.lag, args.ncomp, args.seed) |
| cvm, C, a0, dilation = training_context(msm, data, train, args.lag) |
|
|
| fit_args = SimpleNamespace( |
| target="dilation", |
| target_q=args.target_q, |
| cleft_thr=0.45, |
| horizon=None, |
| fk_variant="rho", |
| comm_l_sweep=[], |
| horizon_sweep=[], |
| alpha_rho_sweep=args.alpha_rho_sweep, |
| ) |
| selectors, q, B, a0, horizon = build_selectors( |
| msm, cvm, C, a0, fit_args, dil=dilation |
| ) |
|
|
| orig_state, drug, traj = frame_table(data, held_dt, held, args) |
| s2a = {int(s): i for i, s in enumerate(msm.active_symbols)} |
| state = np.array([s2a.get(int(s), -1) for s in orig_state]) |
| keep = state >= 0 |
| state, drug, traj = state[keep], drug[keep], traj[keep] |
| n = msm.n_active |
| mu, cnt = state_means(state, drug, n) |
| po, _ = state_means(state, (drug > args.drug_thresh).astype(float), n) |
|
|
| eq = enrich( |
| selectors["equilibrium"], |
| mu, |
| cnt, |
| po, |
| drug.mean(), |
| (drug > args.drug_thresh).mean(), |
| ) |
| fold = { |
| "fold": fi, |
| "held_trajectories": held, |
| "held_names": [data["meta"][i]["name"] for i in held], |
| "n_frames": int(len(drug)), |
| "n_active": int(n), |
| "state_coverage": int((cnt > 0).sum()), |
| "B_size": int(len(B)), |
| "a0": int(a0), |
| "horizon": int(horizon), |
| "selectors": {}, |
| } |
| for name in selector_names: |
| if name not in selectors: |
| continue |
| e = enrich( |
| selectors[name], |
| mu, |
| cnt, |
| po, |
| drug.mean(), |
| (drug > args.drug_thresh).mean(), |
| ) |
| if e is None: |
| continue |
| e["ratio_to_equilibrium"] = e["E_drug"] / eq["E_drug"] |
| e["open_ratio_to_equilibrium"] = e["P_open"] / eq["P_open"] |
| fold["selectors"][name] = e |
| result["folds"].append(fold) |
| fk = fold["selectors"].get("fk_doob", {}) |
| print( |
| f" n={len(drug)} coverage={fk.get('coverage', np.nan):.2f} " |
| f"FK/equilibrium={fk.get('ratio_to_equilibrium', np.nan):.3f}" |
| ) |
|
|
| summary = {} |
| for name in selector_names: |
| vals = [ |
| f["selectors"][name]["ratio_to_equilibrium"] |
| for f in result["folds"] |
| if name in f["selectors"] |
| ] |
| opens = [ |
| f["selectors"][name]["open_ratio_to_equilibrium"] |
| for f in result["folds"] |
| if name in f["selectors"] |
| ] |
| if not vals: |
| continue |
| arr = np.asarray(vals) |
| if name == "equilibrium": |
| p = 1.0 |
| else: |
| try: |
| p = float(wilcoxon(arr - 1, alternative="two-sided").pvalue) |
| except ValueError: |
| p = 1.0 |
| summary[name] = { |
| "n_folds": len(vals), |
| "mean_ratio_to_equilibrium": float(arr.mean()), |
| "median_ratio_to_equilibrium": float(np.median(arr)), |
| "fold_percentile_interval": percentile_ci(arr), |
| "folds_above_1": int((arr > 1).sum()), |
| "wilcoxon_p_vs_1": p, |
| "mean_open_ratio_to_equilibrium": float(np.mean(opens)), |
| "values": vals, |
| } |
| result["summary"] = summary |
| args.out.write_text(json.dumps(result, indent=2, default=float)) |
|
|
| print("\nTRAJECTORY-HELD-OUT RESULT") |
| for name, s in summary.items(): |
| print( |
| f" {name:<18} {s['mean_ratio_to_equilibrium']:.3f}x " |
| f"(median {s['median_ratio_to_equilibrium']:.3f}; " |
| f"{s['folds_above_1']}/{s['n_folds']} folds >1; " |
| f"p={s['wilcoxon_p_vs_1']:.3g})" |
| ) |
| print(f"-> {args.out}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
