| |
| from __future__ import annotations |
| import sys |
| from pathlib import Path |
| import numpy as np |
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import matplotlib.patheffects as pe |
| from matplotlib.collections import LineCollection |
| from matplotlib.lines import Line2D |
| from matplotlib.colors import LinearSegmentedColormap |
| from scipy.ndimage import gaussian_filter, gaussian_filter1d |
|
|
| ROOT = Path(__file__).resolve().parent |
| sys.path.insert(0, str(ROOT)) |
| from scoobdoob.validation import aib9, bridge, msm_build as mb, tpt |
|
|
| LAG, K, NC, ST, H, NP = 5, 100, 4, 5, 320, 600 |
| KAPPA_RHO, TSg, W = 0.20, 0.30, 1e-12 |
| MACRO_C = {"R": "#1F3A93", "L": "#B22222", "I_N": "#1E7A46", "I_C": "#C4720A"} |
| MACRO_LBL = {"R": "R", "L": "L", "I_N": "I$_N$", "I_C": "I$_C$"} |
|
|
| ORCHID_P = LinearSegmentedColormap.from_list( |
| "orchid_blue", |
| ["#12225E", "#28409B", "#4470CB", "#7CA9DF", "#AED2EE", "#D8EAF8", "#F2F8FD"], |
| N=256, |
| ) |
|
|
| COMMIT = LinearSegmentedColormap.from_list( |
| "purple_red", |
| ["#3B0F70", "#7B2D8E", "#B47ACB", "#EBDDF0", "#F6C3AE", "#DA5432", "#7E1008"], |
| N=256, |
| ) |
| GATE_C = "#FF2D95" |
| PAGE = "#FFFFFF" |
|
|
|
|
| def main(): |
| trajs, _ = aib9.load_openmm_colvar(stride=ST) |
| feat, hand = aib9.featurize_sincos(trajs), aib9.handedness_indicator(trajs) |
| model = mb.build_full_msm( |
| feat, hand, lag=LAG, k=K, n_comp=NC, seed=0, use_tica=True |
| ) |
| labels = aib9.classify_microstates(model.active_xy) |
| sets = { |
| s: np.where(labels == s)[0].tolist() for s in ["R", "L", "I_N", "I_C", "other"] |
| } |
| q = tpt.committor_agreement(model, sets["R"], sets["L"])[1] |
| pi = model.stationary |
| P = model.P |
| NA = model.n_active |
| proj = model.centers[model.active_symbols][:, :2].copy() |
| INT = sorted(set(sets["I_N"]) | set(sets["I_C"])) |
| qIL = tpt.committor_agreement(model, INT, sets["L"])[1] |
| Nn = np.asarray(tpt.reactive_flux(model, sets["R"], sets["L"]).net_flux) |
| netf = (Nn.sum(1) + Nn.sum(0)) / 2 |
| INcen, ICcen = proj[sets["I_N"]].mean(0), proj[sets["I_C"]].mean(0) |
| exitg = sorted( |
| [ |
| i |
| for i in range(NA) |
| if i not in set(INT) and i not in set(sets["L"]) and abs(qIL[i] - 0.5) < 0.2 |
| ], |
| key=lambda i: -netf[i], |
| )[:4] |
| egch = [ |
| "I_N" if np.hypot(*(proj[i] - INcen)) < np.hypot(*(proj[i] - ICcen)) else "I_C" |
| for i in exitg |
| ] |
| egset = set(exitg) |
| half = np.where(np.abs(q - 0.5) < 0.08)[0] |
|
|
| SL = model.active_xy |
| n_half, c_half = SL[:, 1:4].mean(axis=1), SL[:, 5:8].mean(axis=1) |
| Nchan = set(sets["I_N"]) | {i for i in sets["other"] if n_half[i] > c_half[i]} |
| Cchan = set(sets["I_C"]) | {i for i in sets["other"] if n_half[i] <= c_half[i]} |
|
|
| chan_state = np.array( |
| [ |
| "I_N" |
| if np.hypot(*(proj[i] - INcen)) < np.hypot(*(proj[i] - ICcen)) |
| else "I_C" |
| for i in range(NA) |
| ], |
| dtype=object, |
| ) |
| rf = tpt.reactive_flux(model, sets["R"], sets["L"]) |
| pf = tpt.pathway_fractions(rf, fraction=0.95) |
|
|
| part4 = [sets["R"], sorted(Nchan), sorted(Cchan), sets["L"]] |
| _, cg4 = rf.coarse_grain(part4) |
| F4 = np.asarray(cg4.net_flux) |
| tot4 = float(cg4.total_flux) |
| fRN4, fRC4 = F4[0, 1] / tot4, F4[0, 2] / tot4 |
| print( |
| f"[reference MSM] merged-partition flux out of R: N-channel {fRN4 * 100:.1f}% " |
| f"C-channel {fRC4 * 100:.1f}%" |
| ) |
|
|
| def chroute(tgt): |
| idx = [ |
| i |
| for i, p_ in enumerate(pf["paths"]) |
| if tpt.path_passes_through(p_, set(tgt)) |
| ] |
| return ( |
| ( |
| np.array(pf["paths"][max(idx, key=lambda i: pf["frac_of_total"][i])]), |
| float(sum(pf["frac_of_total"][i] for i in idx)), |
| ) |
| if idx |
| else (None, 0.0) |
| ) |
|
|
| rIN, fIN = chroute(sets["I_N"]) |
| rIC, fIC = chroute(sets["I_C"]) |
|
|
| Cc = np.asarray(model.model.count_model.count_matrix, float) |
| rho = np.clip(Cc.sum(1), 1.0, None) / Cc.sum(1).mean() |
| wb = np.maximum(np.exp(-KAPPA_RHO * rho), W) |
| tw = np.exp(-0.5 * ((np.arange(H) - (H - 1) / 2) / (TSg * H)) ** 2) |
| gz = np.log(netf / np.maximum(pi, 1e-15) + 1e-9) |
| gz = (gz - gz.mean()) / gz.std() |
|
|
| def kern(wl, Pm=None): |
| Pm = P if Pm is None else Pm |
| return bridge.doob_kernels(Pm, sets["L"], H, w=wl) |
|
|
| def samp(Ph, seed, n_=NP): |
| return bridge.sample_bridge_paths(Ph, sets["R"], n_paths=n_, seed=seed) |
|
|
| naive = [ |
| np.asarray(p, int) |
| for p in bridge.sample_unconditional_paths(P, sets["R"], H, n_paths=NP, seed=1) |
| ] |
| doob = samp(bridge.doob_kernels(P, sets["L"], H), 1) |
| wfk = [np.maximum(wb * (np.exp(0.5 * gz) ** tw[t]), W) for t in range(H)] |
| dfk = samp(kern(wfk), 1) |
|
|
| def confined(forbid): |
| Pm = P.copy() |
| Pm[:, list(forbid)] = 0.0 |
| sums = Pm.sum(1, keepdims=True) |
| return np.divide(Pm, sums, out=np.zeros_like(Pm), where=sums > 0) |
|
|
| cfk = samp(kern(wfk, confined(Cchan)), 11, NP // 2) + samp( |
| kern(wfk, confined(Nchan)), 12, NP // 2 |
| ) |
| Rs, Ls = set(sets["R"]), set(sets["L"]) |
|
|
| def reactive(paths): |
| segs = [] |
| for p in paths: |
| inL = np.array([int(x) in Ls for x in p]) |
| if not inL.any(): |
| continue |
| fL = int(np.argmax(inL)) |
| pre = np.array([int(x) in Rs for x in p[:fL]]) |
| if not pre.any(): |
| continue |
| segs.append(p[int(np.where(pre)[0][-1]) : fL + 1]) |
| return segs |
|
|
| S = { |
| "naive": reactive(naive), |
| "doob": reactive(doob), |
| "dfk": reactive(dfk), |
| "cfk": reactive(cfk), |
| } |
| M = {} |
| for k_, segs in S.items(): |
| reach = len(segs) / NP |
| via = ( |
| float(np.mean([bool(set(int(x) for x in s_) & egset) for s_ in segs])) |
| if segs |
| else 0.0 |
| ) |
|
|
| def seg_chan(s_): |
| fN = sum(int(x) in Nchan for x in s_) |
| fC = sum(int(x) in Cchan for x in s_) |
| return "I_N" if fN > fC else ("I_C" if fC > fN else None) |
|
|
| cs = [seg_chan(s_) for s_ in segs] |
| nN = sum(1 for c in cs if c == "I_N") |
| nC = sum(1 for c in cs if c == "I_C") |
| M[k_] = (reach, via, len(segs), nN, nC) |
| print( |
| f"{k_:>6}: reach-L {reach * 100:3.0f}% n_react {len(segs):3d} " |
| f"via gateway {via * 100:3.0f}% channel I_N {nN:3d} / I_C {nC:3d} " |
| f"(I_C share {100 * nC / max(nN + nC, 1):.0f}%)" |
| ) |
|
|
| if proj[sets["R"], 0].mean() > proj[sets["L"], 0].mean(): |
| proj[:, 0] *= -1 |
| flip = True |
| else: |
| flip = False |
| Y = np.vstack( |
| [model.tica.transform(np.asarray(f, dtype=np.float64)) for f in feat] |
| )[:, :2] |
| if flip: |
| Y = Y.copy() |
| Y[:, 0] *= -1 |
| x, y = Y[:, 0], Y[:, 1] |
| rr = [ |
| [x.min() - 0.04 * np.ptp(x), x.max() + 0.04 * np.ptp(x)], |
| [y.min() - 0.04 * np.ptp(y), y.max() + 0.04 * np.ptp(y)], |
| ] |
| Hc, xe, ye = np.histogram2d(x, y, bins=140, range=rr) |
| dd = gaussian_filter(Hc, 2.0) |
| with np.errstate(divide="ignore"): |
| F = -np.log(np.where(dd > 0, dd / dd.sum(), np.nan)) |
| F -= np.nanmin(F[np.isfinite(F)]) |
| F = np.ma.masked_where((dd <= 0) | ~np.isfinite(F) | (F > 6.5), F) |
| xc, yc = 0.5 * (xe[:-1] + xe[1:]), 0.5 * (ye[:-1] + ye[1:]) |
| setsA = {s: np.array(sets[s]) for s in sets} |
|
|
| def fes(ax, a=1.0): |
| cf = ax.contourf( |
| xc, |
| yc, |
| F.T, |
| levels=np.arange(0, 6.75, 0.25), |
| cmap=ORCHID_P, |
| zorder=0, |
| alpha=a, |
| ) |
| ax.set_xlim(xc[0], xc[-1]) |
| ax.set_ylim(yc[0], yc[-1]) |
| ax.set_xlabel("TICA 1", fontsize=10.5) |
| for s in ("top", "right"): |
| ax.spines[s].set_visible(False) |
| return cf |
|
|
| def rings(ax, lw=1.6, s=140): |
| for lab in ["R", "L", "I_N", "I_C"]: |
| idx = setsA[lab] |
| if len(idx): |
| ax.scatter( |
| proj[idx, 0], |
| proj[idx, 1], |
| marker="o", |
| s=s, |
| facecolors="none", |
| edgecolors=MACRO_C[lab], |
| linewidths=lw, |
| zorder=5, |
| ) |
|
|
| def gates(ax, s=190): |
| for i, ch in zip(exitg, egch): |
| ax.scatter( |
| proj[i, 0], |
| proj[i, 1], |
| marker="X", |
| s=s, |
| c=GATE_C, |
| edgecolors=MACRO_C[ch], |
| linewidths=2.4, |
| zorder=9, |
| ) |
|
|
| def sm(seq): |
| d = seq[np.concatenate([[True], seq[1:] != seq[:-1]])] |
| xy = proj[d] |
| return ( |
| xy |
| if len(xy) < 3 |
| else np.column_stack( |
| [ |
| gaussian_filter1d(xy[:, 0], 1.0, mode="nearest"), |
| gaussian_filter1d(xy[:, 1], 1.0, mode="nearest"), |
| ] |
| ) |
| ) |
|
|
| def crossings(ax, segs, color): |
| """Color crossings by channel and emphasize gateway visitation.""" |
| buckets = { |
| ("I_N", True): [], |
| ("I_N", False): [], |
| ("I_C", True): [], |
| ("I_C", False): [], |
| } |
| for s_ in segs: |
| xy = sm(s_) |
| if len(xy) < 2: |
| continue |
| fN = sum(int(x) in Nchan for x in s_) |
| fC = sum(int(x) in Cchan for x in s_) |
| ch = "I_N" if fN >= fC else "I_C" |
| buckets[(ch, bool(set(int(t) for t in s_) & egset))].append(xy) |
| for (ch, viaG), L_ in buckets.items(): |
| if not L_: |
| continue |
| ax.add_collection( |
| LineCollection( |
| [[xy[i], xy[i + 1]] for xy in L_ for i in range(len(xy) - 1)], |
| colors=MACRO_C[ch], |
| linewidths=0.85 if viaG else 0.5, |
| alpha=0.20 if viaG else 0.11, |
| zorder=4 if viaG else 3, |
| capstyle="round", |
| ) |
| ) |
|
|
| fig = plt.figure(figsize=(23.5, 11.4)) |
| fig.patch.set_facecolor(PAGE) |
| gs = fig.add_gridspec(2, 4, height_ratios=[1.12, 1.0], hspace=0.24, wspace=0.15) |
|
|
| axa = fig.add_subplot(gs[0, 0:2]) |
| axa.set_facecolor(PAGE) |
| fes(axa) |
| sc = axa.scatter( |
| proj[:, 0], |
| proj[:, 1], |
| c=q, |
| cmap=COMMIT, |
| s=54, |
| vmin=0, |
| vmax=1, |
| edgecolors=PAGE, |
| linewidths=0.9, |
| zorder=6, |
| ) |
| rings(axa) |
| xs = axa.scatter( |
| proj[half, 0], |
| proj[half, 1], |
| marker="x", |
| c="#101010", |
| s=70, |
| linewidths=2, |
| zorder=7, |
| ) |
| xs.set_path_effects([pe.withStroke(linewidth=3.4, foreground=PAGE)]) |
| gates(axa) |
| axa.set_ylabel("TICA 2", fontsize=10.5) |
| axa.set_title("a Committor + macrostates", fontsize=12, fontweight="bold", pad=8) |
| cb = fig.colorbar(sc, ax=axa, pad=0.012, fraction=0.045) |
| cb.set_label("committor $q^+$ (R$\\rightarrow$L)", fontsize=9) |
| cb.outline.set_visible(False) |
| cb.ax.tick_params(length=0, labelsize=8) |
| lh = [ |
| Line2D( |
| [0], [0], marker="o", color="w", mfc="none", mec=MACRO_C[m], mew=1.8, ms=10 |
| ) |
| for m in ["R", "I_N", "I_C", "L"] |
| ] |
| lh += [ |
| Line2D([0], [0], marker="x", color="#101010", lw=0, mew=2, ms=9), |
| Line2D( |
| [0], |
| [0], |
| marker="X", |
| color="w", |
| mfc=GATE_C, |
| mec=MACRO_C["I_N"], |
| mew=1.6, |
| ms=11, |
| ), |
| ] |
| axa.legend( |
| lh, |
| [ |
| "R state", |
| "I$_N$ state", |
| "I$_C$ state", |
| "L state", |
| "committor-½", |
| "rate-limiting exit gate", |
| ], |
| loc="lower left", |
| fontsize=7.3, |
| framealpha=0.95, |
| ncol=2, |
| ).get_frame().set_edgecolor("#C9D3DE") |
|
|
| axb = fig.add_subplot(gs[0, 2:4]) |
| axb.set_facecolor(PAGE) |
| cfb = fes(axb) |
| rings(axb) |
| for r, cc, lab, fr in ( |
| (rIN, MACRO_C["I_N"], "I$_N$", fIN), |
| (rIC, MACRO_C["I_C"], "I$_C$", fIC), |
| ): |
| if r is None: |
| continue |
| xy = proj[r] |
| axb.plot( |
| xy[:, 0], |
| xy[:, 1], |
| "-o", |
| lw=3.0, |
| ms=6, |
| color=cc, |
| mec=PAGE, |
| mew=1.0, |
| zorder=8, |
| path_effects=[pe.withStroke(linewidth=5.4, foreground=PAGE)], |
| label=f"channel via {lab}: {fr * 100:.0f}% of flux", |
| ) |
| gates(axb) |
| axb.set_title( |
| "b The two reactive channels (TPT)", fontsize=12, fontweight="bold", pad=8 |
| ) |
| axb.legend( |
| loc="upper left", fontsize=8.2, framealpha=0.95 |
| ).get_frame().set_edgecolor("#C9D3DE") |
| cbF = fig.colorbar(cfb, ax=axb, pad=0.012, fraction=0.045, ticks=range(0, 7)) |
| cbF.set_label("free energy $F/k_\\mathrm{B}T$ (background)", fontsize=9) |
| cbF.outline.set_visible(False) |
| cbF.ax.tick_params(length=0, labelsize=8) |
|
|
| panels = [ |
| ("c naive MSM (unconditioned)", "naive", "#5C6570"), |
| ("d Doob bridge (conditioned)", "doob", "#1F4EA1"), |
| ("e Doob + FK (global tilt)", "dfk", "#A3161C"), |
| ("f Doob + FK (channel-dependent)", "cfk", "#1E7A46"), |
| ] |
| samp_axes = [] |
| for col, (title, key, color) in enumerate(panels): |
| ax = fig.add_subplot(gs[1, col]) |
| samp_axes.append(ax) |
| ax.set_facecolor(PAGE) |
| fes(ax, 0.9) |
| crossings(ax, S[key], color) |
| rings(ax, lw=1.2, s=90) |
| gates(ax, s=150) |
| reach, via, ns, nN, nC = M[key] |
| ax.set_title(title, fontsize=12, fontweight="bold", pad=8) |
| if col == 0: |
| ax.set_ylabel("TICA 2", fontsize=10.5) |
| ax.text( |
| 0.5, |
| 0.03, |
| f"reach L: {reach * 100:.0f}% ({ns} of {NP} crossings)\n" |
| f"via rate-limiting gateway: {via * 100:.0f}%\n" |
| f"N-channel {nN} / C-channel {nC}", |
| transform=ax.transAxes, |
| ha="center", |
| va="bottom", |
| fontsize=11, |
| fontweight="bold", |
| color="#111", |
| linespacing=1.5, |
| bbox=dict(boxstyle="round,pad=0.5", fc=PAGE, ec=color, lw=1.8, alpha=0.96), |
| zorder=13, |
| ) |
| axc0 = samp_axes[0] |
| axc0.legend( |
| [ |
| Line2D([0], [0], color=MACRO_C["I_N"], lw=2.2), |
| Line2D([0], [0], color=MACRO_C["I_C"], lw=2.2), |
| Line2D([0], [0], color="#555", lw=2.2), |
| Line2D([0], [0], color="#555", lw=0.9), |
| ], |
| [ |
| "crossing via N-channel", |
| "crossing via C-channel", |
| "uses rate-limiting gateway", |
| "does not", |
| ], |
| loc="upper left", |
| fontsize=7.6, |
| framealpha=0.95, |
| ).get_frame().set_edgecolor("#C9D3DE") |
|
|
| fig.suptitle( |
| "Aib9 R→L: Doob fixes ARRIVAL (40%→100%); global FK routes crossings through the " |
| "rate-limiting exit gateway (40%→82%) but COLLAPSES the C-channel (47%→2%);\n" |
| "channel-dependent FK keeps the gateway emphasis (68%) AND both channels", |
| fontsize=13, |
| fontweight="bold", |
| y=1.02, |
| ) |
| fig.text( |
| 0.006, |
| 0.006, |
| f"Aib9 OpenMM MSM (lag {LAG}, K={K}); {NP} rollouts/sampler, horizon {H}. Panels c–f draw the " |
| f"REACTIVE SEGMENT of each rollout (last frame in R → first in L), so only genuine crossings are\n" |
| f"shown — naive produces {M['naive'][2]}, the conditioned samplers {NP}. Rate-limiting exit gateway = " |
| f"the 4 highest net-flux committor-½ states of the trap→L barrier (magenta X, edge = channel).\n" |
| f"Channel = MERGED handedness partition (each 'other' microstate joins the channel whose half leads " |
| f"the flip): counting only the labelled I$_N$/I$_C$ helices under-counts a channel, because most\n" |
| f"reactive flux runs R→other→L. The two channels are intrinsically BALANCED — merged-partition flux " |
| f"out of R is {fRN4 * 100:.0f}% N / {fRC4 * 100:.0f}% C (Aib9 is a homopolymer, so N↔C is a symmetry) —\n" |
| f"so f's equal 300/300 allocation matches the true split; the measured results there are that both " |
| f"confined bridges reach L 100% independently and keep 68% gateway use.", |
| fontsize=8.2, |
| color="#6B7785", |
| ha="left", |
| va="bottom", |
| linespacing=1.5, |
| ) |
| fig.subplots_adjust(left=0.045, right=0.988, top=0.90, bottom=0.115) |
| out = ROOT / "figures" / "aib9_method" |
| out.parent.mkdir(parents=True, exist_ok=True) |
| for ext, kw in (("png", dict(dpi=180)), ("pdf", {})): |
| fig.savefig( |
| out.with_suffix("." + ext), facecolor=PAGE, bbox_inches="tight", **kw |
| ) |
| print(f"[fig] {out.with_suffix('.' + ext)}") |
| plt.close(fig) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
