#!/usr/bin/env python 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()