| |
| from __future__ import annotations |
| import json, os |
| import numpy as np, torch |
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| 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, tpt, bridge |
| from scoobdoob.validation import mb_ground_truth as gt |
| from scoobdoob.validation import mb_data as V |
|
|
|
|
| plt.rcParams["svg.fonttype"] = "none" |
| plt.rcParams["pdf.fonttype"] = 42 |
|
|
| DATA_DIR = "data/mb_panels_data" |
|
|
| LAG_OPT, N_OPT = 5, 16 |
| TICA_LAG = 20 |
|
|
|
|
| def channel(paths, q, xy, on, thru, start, end, nbins=18): |
| """Dominant route = the maximum-bottleneck (widest) path through the ensemble's OWN empirical transition graph, from the start set to the target set.""" |
| import heapq |
| from collections import defaultdict |
|
|
| W = defaultdict(float) |
| for p in paths: |
| for a, b in zip(p[:-1], p[1:]): |
| if a != b and on[a] and on[b]: |
| W[(int(a), int(b))] += 1.0 |
| if not W: |
| return np.zeros((0, 2)) |
| adj = defaultdict(list) |
| for (a, b), c in W.items(): |
| adj[a].append((b, c)) |
| S, T = set(map(int, start)), set(map(int, end)) |
|
|
| best = defaultdict(float) |
| prev = {} |
| heap = [] |
| for a in S: |
| best[a] = float("inf") |
| heapq.heappush(heap, (-best[a], a)) |
| seen = set() |
| while heap: |
| nb, u = heapq.heappop(heap) |
| nb = -nb |
| if u in seen: |
| continue |
| seen.add(u) |
| if u in T: |
| path = [u] |
| while path[-1] not in S: |
| path.append(prev[path[-1]]) |
| return xy[path[::-1]] |
| for v, c in adj[u]: |
| w = min(nb, c) |
| if w > best[v]: |
| best[v] = w |
| prev[v] = u |
| heapq.heappush(heap, (-w, v)) |
| return np.zeros((0, 2)) |
|
|
|
|
| _CACHE_KEYS = [ |
| "gx", |
| "gy", |
| "gv", |
| "S1", |
| "S2", |
| "Axy", |
| "Bxy", |
| "xlo", |
| "xhi", |
| "ylo", |
| "yhi", |
| "cx", |
| "cy", |
| "q", |
| "keep", |
| "xy", |
| "mep", |
| "arc", |
| "Emep", |
| "ch_raw", |
| "ch_fk", |
| "praw", |
| "pfk", |
| "s1_raw", |
| "s1_fk", |
| "counts", |
| "thru", |
| "on", |
| ] |
| CACHE_NPZ = "data/mb_panels_data/panels_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(): |
| if k in ("praw", "pfk"): |
| out[k] = np.array([np.asarray(p) for p in v], dtype=object) |
| else: |
| out[k] = 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_panels.py --cache-only") |
| z = np.load(path, allow_pickle=True) |
| D = {k: z[k] for k in z.files} |
| for k in ("praw", "pfk"): |
| D[k] = list(D[k]) |
| for k in ("xlo", "xhi", "ylo", "yhi", "s1_raw", "s1_fk"): |
| D[k] = float(D[k]) |
| print(f"[cache] loaded {path}") |
| return D |
|
|
|
|
| def compute(): |
| """Everything numeric. Returns the dict the plot stage needs; touches no matplotlib.""" |
| trajs = V.load_trajs(str(V.DEFAULT_DATA)) |
| model = mb.build_full_msm( |
| trajs, |
| trajs, |
| lag=LAG_OPT, |
| tica_lag=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 |
| _, q, _ = tpt.committor_agreement(model, start, end) |
| rf = tpt.reactive_flux(model, start, end) |
| thru = np.asarray(rf.gross_flux).sum(1) |
| sad = gt.saddle_points() |
| S1 = sad[np.argmin(sad[:, 0])] |
| S2 = sad[np.argmax(sad[:, 0])] |
| q1 = q[np.argmin(np.linalg.norm(xy - S1, axis=1))] |
| q2 = q[np.argmin(np.linalg.norm(xy - S2, axis=1))] |
|
|
| share = thru / thru.sum() |
| on = share > 0.5 / model.n_active |
| on[list(start)] = True |
| on[list(end)] = True |
| Axy = model.active_to_xy(start).mean(0) |
| Bxy = model.active_to_xy(end).mean(0) |
| land = np.vstack([Axy, Bxy, sad]) |
| lo = land.min(0) - 0.35 |
| hi = land.max(0) + 0.35 |
| inbox = ( |
| (xy[:, 0] >= lo[0]) |
| & (xy[:, 0] <= hi[0]) |
| & (xy[:, 1] >= lo[1]) |
| & (xy[:, 1] <= hi[1]) |
| ) |
| on = on & inbox |
| keep = np.where(on)[0] |
| cx, cy = xy[keep, 0], xy[keep, 1] |
|
|
| labels = np.concatenate(model.dtrajs) |
| counts = np.bincount(labels, minlength=model.n_active)[: model.n_active].astype( |
| float |
| ) |
| rho = np.maximum(counts, 1) / counts.mean() |
| react = (q > 0.05) & (q < 0.95) |
| rho_floor = np.percentile(rho[react], 25) |
| w_fk = np.exp(-LAG_OPT * 1.5 * np.maximum(rho, rho_floor)) |
| w_fk[~on] = 0.0 |
| w_raw = np.ones_like(w_fk) |
| w_raw[~on] = 0.0 |
|
|
| def rollout(w): |
| Ph = bridge.doob_kernels(model.P, end, N_OPT, w=w) |
| return bridge.sample_bridge_paths(Ph, start, n_paths=3000, seed=0) |
|
|
| praw = rollout(w_raw) |
| pfk = rollout(w_fk) |
| ch_raw = channel(praw, q, xy, on, thru, start, end) |
| ch_fk = channel(pfk, q, xy, on, thru, start, end) |
| n1 = np.linalg.norm(xy - S1, axis=1) < 0.25 |
| s1_raw = np.mean([any(n1[s] for s in p) for p in praw]) * 100 |
| s1_fk = np.mean([any(n1[s] for s in p) for p in pfk]) * 100 |
|
|
| M = np.load("data/mb_mep.npz") |
| mep = M["mep"] |
| Emep = M["Emep"] |
| arc = M["arc"] |
| _keep_mep = mep[:, 1] >= -0.05 |
| mep = mep[_keep_mep] |
| xlo, xhi = cx.min() - 0.12, cx.max() + 0.12 |
| ylo, yhi = cy.min() - 0.12, cy.max() + 0.12 |
| 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. Reads the cache dict; recomputes nothing. Safe to re-run while restyling.""" |
| 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"] |
| cx = D["cx"] |
| cy = D["cy"] |
| q = D["q"] |
| keep = D["keep"] |
| xy = D["xy"] |
| mep = D["mep"] |
| arc = D["arc"] |
| Emep = D["Emep"] |
| ch_raw = D["ch_raw"] |
| ch_fk = D["ch_fk"] |
| praw = D["praw"] |
| pfk = D["pfk"] |
| s1_raw = D["s1_raw"] |
| s1_fk = D["s1_fk"] |
| counts = D["counts"] |
| thru = D["thru"] |
| on = D["on"] |
|
|
| fig = plt.figure(figsize=(15, 5.4)) |
| gs = GridSpec(1, 4, width_ratios=[1, 1, 1, 0.9], wspace=0.28) |
|
|
| def base(ax): |
| ax.contour(gx, gy, gv, levels=22, colors="0.82", linewidths=0.5, zorder=0) |
| 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.set_xlim(xlo, xhi) |
| ax.set_ylim(ylo, yhi) |
| ax.set_xlabel("x") |
|
|
| ax = fig.add_subplot(gs[0]) |
| base(ax) |
| sc = ax.scatter(cx, cy, c=q[keep], cmap="coolwarm", s=22, zorder=3, vmin=0, vmax=1) |
| ax.plot( |
| mep[:, 0], |
| mep[:, 1], |
| "--", |
| color="0.25", |
| lw=1.6, |
| zorder=4, |
| label="min-energy path", |
| ) |
| ax.set_ylabel("y") |
| ax.set_title("a MSM microstates + committor", fontsize=10.5, loc="left") |
| ax.legend(fontsize=7.5, loc="lower right", frameon=False) |
| 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", |
| ) |
| cax = ax.inset_axes([0.60, 0.90, 0.36, 0.045]) |
| cb = fig.colorbar(sc, cax=cax, orientation="horizontal") |
| cb.set_label("committor $q^+$", fontsize=7.5) |
| cb.ax.tick_params(labelsize=6.5) |
|
|
| ax = fig.add_subplot(gs[1]) |
| base(ax) |
| ax.scatter(cx, cy, s=9, color="0.72", zorder=1) |
| for p in praw[:150]: |
| pp = xy[list(p)] |
| ax.plot(pp[:, 0], pp[:, 1], "-", color="#4477aa", lw=0.5, alpha=0.06, zorder=2) |
| ax.plot(mep[:, 0], mep[:, 1], "--", color="0.35", lw=1.3, zorder=4) |
| ax.plot( |
| ch_raw[:, 0], |
| ch_raw[:, 1], |
| "-o", |
| color="#22447a", |
| lw=2.6, |
| ms=4, |
| zorder=5, |
| label="dominant route (max-flux)", |
| ) |
| ax.set_title(f"b Without FK ({s1_raw:.0f}% via S1)", fontsize=10.5, loc="left") |
| ax.legend(fontsize=7.5, loc="lower right", frameon=False) |
|
|
| ax = fig.add_subplot(gs[2]) |
| base(ax) |
| ax.scatter(cx, cy, s=9, color="0.72", zorder=1) |
| for p in pfk[:150]: |
| pp = xy[list(p)] |
| ax.plot(pp[:, 0], pp[:, 1], "-", color="#cc6677", lw=0.5, alpha=0.06, zorder=2) |
| ax.plot(mep[:, 0], mep[:, 1], "--", color="0.35", lw=1.3, zorder=4) |
| ax.plot( |
| ch_fk[:, 0], |
| ch_fk[:, 1], |
| "-o", |
| color="#a03050", |
| lw=2.6, |
| ms=4, |
| zorder=5, |
| label="dominant route (max-flux)", |
| ) |
| ax.set_title(f"c With floored FK ({s1_fk:.0f}% via S1)", fontsize=10.5, loc="left") |
| ax.legend(fontsize=7.5, loc="lower right", frameon=False) |
|
|
| ax = fig.add_subplot(gs[3]) |
| ax.plot(arc, Emep, "-", color="0.2", lw=2) |
| ipk = np.argmax(Emep) |
| ax.scatter([arc[ipk]], [Emep[ipk]], s=90, color="black", zorder=5) |
| ax.annotate( |
| "S1 (rate-limiting)\nbarrier 106 = 6.4 $k_BT$", |
| (arc[ipk], Emep[ipk]), |
| xytext=(arc[ipk] - 0.05, Emep[ipk] - 30), |
| fontsize=8, |
| ha="center", |
| arrowprops=dict(arrowstyle="->"), |
| ) |
|
|
| iS2 = np.argmin(np.linalg.norm(mep - S2, axis=1)) |
| ax.scatter([arc[iS2]], [Emep[iS2]], s=70, color="0.5", zorder=5) |
| ax.annotate( |
| "S2\nbarrier 9 = 0.5 $k_BT$", |
| (arc[iS2], Emep[iS2]), |
| xytext=(arc[iS2] - 0.05, Emep[iS2] - 34), |
| fontsize=8, |
| color="0.4", |
| ha="center", |
| arrowprops=dict(arrowstyle="->", color="0.4"), |
| ) |
| ax.set_xlabel("path length A$\\to$B") |
| ax.set_ylabel("potential energy (MB units)") |
| ax.set_title("d barrier heights along the MEP", fontsize=10.5, loc="left") |
| ax.spines[["top", "right"]].set_visible(False) |
|
|
| fig.suptitle( |
| "The transition is obligated through S1 (d); the raw MSM under-samples it (b), " |
| "floored FK recovers it (c). $\\tau$=5, N=16, $\\alpha$=1.5, TICA lag 20", |
| fontsize=11.5, |
| y=1.0, |
| ) |
| fig.tight_layout(rect=[0, 0, 1, 0.95]) |
| fig.savefig("figures/mb_panels.svg", bbox_inches="tight") |
| fig.savefig("figures/mb_panels.png", dpi=200, bbox_inches="tight") |
| fig.savefig("figures/mb_panels.pdf", bbox_inches="tight") |
|
|
| print( |
| f"[stats] S1 via raw {s1_raw:.0f}% -> FK {s1_fk:.0f}%; MEP peak at S1 (+106 vs S2 +9)" |
| ) |
| print("[done] -> figures/mb_panels.svg / .png / .pdf") |
|
|
|
|
| def main(): |
| import argparse |
|
|
| ap = argparse.ArgumentParser(description="Muller-Brown panel 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 the 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() |
|
|
