| import numpy as np |
|
|
| from scoobdoob import _conditioning |
| from scoobdoob.validation import tpt |
|
|
| STRIDE, SEED_RELAX = 10, 50 |
| ALPHA_V, TAU = 3.0, 0.12 |
| TIME_SIGMA, COMM_L, COMM_S, W_FLOOR = 0.30, 3.0, 0.35, 1e-12 |
|
|
|
|
| def fk_weights( |
| C, q, horizon, comm_l=COMM_L, comm_s=COMM_S, variant="rho", alpha_rho=None |
| ): |
| """(untilted, committor-tilted) FK weight lists, one entry per horizon step.""" |
| row = np.asarray(C, dtype=np.float64).sum(axis=1) |
| a_eff = (TAU * ALPHA_V) if alpha_rho is None else float(alpha_rho) |
| if variant == "rho": |
| rho = np.clip(row, 1.0, None) / row.mean() |
| w_base = np.exp(-a_eff * rho) |
| else: |
| w_base = np.exp(-a_eff * (1.0 / (row + 1.0))) |
| q_pen = np.exp(comm_l * (-((q - 0.5) ** 2)) / comm_s**2) |
| t = np.arange(horizon) |
| time_win = np.exp(-0.5 * ((t - (horizon - 1) / 2) / (TIME_SIGMA * horizon)) ** 2) |
| wA = [np.maximum(w_base, W_FLOOR) for _ in range(horizon)] |
| wB = [np.maximum(w_base * (q_pen ** time_win[i]), W_FLOOR) for i in range(horizon)] |
| return wA, wB |
|
|
|
|
| def crossing_horizon(m, a0, B, cap=20000): |
| """Return the bounded horizon and reference MFPT in lag steps.""" |
| P = np.asarray(m.P, dtype=np.float64) |
| n = P.shape[0] |
| Bs = set(int(b) for b in B) |
| T = np.array([i for i in range(n) if i not in Bs]) |
| if not len(T): |
| return cap |
| try: |
| mfpt = np.zeros(n) |
| mfpt[T] = np.linalg.solve(np.eye(len(T)) - P[np.ix_(T, T)], np.ones(len(T))) |
| h = int(np.clip(round(mfpt[a0]), 5, cap)) |
| return h, float(mfpt[a0]) |
| except Exception: |
| return cap, float("nan") |
|
|
|
|
| def state_means(st, dg, n_states): |
| """E[drug | state] and per-state counts, on whatever frame subset is passed in.""" |
| mu = np.full(n_states, np.nan) |
| cnt = np.zeros(n_states, dtype=int) |
| for s in range(n_states): |
| m = st == s |
| cnt[s] = m.sum() |
| if cnt[s]: |
| mu[s] = dg[m].mean() |
| return mu, cnt |
|
|
|
|
| def target_set(cvm, dil, a0, n, a, dil_override=None): |
| """The B set the tilt aims at.""" |
| if a.target == "cleft": |
| B = np.array([b for b in np.where(cvm[:, 1] > a.cleft_thr)[0] if b != a0]) |
| if not len(B): |
| raise SystemExit( |
| f"empty target set B at cleft_thr={a.cleft_thr}. Under a pocket-feature MSM the " |
| f"per-state mean cleft CV never reaches this threshold because pocket states mix " |
| f"open and closed frames of the GLOBAL coordinate. Use --target dilation." |
| ) |
| return B |
| d = dil if dil_override is None else dil_override |
| ok = np.isfinite(d) |
| thr = np.quantile(d[ok], 1.0 - a.target_q) |
| B = np.array([b for b in np.where(ok & (d >= thr))[0] if b != a0]) |
| if not len(B): |
| raise SystemExit("empty dilation target set") |
| return B |
|
|
|
|
| def build_selectors(m, cvm, C, a0, a, dil=None, B=None): |
| """Weight vectors over ACTIVE states for every selector, each summing to 1.""" |
| P = np.asarray(m.P, dtype=np.float64) |
| pi = np.asarray(m.stationary, dtype=np.float64) |
| n = len(pi) |
|
|
| if B is None: |
| B = target_set(cvm, dil, a0, n, a) |
| q = np.asarray(tpt.forward_committor(m, [a0], list(B))) |
|
|
| sel = {} |
| sel["equilibrium"] = pi / pi.sum() |
| sel["uniform_states"] = np.full(n, 1.0 / n) |
|
|
| half = np.abs(q - 0.5) <= 0.1 |
| if half.sum(): |
| w = pi * half |
| sel["committor_half"] = w / w.sum() |
|
|
| try: |
| rf = tpt.reactive_flux(m, [a0], list(B)) |
| gross = np.asarray(rf.gross_flux) |
| carried = gross.sum(axis=1) + gross.sum(axis=0) |
| if carried.sum() > 0: |
| sel["tpt_flux"] = carried / carried.sum() |
| except Exception as e: |
| print(f" [warn] TPT flux unavailable: {e}") |
|
|
| horizon = a.horizon |
| mfpt = float("nan") |
| if horizon is None: |
| horizon, mfpt = crossing_horizon(m, a0, B) |
| print(f" [fk] horizon = {horizon} lag steps (MFPT a0->B = {mfpt:.1f})") |
|
|
| def occupancy(wl, H): |
| Ph, _ = _conditioning.build(P, wl, list(B), H, starts=[a0]) |
| mu = np.zeros(n) |
| mu[a0] = 1.0 |
| acc = np.zeros(n) |
| for Pn in Ph: |
| mu = mu @ np.asarray(Pn) |
| s = mu.sum() |
| if not np.isfinite(s) or s <= 0: |
| raise FloatingPointError( |
| "Conditioned occupancy lost all probability mass" |
| ) |
| mu = mu / s |
| acc += mu |
| return acc / acc.sum() |
|
|
| wA, wB = fk_weights(C, q, horizon, variant=a.fk_variant) |
| for name, wl in (("doob_untilted", wA), ("fk_doob", wB)): |
| try: |
| sel[name] = occupancy(wl, horizon) |
| except Exception as e: |
| print(f" [warn] {name} failed: {e}") |
|
|
| for a_rho in getattr(a, "alpha_rho_sweep", []) or []: |
| try: |
| wa, _ = fk_weights(C, q, horizon, variant=a.fk_variant, alpha_rho=a_rho) |
| sel[f"alg1[a={a_rho:g}]"] = occupancy(wa, horizon) |
| except Exception as e: |
| print(f" [warn] alg1[a={a_rho:g}] failed: {e}") |
|
|
| for cl in a.comm_l_sweep: |
| for hf in a.horizon_sweep: |
| H = max(5, int(round(horizon * hf))) |
| try: |
| _, wb = fk_weights(C, q, H, comm_l=cl, variant=a.fk_variant) |
| sel[f"fk_doob[L={cl},H={H}]"] = occupancy(wb, H) |
| except Exception: |
| pass |
|
|
| return sel, q, B, a0, horizon |
|
|
|
|
| def enrich(weights, mu_state, cnt, p_open_state, base_mean, base_open): |
| """Expected druggability and open-probability under a state-weight vector.""" |
| ok = np.isfinite(mu_state) & (cnt > 0) |
| w = np.where(ok, weights, 0.0) |
| if w.sum() <= 0: |
| return None |
| w = w / w.sum() |
| em = float((w * np.nan_to_num(mu_state)).sum()) |
| eo = float((w * np.nan_to_num(p_open_state)).sum()) |
| return { |
| "E_drug": em, |
| "P_open": eo, |
| "enrich_mean": em / base_mean if base_mean > 0 else float("nan"), |
| "enrich_open": eo / base_open if base_open > 0 else float("nan"), |
| "coverage": float(weights[ok].sum() / max(weights.sum(), 1e-300)), |
| } |
|
|
