| from __future__ import annotations |
|
|
| import numpy as np |
|
|
| from . import msm_build as mb |
| from . import tpt, aib9 |
|
|
|
|
| def analyze_point(feat, hand, K, lag, n_comp=4, seed=0): |
| """One (K, lag) build -> dict of the key stability metrics.""" |
| out = {"K": int(K), "lag": int(lag), "ok": False} |
| try: |
| m = mb.build_full_msm( |
| feat, hand, lag=lag, k=K, n_comp=n_comp, use_tica=True, seed=seed |
| ) |
| except Exception as e: |
| out["error"] = str(e) |
| return out |
| out["n_active"] = int(m.n_active) |
| out["lam1"] = float(np.real(m.model.eigenvalues(k=1)[0])) |
| labels = aib9.classify_microstates(m.active_xy) |
| sets = { |
| k: np.where(labels == k)[0].tolist() for k in ["R", "L", "I_N", "I_C", "other"] |
| } |
| out["n_R"], out["n_L"] = len(sets["R"]), len(sets["L"]) |
| if not sets["R"] or not sets["L"]: |
| return out |
| q = tpt.forward_committor(m, sets["R"], sets["L"]) |
| half = tpt.committor_half_states(q, tol=0.1) |
| inter = set(sets["I_N"]) | set(sets["I_C"]) |
| pi = m.stationary |
| out["committor_half_mass"] = ( |
| float( |
| pi[[h for h in half if h in inter]].sum() / max(pi[list(half)].sum(), 1e-12) |
| ) |
| if len(half) |
| else 0.0 |
| ) |
| rf = tpt.reactive_flux(m, sets["R"], sets["L"]) |
| _, cg = rf.coarse_grain([sets[k] for k in ["R", "I_N", "I_C", "other", "L"]]) |
| F = np.asarray(cg.net_flux) |
| tot = float(cg.total_flux) |
| out["direct"] = float(F[0, 4] / tot) |
| out["flux_IN"] = float(F[0, 1] / tot) |
| out["flux_IC"] = float(F[0, 2] / tot) |
| try: |
| out["ck2"] = mb.ck_test(m.dtrajs, lag=lag, k=K, factors=(2,))[2] |
| except Exception: |
| out["ck2"] = float("nan") |
| out["ok"] = True |
| return out |
|
|
|
|
| def sweep(feat, hand, Ks, lags, n_comp=4, seed=0): |
| return [analyze_point(feat, hand, K, lag, n_comp, seed) for K in Ks for lag in lags] |
|
|
|
|
| def its_convergence(feat, hand, K, lags, n_comp=4, n_its=3, seed=0): |
| """ITS vs lag at a FIXED discretization (vary only the count lag).""" |
| ref_lag = lags[len(lags) // 2] |
| m = mb.build_full_msm( |
| feat, hand, lag=ref_lag, k=K, n_comp=n_comp, use_tica=True, seed=seed |
| ) |
| lg, its = mb.implied_timescales(m.dtrajs, lags, k=K, n_its=n_its) |
| return lg, its |
|
|
|
|
| def summarize(results): |
| """Range of the two decision metrics across all valid sweep points.""" |
| ok = [r for r in results if r.get("ok")] |
| if not ok: |
| return {} |
| ch = [r["committor_half_mass"] for r in ok] |
| di = [r["direct"] for r in ok] |
| return { |
| "n_points": len(ok), |
| "committor_half_mass_min": float(min(ch)), |
| "committor_half_mass_max": float(max(ch)), |
| "direct_min": float(min(di)), |
| "direct_max": float(max(di)), |
| "both_channels_always": bool( |
| all(r["flux_IN"] > 0 and r["flux_IC"] > 0 for r in ok) |
| ), |
| } |
|
|
|
|
| def make_figure(results, its_lags, its, K_default, lag_default, dt_ps, outpath): |
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
| plt.rcParams.update( |
| { |
| "font.family": "DejaVu Sans", |
| "mathtext.fontset": "custom", |
| "mathtext.rm": "DejaVu Sans", |
| "mathtext.it": "DejaVu Sans:italic", |
| "mathtext.bf": "DejaVu Sans:bold", |
| "mathtext.sf": "DejaVu Sans", |
| "mathtext.cal": "DejaVu Sans", |
| "pdf.fonttype": 42, |
| "ps.fonttype": 42, |
| } |
| ) |
|
|
| Ks = sorted({r["K"] for r in results}) |
| lags = sorted({r["lag"] for r in results}) |
| grid = {(r["K"], r["lag"]): r for r in results} |
|
|
| def heat(metric): |
| M = np.full((len(Ks), len(lags)), np.nan) |
| for i, K in enumerate(Ks): |
| for j, lg in enumerate(lags): |
| r = grid.get((K, lg)) |
| if r and r.get("ok"): |
| M[i, j] = r[metric] |
| return M |
|
|
| fig, ax = plt.subplots(1, 3, figsize=(13.2, 4.1), layout="constrained") |
|
|
| for c in range(its.shape[1]): |
| ax[0].plot( |
| np.array(its_lags) * dt_ps, |
| its[:, c] * dt_ps, |
| "-o", |
| label=rf"$t_{c + 1}$", |
| markersize=4, |
| ) |
| ax[0].axvline( |
| lag_default * dt_ps, |
| color="gray", |
| ls="--", |
| label=f"{lag_default * dt_ps:.0f} ps lag", |
| ) |
| ax[0].set_xlabel("Lag time (ps)") |
| ax[0].set_ylabel("Implied timescale (ps)") |
| ax[0].set_yscale("log") |
| ax[0].set_title("Implied timescales") |
| ax[0].legend(fontsize=10, frameon=False) |
|
|
| for a, metric, title, vmax, decimals, ticks in [ |
| ( |
| ax[1], |
| "committor_half_mass", |
| "Intermediate overlap", |
| 100.0, |
| 1, |
| [0, 25, 50, 75, 100], |
| ), |
| ( |
| ax[2], |
| "direct", |
| r"Direct $R$-to-$L$ flux", |
| 1.1, |
| 2, |
| [0, 0.25, 0.5, 0.75, 1.0, 1.1], |
| ), |
| ]: |
| M = 100.0 * heat(metric) |
| finite = M[np.isfinite(M)] |
| if finite.size and finite.max() > vmax: |
| vmax = float(finite.max()) |
| ticks = np.linspace(0, vmax, 5) |
| im = a.imshow( |
| M, |
| aspect="auto", |
| cmap="viridis" if metric != "direct" else "magma_r", |
| vmin=0, |
| vmax=vmax, |
| origin="lower", |
| ) |
| a.set_xticks(range(len(lags))) |
| a.set_xticklabels(lags) |
| a.set_yticks(range(len(Ks))) |
| a.set_yticklabels(Ks) |
| a.set_xlabel("Lag time (ps)") |
| a.set_ylabel("Microstates ($K$)") |
| for i in range(len(Ks)): |
| for j in range(len(lags)): |
| if np.isfinite(M[i, j]): |
| rgb = im.cmap(im.norm(M[i, j]))[:3] |
| lightness = np.dot(rgb, [0.2126, 0.7152, 0.0722]) |
| a.text( |
| j, |
| i, |
| f"{M[i, j]:.{decimals}f}", |
| ha="center", |
| va="center", |
| color="black" if lightness > 0.5 else "white", |
| fontsize=11, |
| ) |
|
|
| a.scatter( |
| [lags.index(lag_default) + 0.30], |
| [Ks.index(K_default) + 0.29], |
| marker="*", |
| c="red", |
| s=105, |
| edgecolors="white", |
| linewidths=0.6, |
| zorder=3, |
| ) |
| cbar = fig.colorbar(im, ax=a, ticks=ticks, fraction=0.048, pad=0.03) |
| cbar.ax.set_title("%", fontsize=11, pad=6) |
| cbar.ax.tick_params(labelsize=10) |
| a.set_title(title) |
|
|
| for letter, a in zip("ABC", ax): |
| a.tick_params(labelsize=10) |
| a.xaxis.label.set_size(11) |
| a.yaxis.label.set_size(11) |
| a.title.set_fontsize(12) |
| a.text( |
| -0.16, |
| 1.07, |
| letter, |
| transform=a.transAxes, |
| fontsize=15, |
| fontweight="bold", |
| va="bottom", |
| ) |
|
|
| from pathlib import Path as _P |
|
|
| _o = _P(outpath) |
| _o.parent.mkdir(parents=True, exist_ok=True) |
| fig.savefig(_o.with_suffix(".png"), dpi=220) |
| fig.savefig(_o.with_suffix(".pdf"), bbox_inches="tight") |
| plt.close(fig) |
|
|
