| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import matplotlib |
|
|
| matplotlib.use("Agg") |
|
|
| from scoobdoob.validation import msm_build as mb |
| from scoobdoob.validation import tpt, bridge, aib9 |
|
|
| ROOT = Path(__file__).resolve().parent |
| DEFAULT_DATA = aib9.OPENMM_DIR |
| OUTDIR = ROOT / "data" / "validation_aib9" |
|
|
| LAG = 5 |
| K = 100 |
| N_COMP = 4 |
| HORIZON = 60 |
| TOL_Q = 0.10 |
|
|
|
|
| def gate_msm(feat, hand, temps): |
| print("\n" + "=" * 70) |
| print("PART 2.0/2.1 Build Aib9 MSM (36D sin/cos -> TICA -> connected -> rev MLE)") |
| print("=" * 70) |
| model = mb.build_full_msm( |
| feat, hand, lag=LAG, k=K, n_comp=N_COMP, seed=0, use_tica=True |
| ) |
| n_off = K - len(model.active_symbols) |
| lam1 = float(np.real(model.model.eigenvalues(k=1)[0])) |
| print( |
| f" microstates K={K}, connected set={model.n_active} " |
| f"({n_off} dropped), lambda_1={lam1:.5f}" |
| ) |
| lags = [2, 5, 10, 20, 40] |
| lg, its = mb.implied_timescales(model.dtrajs, lags, k=K, n_its=3) |
| ck = mb.ck_test(model.dtrajs, lag=LAG, k=K, factors=(2, 3)) |
| print( |
| f" ITS(tau*={LAG}, frames) slowest 3 = " |
| f"{np.array2string(its[lags.index(LAG)], precision=1)}" |
| ) |
| print(f" CK error: 2tau={ck[2]:.3f}, 3tau={ck[3]:.3f}") |
| return model, { |
| "n_active": int(model.n_active), |
| "n_dropped": int(n_off), |
| "lambda1": lam1, |
| "ck": {int(k): v for k, v in ck.items()}, |
| } |
|
|
|
|
| def gate_endpoints(model): |
| print("\n" + "=" * 70) |
| print("GATE 2.2 Endpoints R / L from helical handedness") |
| print("=" * 70) |
| labels = aib9.classify_microstates(model.active_xy) |
| sets = { |
| k: np.where(labels == k)[0].tolist() for k in ["R", "L", "I_N", "I_C", "other"] |
| } |
| print( |
| f" microstate labels: R={len(sets['R'])} L={len(sets['L'])} " |
| f"I_N={len(sets['I_N'])} I_C={len(sets['I_C'])} other={len(sets['other'])}" |
| ) |
| pL_R = model.active_xy[sets["R"]][:, aib9.CORE].mean() |
| pL_L = model.active_xy[sets["L"]][:, aib9.CORE].mean() |
| print( |
| f" R: core P(left)={pL_R:.2f} (right helix); " |
| f"L: core P(left)={pL_L:.2f} (left helix)" |
| ) |
|
|
| A2, B2, _ = mb.q2_sign_endpoints(model, k_each=6) |
| a2_isR = np.mean([labels[i] in ("R",) for i in A2]) |
| b2_isL = np.mean([labels[i] in ("L",) for i in B2]) |
| print( |
| f" q2-sign split agreement with R/L handedness: " |
| f"{max(a2_isR, 1 - a2_isR) * 100:.0f}% / {max(b2_isL, 1 - b2_isL) * 100:.0f}%" |
| ) |
| return labels, sets |
|
|
|
|
| def gate_committor(model, sets, labels): |
| print("\n" + "=" * 70) |
| print("GATE 2.3 Committor-1/2 vs KNOWN intermediates I_N, I_C (THE GATE)") |
| print("=" * 70) |
| R, L = sets["R"], sets["L"] |
| agree, q, _ = tpt.committor_agreement(model, R, L) |
| print(f" committor solver cross-check (deeptime vs linsolve): max|dq|={agree:.2e}") |
| print(f" q on R ~ {q[R].mean():.3f}, q on L ~ {q[L].mean():.3f}") |
|
|
| half = tpt.committor_half_states(q, tol=TOL_Q) |
| inter = set(sets["I_N"]) | set(sets["I_C"]) |
| half_in_inter = [h for h in half if h in inter] |
| frac_states = len(half_in_inter) / max(len(half), 1) |
| pi = model.stationary |
| mass_frac = ( |
| pi[half_in_inter].sum() / max(pi[list(half)].sum(), 1e-12) if len(half) else 0.0 |
| ) |
|
|
| in_basin = [h for h in half if labels[h] in ("R", "L")] |
| qIN = q[sets["I_N"]].mean() if sets["I_N"] else np.nan |
| qIC = q[sets["I_C"]].mean() if sets["I_C"] else np.nan |
| n_half_IN = sum(h in set(sets["I_N"]) for h in half) |
| n_half_IC = sum(h in set(sets["I_C"]) for h in half) |
| print( |
| f" committor-1/2 ensemble: {len(half)} states " |
| f"({n_half_IN} in I_N, {n_half_IC} in I_C)" |
| ) |
| print( |
| f" overlap with intermediates: {frac_states * 100:.0f}% of states, " |
| f"{mass_frac * 100:.0f}% of committor-1/2 mass" |
| ) |
| print(f" committor-1/2 states inside an endpoint (R/L) basin: {len(in_basin)}") |
| sym = "symmetric" if abs(qIN - qIC) < 0.1 else "asymmetric" |
| print( |
| f" mean committor of I_N={qIN:.2f}, I_C={qIC:.2f} " |
| f"-> the two channels are {sym}. Aib is achiral, so a symmetric split " |
| f"is the physically expected result." |
| ) |
| passed = (agree < 1e-6) and (mass_frac >= 0.5) and (len(in_basin) == 0) |
| print( |
| f" -> {'PASS' if passed else 'FAIL'} " |
| f"(committor-1/2 coincides with the known intermediates)" |
| ) |
| return ( |
| passed, |
| { |
| "committor_agreement": agree, |
| "n_half": int(len(half)), |
| "overlap_frac_states": frac_states, |
| "overlap_frac_mass": float(mass_frac), |
| "n_half_in_basin": len(in_basin), |
| "q_IN": float(qIN), |
| "q_IC": float(qIC), |
| "pass": bool(passed), |
| }, |
| q, |
| half, |
| ) |
|
|
|
|
| def gate_flux(model, sets, q, half): |
| print("\n" + "=" * 70) |
| print("GATE 2.4 Reactive flux + bridge rollouts route through intermediates") |
| print("=" * 70) |
| R, L = sets["R"], sets["L"] |
| IN, IC = set(sets["I_N"]), set(sets["I_C"]) |
|
|
| rf = tpt.reactive_flux(model, R, L) |
| order = ["R", "I_N", "I_C", "other", "L"] |
| _, cg = rf.coarse_grain([sets[k] for k in order]) |
| F = np.asarray(cg.net_flux) |
| tot = float(cg.total_flux) |
| f_IN = float(F[0, 1] / tot) |
| f_IC = float(F[0, 2] / tot) |
| f_other = float(F[0, 3] / tot) |
| f_direct = float(F[0, 4] / tot) |
| f_thru = f_IN + f_IC + f_other |
| print(f" total reactive flux = {tot:.3e}") |
| print( |
| f" flux routed through resolved intermediates: {f_thru * 100:.0f}% " |
| f"(direct R->L: {f_direct * 100:.0f}% -- small once swap artifacts removed)" |
| ) |
| print( |
| f" TWO CHANNELS: flip-from-N (I_N) {f_IN * 100:.0f}% | " |
| f"flip-from-C (I_C) {f_IC * 100:.0f}% | middle/other {f_other * 100:.0f}%" |
| ) |
|
|
| band = set(np.where((q > 0.2) & (q < 0.8))[0].tolist()) |
| Ph = bridge.doob_kernels(model.P, L, HORIZON) |
| paths = bridge.sample_bridge_paths(Ph, R, n_paths=300, seed=0) |
| succ = bridge.success_with_ci(paths, L, seed=0) |
| visit_band = np.mean([bool(set(map(int, p)) & band) for p in paths]) |
| visit_IN = np.mean([bool(set(map(int, p)) & IN) for p in paths]) |
| visit_IC = np.mean([bool(set(map(int, p)) & IC) for p in paths]) |
| print( |
| f" bridge rollouts R->L: success {succ['rate'] * 100:.0f}% " |
| f"CI[{succ['ci'][0] * 100:.0f},{succ['ci'][1] * 100:.0f}], " |
| f"cross committor band {visit_band * 100:.0f}%" |
| ) |
| print( |
| f" rollouts touching I_N {visit_IN * 100:.0f}%, I_C {visit_IC * 100:.0f}% " |
| f"(both channels sampled)" |
| ) |
|
|
| passed = (f_IN > 0 and f_IC > 0) and (succ["rate"] > 0.9) |
| print( |
| f" -> {'PASS' if passed else 'FAIL'} " |
| f"(multi-channel reactive ensemble through both known intermediates)" |
| ) |
| return ( |
| passed, |
| { |
| "total_flux": tot, |
| "flux_through_resolved_inter": f_thru, |
| "flux_direct": f_direct, |
| "flux_IN": f_IN, |
| "flux_IC": f_IC, |
| "flux_other": f_other, |
| "rollout_success": succ["rate"], |
| "rollout_cross_band": float(visit_band), |
| "rollout_visit_IN": float(visit_IN), |
| "rollout_visit_IC": float(visit_IC), |
| "pass": bool(passed), |
| }, |
| paths, |
| ) |
|
|
|
|
| def run_robustness(feat, hand): |
| """Opt-in K x lag + ITS-convergence robustness report (gate stability).""" |
| from scoobdoob.validation import robustness as rob |
|
|
| print("\n" + "=" * 70) |
| print("ROBUSTNESS Are the conclusions stable to K and lag?") |
| print("=" * 70) |
| Ks, lags = [75, 100, 150, 200], [2, 5, 10, 20] |
| results = rob.sweep(feat, hand, Ks, lags, n_comp=N_COMP) |
| print( |
| f" {'K':>4} {'lag':>4} {'active':>6} {'lam1':>5} {'CK2t':>5} " |
| f"{'chalf%':>7} {'direct%':>7} {'I_N%':>5} {'I_C%':>5}" |
| ) |
| for r in results: |
| if r.get("ok"): |
| print( |
| f" {r['K']:4d} {r['lag']:4d} {r['n_active']:6d} {r['lam1']:5.2f} " |
| f"{r['ck2']:5.2f} {r['committor_half_mass'] * 100:6.0f}% " |
| f"{r['direct'] * 100:6.0f}% {r['flux_IN'] * 100:5.0f} {r['flux_IC'] * 100:5.0f}" |
| ) |
| summ = rob.summarize(results) |
| print( |
| f"\n committor-1/2 overlap across all {summ['n_points']} (K,lag): " |
| f"{summ['committor_half_mass_min'] * 100:.0f}-{summ['committor_half_mass_max'] * 100:.0f}%" |
| ) |
| print( |
| f" spurious direct flux across all (K,lag): " |
| f"{summ['direct_min'] * 100:.0f}-{summ['direct_max'] * 100:.0f}%" |
| ) |
| print( |
| f" both channels carry flux at every (K,lag): {summ['both_channels_always']}" |
| ) |
| its_lags = [1, 2, 3, 5, 8, 12, 20, 30] |
| lg, its = rob.its_convergence(feat, hand, K=K, lags=its_lags, n_comp=N_COMP) |
| OUTDIR.mkdir(parents=True, exist_ok=True) |
| rob.make_figure( |
| results, |
| lg, |
| its, |
| K_default=K, |
| lag_default=LAG, |
| dt_ps=1.0, |
| outpath=ROOT / "figures" / "validation_aib9" / "aib9_robustness.png", |
| ) |
| print(f" robustness figure -> {OUTDIR / 'aib9_robustness.png'}") |
| with open(OUTDIR / "robustness.json", "w") as f: |
| json.dump( |
| { |
| "grid": results, |
| "summary": summ, |
| "its_lags": [int(x) for x in lg], |
| "its": np.where(np.isfinite(its), its, None).tolist(), |
| }, |
| f, |
| indent=2, |
| default=float, |
| ) |
| print(f" robustness report -> {OUTDIR / 'robustness.json'}") |
| return summ |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--data", default=DEFAULT_DATA) |
| ap.add_argument( |
| "--robustness", |
| action="store_true", |
| help="also run the K x lag + ITS-convergence stability report", |
| ) |
| ap.add_argument("--colvar_stride", type=int, default=5) |
| args = ap.parse_args() |
|
|
| trajs, temps = aib9.load_openmm_colvar(args.data, stride=args.colvar_stride) |
| print(f"loaded {sum(len(t) for t in trajs)} frames") |
| feat = aib9.featurize_sincos(trajs) |
| hand = aib9.handedness_indicator(trajs) |
|
|
| model, qc = gate_msm(feat, hand, temps) |
| labels, sets = gate_endpoints(model) |
| p23, r23, q, half = gate_committor(model, sets, labels) |
| p24, r24, paths = gate_flux(model, sets, q, half) |
|
|
| rob_summary = run_robustness(feat, hand) if args.robustness else None |
|
|
| print("\n" + "=" * 70) |
| print("SUMMARY") |
| print("=" * 70) |
| for name, ok in [ |
| ("2.3 committor-1/2 <-> known intermediates", p23), |
| ("2.4 reactive flux / rollouts through intermediates", p24), |
| ]: |
| print(f" [{'PASS' if ok else 'FAIL'}] {name}") |
| gate2 = p23 |
| print( |
| f"\n GATE 2 (2.3) -> {'PASS' if gate2 else 'FAIL'} " |
| f"(structural-overlap criterion {'met' if gate2 else 'not met'})" |
| ) |
|
|
| OUTDIR.mkdir(parents=True, exist_ok=True) |
| report = { |
| "config": { |
| "lag": LAG, |
| "K": K, |
| "n_comp": N_COMP, |
| "horizon": HORIZON, |
| "tol_q": TOL_Q, |
| "source": "openmm", |
| }, |
| "qc": qc, |
| "gate_2_3": r23, |
| "gate_2_4": r24, |
| "gate_2_pass": bool(gate2), |
| "robustness": rob_summary, |
| } |
| with open(OUTDIR / "report.json", "w") as f: |
| json.dump(report, f, indent=2, default=float) |
| print(f" report -> {OUTDIR / 'report.json'}") |
| return 0 if gate2 else 1 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
