| |
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from scoobdoob.validation import aib9, msm_build as mb, tpt |
|
|
|
|
| ROOT = Path(__file__).resolve().parent |
| LAG, K, N_COMP, STRIDE, HORIZON = 5, 100, 4, 5, 320 |
| N_PATHS = 600 |
| PATH_SEEDS = (0, 1, 2) |
| KAPPAS = (0.0, 0.10, 0.20, 0.36, 0.50, 0.75, 1.00) |
| GATE_BETA, TIME_SIGMA, W_FLOOR = 0.5, 0.30, 1e-12 |
|
|
|
|
| def conditioned_kernels(P, target, weights): |
| n_states = P.shape[0] |
| terminal = np.zeros(n_states) |
| terminal[target] = 1.0 / len(target) |
| backward = terminal.copy() |
| messages = [backward.copy()] |
| for t in range(HORIZON - 1, -1, -1): |
| raw = P @ (weights[t] * backward) |
| total = raw.sum() |
| backward = raw / total if total > 0 else np.full(n_states, 1.0 / n_states) |
| messages.append(backward.copy()) |
| messages = messages[::-1] |
|
|
| kernels = [] |
| for t in range(HORIZON): |
| numerator = P * (weights[t][None, :] * messages[t + 1][None, :]) |
| denominator = (P @ (weights[t] * messages[t + 1]))[:, None] |
| kernel = np.maximum( |
| np.nan_to_num(np.divide(numerator, np.maximum(denominator, 1e-300))), |
| 0.0, |
| ) |
| row_sum = kernel.sum(axis=1, keepdims=True) |
| bad = (row_sum <= 1e-300).ravel() |
| if bad.any(): |
| kernel[bad] = P[bad] |
| row_sum = kernel.sum(axis=1, keepdims=True) |
| kernels.append(kernel / np.maximum(row_sum, 1e-300)) |
| return kernels |
|
|
|
|
| def sample(kernels, starts, seed): |
| rng = np.random.default_rng(seed) |
| paths = [] |
| for _ in range(N_PATHS): |
| state = int(rng.choice(starts)) |
| path = [state] |
| for kernel in kernels: |
| probabilities = kernel[state] |
| state = int( |
| rng.choice(len(probabilities), p=probabilities / probabilities.sum()) |
| ) |
| path.append(state) |
| paths.append(np.asarray(path, dtype=int)) |
| return paths |
|
|
|
|
| def main(): |
| trajectories, _ = aib9.load_openmm_colvar(stride=STRIDE) |
| features = aib9.featurize_sincos(trajectories) |
| handedness = aib9.handedness_indicator(trajectories) |
| model = mb.build_full_msm( |
| features, |
| handedness, |
| lag=LAG, |
| k=K, |
| n_comp=N_COMP, |
| seed=0, |
| use_tica=True, |
| ) |
|
|
| labels = aib9.classify_microstates(model.active_xy) |
| sets = { |
| name: np.where(labels == name)[0].tolist() |
| for name in ("R", "L", "I_N", "I_C", "other") |
| } |
| P = np.asarray(model.P) |
| pi = np.asarray(model.stationary) |
| projection = model.centers[model.active_symbols][:, :2] |
|
|
| intermediate = sorted(set(sets["I_N"]) | set(sets["I_C"])) |
| q_intermediate_to_left = tpt.committor_agreement(model, intermediate, sets["L"])[1] |
| net_flux_matrix = np.asarray( |
| tpt.reactive_flux(model, sets["R"], sets["L"]).net_flux |
| ) |
| state_flux = (net_flux_matrix.sum(1) + net_flux_matrix.sum(0)) / 2 |
|
|
| n_center = projection[sets["I_N"]].mean(0) |
| c_center = projection[sets["I_C"]].mean(0) |
| gateways = sorted( |
| [ |
| i |
| for i in range(model.n_active) |
| if i not in set(intermediate) |
| and i not in set(sets["L"]) |
| and abs(q_intermediate_to_left[i] - 0.5) < 0.2 |
| ], |
| key=lambda i: -state_flux[i], |
| )[:4] |
| gateway_set = set(gateways) |
|
|
| state_handedness = model.active_xy |
| n_half = state_handedness[:, 1:4].mean(axis=1) |
| c_half = state_handedness[:, 5:8].mean(axis=1) |
| n_channel = set(sets["I_N"]) | {i for i in sets["other"] if n_half[i] > c_half[i]} |
| c_channel = set(sets["I_C"]) | {i for i in sets["other"] if n_half[i] <= c_half[i]} |
|
|
| counts = np.asarray(model.model.count_model.count_matrix, dtype=float) |
| rho = np.clip(counts.sum(1), 1.0, None) / counts.sum(1).mean() |
| z = np.log(state_flux / np.maximum(pi, 1e-15) + 1e-9) |
| z = (z - z.mean()) / z.std() |
| time_window = np.exp( |
| -0.5 * ((np.arange(HORIZON) - (HORIZON - 1) / 2) / (TIME_SIGMA * HORIZON)) ** 2 |
| ) |
|
|
| results = [] |
| for kappa in KAPPAS: |
| density = np.maximum(np.exp(-kappa * rho), W_FLOOR) |
| weights = [ |
| np.maximum( |
| density * np.exp(GATE_BETA * z * time_window[t]), |
| W_FLOOR, |
| ) |
| for t in range(HORIZON) |
| ] |
| kernels = conditioned_kernels(P, sets["L"], weights) |
|
|
| seed_metrics = [] |
| for seed in PATH_SEEDS: |
| paths = sample(kernels, sets["R"], seed) |
| reaches = [] |
| gateway_hits = [] |
| channels = [] |
| visited = set() |
| for path in paths: |
| visited.update(map(int, path)) |
| in_left = np.array([int(x) in set(sets["L"]) for x in path]) |
| reaches.append(bool(in_left.any())) |
| if not in_left.any(): |
| continue |
| first_left = int(np.argmax(in_left)) |
| before = np.array([int(x) in set(sets["R"]) for x in path[:first_left]]) |
| if not before.any(): |
| continue |
| segment = path[int(np.where(before)[0][-1]) : first_left + 1] |
| gateway_hits.append(bool(set(map(int, segment)) & gateway_set)) |
| n_count = sum(int(x) in n_channel for x in segment) |
| c_count = sum(int(x) in c_channel for x in segment) |
| channels.append( |
| "I_N" |
| if n_count > c_count |
| else "I_C" |
| if c_count > n_count |
| else "other" |
| ) |
|
|
| n_paths = sum(channel == "I_N" for channel in channels) |
| c_paths = sum(channel == "I_C" for channel in channels) |
| seed_metrics.append( |
| { |
| "reach_pct": 100 * np.mean(reaches), |
| "gateway_pct": 100 * np.mean(gateway_hits), |
| "n_channel": n_paths, |
| "c_channel": c_paths, |
| "c_share_pct": 100 * c_paths / max(n_paths + c_paths, 1), |
| "distinct_states": len(visited), |
| } |
| ) |
|
|
| record = {"kappa": kappa, "path_seeds": list(PATH_SEEDS)} |
| for key in seed_metrics[0]: |
| values = np.asarray([metric[key] for metric in seed_metrics], dtype=float) |
| record[key] = { |
| "mean": float(values.mean()), |
| "sd": float(values.std(ddof=1)), |
| } |
| results.append(record) |
| print( |
| f"kappa={kappa:4.2f} reach={record['reach_pct']['mean']:5.1f}% " |
| f"gateway={record['gateway_pct']['mean']:5.1f}% " |
| f"C-share={record['c_share_pct']['mean']:5.1f}% " |
| f"states={record['distinct_states']['mean']:4.0f}" |
| ) |
|
|
| output = { |
| "model": { |
| "lag_frames": LAG, |
| "frame_interval_ps": 1.0, |
| "horizon": HORIZON, |
| "n_microstates": K, |
| "clustering_seed": 0, |
| }, |
| "sampler": { |
| "gate_beta": GATE_BETA, |
| "time_sigma": TIME_SIGMA, |
| "n_paths_per_seed": N_PATHS, |
| }, |
| "results": results, |
| } |
| out_path = ROOT / "data" / "aib9_fes" / "aib9_kappa_sweep.json" |
| out_path.write_text(json.dumps(output, indent=2)) |
| print(f"wrote {out_path}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
