| |
| from __future__ import annotations |
| import json, os |
| import numpy as np, torch |
| import matplotlib.pyplot as plt |
| from matplotlib.gridspec import GridSpec |
|
|
| from scoobdoob.mb_potential import mb_potential_xy |
| from scoobdoob.validation import msm_build as mb, bridge |
| from scoobdoob.validation import mb_ground_truth as gt |
| from scoobdoob.validation import mb_data as V |
| import make_mb_panels as MP |
|
|
| OUT = "figures/mb_unconditional" |
| DATA_DIR = "data/mb_unconditional_data" |
| HORIZONS = [16, 32, 64, 128, 256, 512] |
|
|
|
|
| _CACHE_KEYS = [ |
| "gx", |
| "gy", |
| "gv", |
| "S1", |
| "S2", |
| "Axy", |
| "Bxy", |
| "xlo", |
| "xhi", |
| "ylo", |
| "yhi", |
| "xy", |
| "allxy", |
| "paths0", |
| "reach0", |
| "succ", |
| "N0", |
| ] |
| CACHE_NPZ = f"{DATA_DIR}/uncond_cache.npz" |
|
|
|
|
| def save_cache(D, path=CACHE_NPZ): |
| os.makedirs(os.path.dirname(path), exist_ok=True) |
| out = {} |
| for k, v in D.items(): |
| out[k] = ( |
| np.array([np.asarray(x) for x in v], dtype=object) |
| if k in ("paths0",) |
| else np.asarray(v) |
| ) |
| np.savez_compressed(path, **out) |
| print(f"[cache] -> {path}") |
|
|
|
|
| def load_cache(path=CACHE_NPZ): |
| if not os.path.exists(path): |
| raise SystemExit( |
| f"missing {path}\nrun: python make_mb_unconditional.py --cache-only" |
| ) |
| z = np.load(path, allow_pickle=True) |
| D = {k: z[k] for k in z.files} |
| D["paths0"] = list(D["paths0"]) |
| for k in ("xlo", "xhi", "ylo", "yhi", "N0"): |
| D[k] = float(D[k]) if k != "N0" else int(D[k]) |
| print(f"[cache] loaded {path}") |
| return D |
|
|
|
|
| def compute(): |
| """Everything numeric; no matplotlib.""" |
| trajs = V.load_trajs(str(V.DEFAULT_DATA)) |
| model = mb.build_full_msm( |
| trajs, |
| trajs, |
| lag=MP.LAG_OPT, |
| tica_lag=MP.TICA_LAG, |
| k=V.K, |
| n_comp=2, |
| seed=0, |
| use_tica=True, |
| ) |
| A, B, _ = mb.q2_sign_endpoints(model) |
| start, end = V.orient_endpoints(model, A, B) |
| xy = model.state_xy |
| sad = gt.saddle_points() |
| S1 = sad[np.argmin(sad[:, 0])] |
| S2 = sad[np.argmax(sad[:, 0])] |
| Axy = model.active_to_xy(start).mean(0) |
| Bxy = model.active_to_xy(end).mean(0) |
| endset = set(end) |
|
|
| def reach_frac(paths): |
| return np.mean([any(s in endset for s in p) for p in paths]) * 100 |
|
|
| N0 = MP.N_OPT |
| paths0 = bridge.sample_unconditional_paths(model.P, start, N0, n_paths=3000, seed=1) |
| reach0 = reach_frac(paths0) |
|
|
| succ = [] |
| for N in HORIZONS: |
| ps = bridge.sample_unconditional_paths(model.P, start, N, n_paths=3000, seed=1) |
| succ.append(reach_frac(ps)) |
|
|
| allxy = model.active_xy |
| xlo, xhi = allxy[:, 0].min() - 0.2, allxy[:, 0].max() + 0.2 |
| ylo, yhi = allxy[:, 1].min() - 0.2, allxy[:, 1].max() + 0.2 |
| gx, gy = np.meshgrid(np.linspace(xlo, xhi, 300), np.linspace(ylo, yhi, 300)) |
| with torch.no_grad(): |
| gv = np.clip( |
| mb_potential_xy(torch.tensor(gx), torch.tensor(gy)).numpy(), -150, 100 |
| ) |
|
|
| return {k: v for k, v in locals().items() if k in _CACHE_KEYS} |
|
|
|
|
| def render(D): |
| os.makedirs("figures", exist_ok=True) |
| """Plotting only; recomputes nothing.""" |
| gx = D["gx"] |
| gy = D["gy"] |
| gv = D["gv"] |
| S1 = D["S1"] |
| S2 = D["S2"] |
| Axy = D["Axy"] |
| Bxy = D["Bxy"] |
| xlo = D["xlo"] |
| xhi = D["xhi"] |
| ylo = D["ylo"] |
| yhi = D["yhi"] |
| xy = D["xy"] |
| allxy = D["allxy"] |
| paths0 = D["paths0"] |
| reach0 = D["reach0"] |
| succ = D["succ"] |
| N0 = D["N0"] |
|
|
| fig = plt.figure(figsize=(10.6, 5.0)) |
| gs = GridSpec(1, 2, width_ratios=[1.25, 1], wspace=0.30) |
|
|
| ax = fig.add_subplot(gs[0]) |
| ax.contour(gx, gy, gv, levels=25, colors="0.82", linewidths=0.5, zorder=0) |
| ax.scatter( |
| allxy[:, 0], allxy[:, 1], s=9, color="0.72", zorder=1, label="microstates" |
| ) |
| for p in paths0[:250]: |
| pp = xy[list(p)] |
| ax.plot(pp[:, 0], pp[:, 1], "-", color="#4477aa", lw=0.5, alpha=0.06, zorder=2) |
| ax.scatter(*S1, marker="X", s=150, color="black", zorder=6) |
| ax.scatter(*S2, marker="X", s=150, color="0.5", zorder=6) |
| ax.scatter(*Axy, marker="*", s=260, color="#228833", ec="white", zorder=7) |
| ax.scatter(*Bxy, marker="*", s=260, color="#ee6677", ec="white", zorder=7) |
| ax.annotate( |
| "S1", S1, xytext=(S1[0] - 0.13, S1[1] + 0.10), fontsize=8.5, weight="bold" |
| ) |
| ax.annotate( |
| "S2", |
| S2, |
| xytext=(S2[0] + 0.06, S2[1] + 0.10), |
| fontsize=8.5, |
| color="0.4", |
| weight="bold", |
| ) |
| ax.text( |
| Axy[0], |
| Axy[1] + 0.13, |
| "A", |
| ha="center", |
| fontsize=11, |
| color="#228833", |
| weight="bold", |
| ) |
| ax.text( |
| Bxy[0], |
| Bxy[1] - 0.17, |
| "B", |
| ha="center", |
| fontsize=11, |
| color="#ee6677", |
| weight="bold", |
| ) |
| ax.set_xlim(xlo, xhi) |
| ax.set_ylim(ylo, yhi) |
| ax.set_xlabel("x") |
| ax.set_ylabel("y") |
| ax.text( |
| 0.03, |
| 0.03, |
| f"{reach0:.0f}% reach B", |
| transform=ax.transAxes, |
| fontsize=10, |
| fontweight="bold", |
| bbox=dict(boxstyle="round", fc="white", ec="#4477aa", alpha=0.9), |
| ) |
| ax.legend(loc="upper right", fontsize=8, frameon=False) |
| ax.set_title(f"A Unconditional MSM rollouts (N={N0})", fontsize=10.5, loc="left") |
|
|
| ax = fig.add_subplot(gs[1]) |
| ax.axhline(100, ls=":", color="0.6", lw=1) |
| ax.text( |
| HORIZONS[-1], |
| 100, |
| "conditioned = 100%\nat N=16", |
| va="top", |
| ha="right", |
| fontsize=8, |
| color="0.45", |
| ) |
| ax.plot(HORIZONS, succ, "-o", color="#4477aa", lw=1.8, ms=5) |
| ax.set_xscale("log", base=2) |
| ax.set_xticks(HORIZONS) |
| ax.set_xticklabels(HORIZONS) |
| ax.set_xlabel("horizon N (lag steps)") |
| ax.set_ylabel("% of rollouts that reach B") |
| ax.set_ylim(-4, 108) |
| ax.set_title("B Unconditional reaction vs horizon", fontsize=10.5, loc="left") |
| ax.spines[["top", "right"]].set_visible(False) |
| for xN, yv in zip(HORIZONS, succ): |
| ax.annotate( |
| f"{yv:.0f}%", |
| (xN, yv), |
| xytext=(0, 7), |
| textcoords="offset points", |
| fontsize=7.5, |
| ha="center", |
| color="#22447a", |
| ) |
|
|
| fig.suptitle( |
| "Naive unconditional MSM rollouts almost never react; conditioning is " |
| "required. $\\tau$=5, TICA lag 20", |
| fontsize=11, |
| y=1.0, |
| ) |
| fig.tight_layout(rect=[0, 0, 1, 0.95]) |
| for ext in ("svg", "png", "pdf"): |
| fig.savefig(f"{OUT}.{ext}", dpi=200, bbox_inches="tight") |
|
|
| print( |
| f"[stats] N={N0}: {reach0:.1f}% reach B; " |
| f"horizon {list(zip(HORIZONS, [round(s, 1) for s in succ]))}" |
| ) |
| print(f"[done] -> {OUT}.svg/.png/.pdf") |
|
|
|
|
| def main(): |
| import argparse |
|
|
| ap = argparse.ArgumentParser(description="MB unconditional-rollout figure.") |
| ap.add_argument( |
| "--cache-only", action="store_true", help="compute and cache, do not plot" |
| ) |
| ap.add_argument( |
| "--from-cache", |
| action="store_true", |
| help="plot from cache without recomputing (fast; for restyling)", |
| ) |
| ap.add_argument("--cache", default=CACHE_NPZ) |
| a = ap.parse_args() |
| if a.from_cache: |
| render(load_cache(a.cache)) |
| return |
| D = compute() |
| save_cache(D, a.cache) |
| if a.cache_only: |
| return |
| render(D) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
