| from pathlib import Path |
| import argparse |
| import csv |
| import hashlib |
| import importlib.metadata |
| import itertools |
| import json |
| import platform |
| import sys |
| import time |
|
|
| import numpy as np |
|
|
| OUT = Path(__file__).resolve().parent |
| ROOT = OUT.parents[1] |
| sys.path.insert(0, str(ROOT)) |
| from scoobdoob import _conditioning |
|
|
| CONDITIONS = [ |
| ("doob", None), |
| ("none", 0), |
| ("p10", 10), |
| ("p25", 25), |
| ("p50", 50), |
| ("p75", 75), |
| ] |
| SETTINGS = dict( |
| cluster_seeds=list(range(20)), |
| k=100, |
| tica_lag=20, |
| msm_lag=5, |
| horizon=16, |
| alpha_per_transition=7.5, |
| floor_percentiles=[10, 25, 50, 75], |
| no_floor_control=True, |
| start_law="uniform on the original source set; fixed-start mixture", |
| s1_radius=0.25, |
| ) |
|
|
|
|
| def dump(path, obj): |
| path.write_text(json.dumps(obj, indent=2, allow_nan=False) + "\n") |
|
|
|
|
| def digest(path): |
| return hashlib.sha256(Path(path).read_bytes()).hexdigest() |
|
|
|
|
| def exact_metrics(kernels, mu, region, targets, corridor): |
| """Forward full law and killed laws, including the state at step zero.""" |
| m, n = len(mu), len(kernels) |
| mass = mu.copy() |
| survival = mu.copy() |
| survival[region] = 0 |
| avoiding = np.broadcast_to(mu, (m, m)).copy() |
| np.fill_diagonal(avoiding, 0) |
| s1_occupancy = float(mass[region].sum()) |
| self_steps = 0.0 |
| max_outside = float(mass[~corridor].sum()) |
| max_outside_arrivals = 0.0 |
| for K in kernels: |
| self_steps += float(mass @ np.diag(K)) |
| mass = mass @ K |
| np.testing.assert_allclose(mass.sum(), 1, atol=2e-12, rtol=0) |
| max_outside = max(max_outside, float(mass[~corridor].sum())) |
| max_outside_arrivals = max(max_outside_arrivals, float(mass[~corridor].sum())) |
| s1_occupancy += float(mass[region].sum()) |
| survival = survival @ K |
| survival[region] = 0 |
| avoiding = avoiding @ K |
| np.fill_diagonal(avoiding, 0) |
| hit_each = np.clip(1 - avoiding.sum(axis=1), 0, 1) |
| return dict( |
| s1_probability=float(np.clip(1 - survival.sum(), 0, 1)), |
| expected_distinct_states=float(hit_each.sum()), |
| expected_self_transitions=self_steps, |
| self_transition_fraction=self_steps / n, |
| expected_s1_state_observations=s1_occupancy, |
| terminal_probability=float(mass[targets].sum()), |
| max_outside_corridor_probability=max_outside, |
| max_outside_arrival_mask_probability=max_outside_arrivals, |
| ) |
|
|
|
|
| def reference_hit_probability(P, w, mu, region, targets, horizon): |
| """Independent ratio of weighted path sums, separately for each start.""" |
| full = np.zeros(len(P)) |
| full[targets] = 1 |
| avoided = full.copy() |
| avoided[region] = 0 |
| w_avoid = w.copy() |
| w_avoid[region] = 0 |
| for _ in range(horizon): |
| full = P @ (w * full) |
| avoided = P @ (w_avoid * avoided) |
| scale = float(full.max()) |
| if scale <= 0: |
| raise RuntimeError("No feasible weighted continuation") |
| full /= scale |
| avoided /= scale |
| initial = mu > 0 |
| if np.any(full[initial] <= 0): |
| raise RuntimeError("Infeasible prescribed start") |
| avoided[region] = 0 |
| return float(1 - mu[initial] @ (avoided[initial] / full[initial])) |
|
|
|
|
| def validate_toy(): |
| """Brute-force the weighted reference law, not a second kernel simulation.""" |
| P = np.array([[0.5, 0.4, 0.1], [0.1, 0.6, 0.3], [0.2, 0.3, 0.5]]) |
| w = np.array([0.7, 0.3, 1.2]) |
| mu = np.array([0.3, 0.7, 0.0]) |
| region = np.array([False, True, False]) |
| targets = np.array([2]) |
| horizon = 4 |
| kernels, _ = _conditioning.build(P, w, targets, horizon, starts=[0, 1]) |
| got = exact_metrics(kernels, mu, region, targets, np.ones(3, dtype=bool)) |
| expected = np.zeros(4) |
| for i in [0, 1]: |
| norm = 0.0 |
| moments = np.zeros(4) |
| for rest in itertools.product(range(3), repeat=horizon): |
| if rest[-1] != 2: |
| continue |
| path = (i,) + rest |
| prob = float(np.prod([P[a, b] * w[b] for a, b in zip(path[:-1], path[1:])])) |
| norm += prob |
| moments += prob * np.array( |
| [ |
| any(region[s] for s in path), |
| len(set(path)), |
| sum(a == b for a, b in zip(path[:-1], path[1:])), |
| sum(region[s] for s in path), |
| ], |
| dtype=float, |
| ) |
| expected += mu[i] * moments / norm |
| actual = [ |
| got[k] |
| for k in [ |
| "s1_probability", |
| "expected_distinct_states", |
| "expected_self_transitions", |
| "expected_s1_state_observations", |
| ] |
| ] |
| np.testing.assert_allclose(actual, expected, atol=2e-12, rtol=0) |
| other = reference_hit_probability(P, w, mu, region, targets, horizon) |
| np.testing.assert_allclose(other, expected[0], atol=2e-12, rtol=0) |
|
|
| masked_w = np.array([0.0, 0.3, 1.2]) |
| K, _ = _conditioning.build(P, masked_w, targets, 1, starts=[0]) |
| edge = exact_metrics(K, np.array([1.0, 0.0, 0.0]), region, targets, masked_w > 0) |
| np.testing.assert_allclose( |
| [ |
| edge["terminal_probability"], |
| edge["expected_distinct_states"], |
| edge["max_outside_corridor_probability"], |
| edge["max_outside_arrival_mask_probability"], |
| ], |
| [1.0, 2.0, 1.0, 0.0], |
| atol=2e-12, |
| rtol=0, |
| ) |
| return dict( |
| exhaustive_reference_check=True, |
| outside_mask_start_check=True, |
| max_absolute_error=float(np.max(np.abs(np.array(actual) - expected))), |
| ) |
|
|
|
|
| def run(seed): |
| start_time = time.time() |
| toy = validate_toy() |
| from scoobdoob.validation import msm_build as mb, tpt |
| from scoobdoob.validation import mb_ground_truth as gt |
| from scoobdoob.validation import mb_data as V |
|
|
| files = sorted((ROOT / "data/mb_reactive").glob("traj_*.npy")) |
| if len(files) != 12: |
| raise RuntimeError( |
| "Expected twelve trajectory files in data/mb_reactive; run python -m scoobdoob.validation.mb_data first" |
| ) |
| trajs = [np.load(p).astype(np.float64) for p in files] |
| model = mb.build_full_msm( |
| trajs, trajs, lag=5, tica_lag=20, k=100, n_comp=2, seed=seed, use_tica=True |
| ) |
|
|
| np.testing.assert_array_equal(model.active_symbols, np.arange(100)) |
| A, B, _ = mb.q2_sign_endpoints(model) |
| starts, targets = V.orient_endpoints(model, A, B) |
| starts, targets = np.array(starts), np.array(targets) |
| xy = model.active_xy |
| agreement, q, _ = tpt.committor_agreement(model, starts.tolist(), targets.tolist()) |
| flux = tpt.reactive_flux(model, starts.tolist(), targets.tolist()) |
| through = np.asarray(flux.gross_flux).sum(axis=1) |
| corridor = through / through.sum() > 0.5 / model.n_active |
| corridor[starts] = True |
| corridor[targets] = True |
| saddles = gt.saddle_points() |
| S1 = saddles[np.argmin(saddles[:, 0])] |
| landmarks = np.vstack([xy[starts].mean(0), xy[targets].mean(0), saddles]) |
| lo, hi = landmarks.min(0) - 0.35, landmarks.max(0) + 0.35 |
| corridor &= np.all((xy >= lo) & (xy <= hi), axis=1) |
|
|
| counts = model.counts[model.active_symbols].astype(float) |
| rho = np.maximum(counts, 1) / counts.mean() |
| react = (q > 0.05) & (q < 0.95) |
| region = np.linalg.norm(xy - S1, axis=1) < 0.25 |
| if not np.any(region) or not np.any(react): |
| raise RuntimeError("Missing S1 or reference committor region") |
| mu = np.zeros(model.n_active) |
| mu[starts] = 1 / len(starts) |
| saved_match = None |
| if seed == 0: |
| with np.load( |
| ROOT / "data/mb_panels_data/panels_cache.npz", allow_pickle=True |
| ) as z: |
| saved_match = dict( |
| max_state_coordinate_difference=float(np.max(np.abs(xy - z["xy"]))), |
| counts_equal=bool(np.array_equal(counts, z["counts"])), |
| corridor_equal=bool(np.array_equal(corridor, z["on"])), |
| ) |
| np.savez_compressed( |
| OUT / f"model_{seed:02d}.npz", |
| P=model.P, |
| pi=model.stationary, |
| q=q, |
| xy=xy, |
| counts=counts, |
| rho=rho, |
| corridor=corridor, |
| source=starts, |
| target=targets, |
| S1_region=region, |
| mu=mu, |
| ) |
| records = [] |
| for name, pct in CONDITIONS: |
| floor = ( |
| None |
| if pct is None |
| else (0.0 if pct == 0 else float(np.percentile(rho[react], pct))) |
| ) |
| w = ( |
| np.ones_like(rho) |
| if name == "doob" |
| else np.exp(-7.5 * np.maximum(rho, floor)) |
| ) |
| if np.any(w <= 0): |
| raise FloatingPointError("A positive density weight underflowed") |
| w[~corridor] = 0 |
| kernels, h = _conditioning.build(model.P, w, targets, 16, starts=starts) |
| for K in kernels: |
| if np.any((K > 0) & (model.P == 0)): |
| raise RuntimeError("New edge introduced") |
| record = exact_metrics(kernels, mu, region, targets, corridor) |
| independent = reference_hit_probability(model.P, w, mu, region, targets, 16) |
| error = abs(independent - record["s1_probability"]) |
| if error > 2e-10: |
| raise AssertionError(f"Two path-law calculations disagree by {error}") |
| record.update( |
| condition=name, |
| floor_percentile=pct, |
| density_floor=floor, |
| fraction_corridor_states_below_floor=( |
| None if floor is None else float(np.mean(rho[corridor] < floor)) |
| ), |
| independent_recursion_error=error, |
| ) |
| records.append(record) |
| print( |
| f"seed={seed} {name} S1={record['s1_probability']:.8f} " |
| f"distinct/path={record['expected_distinct_states']:.6f}", |
| flush=True, |
| ) |
| source_names = [ |
| "make_mb_panels.py", |
| "mb_stats.py", |
| "scoobdoob/validation/msm_build.py", |
| "scoobdoob/_conditioning.py", |
| "scoobdoob/validation/mb_data.py", |
| ] |
| sources = {str(p.relative_to(ROOT)): digest(p) for p in files} |
| sources.update({p: digest(ROOT / p) for p in source_names}) |
| sources[str(Path(__file__).relative_to(ROOT))] = digest(__file__) |
| versions = { |
| p: importlib.metadata.version(p) |
| for p in ["numpy", "scipy", "scikit-learn", "deeptime"] |
| } |
| dump( |
| OUT / f"result_{seed:02d}.json", |
| dict( |
| seed=seed, |
| settings=SETTINGS, |
| records=records, |
| n_frames=sum(len(t) for t in trajs), |
| n_active=model.n_active, |
| corridor_states=int(corridor.sum()), |
| s1_states=int(region.sum()), |
| source_states_outside_arrival_mask=starts[~corridor[starts]].tolist(), |
| target_states_outside_arrival_mask=targets[~corridor[targets]].tolist(), |
| reference_committor_agreement=float(agreement), |
| seed0_saved_figure_check=saved_match, |
| validation=toy, |
| source_sha256=sources, |
| versions=versions, |
| python=platform.python_version(), |
| elapsed_seconds=time.time() - start_time, |
| ), |
| ) |
|
|
|
|
| def summarize(): |
| results = [ |
| json.loads((OUT / f"result_{s:02d}.json").read_text()) for s in range(20) |
| ] |
| if any(r["settings"] != SETTINGS for r in results): |
| raise RuntimeError("Inconsistent settings") |
| if len({json.dumps(r["source_sha256"], sort_keys=True) for r in results}) != 1: |
| raise RuntimeError("Source/data hashes changed across jobs") |
| records = [ |
| {**row, "cluster_seed": r["seed"]} for r in results for row in r["records"] |
| ] |
| with (OUT / "all_results.csv").open("w", newline="") as f: |
| writer = csv.DictWriter(f, fieldnames=list(records[0])) |
| writer.writeheader() |
| writer.writerows(records) |
| values = { |
| name: np.array( |
| [ |
| next(x for x in r["records"] if x["condition"] == name)[ |
| "s1_probability" |
| ] |
| for r in results |
| ] |
| ) |
| * 100 |
| for name, _ in CONDITIONS |
| } |
| stats = {} |
| for name, _ in CONDITIONS: |
| y = values[name] |
| delta = y - values["doob"] |
| delta25 = y - values["p25"] |
| by = [next(x for x in r["records"] if x["condition"] == name) for r in results] |
| stats[name] = dict( |
| s1_mean_pct=float(y.mean()), |
| s1_sd_pct=float(y.std(ddof=1)), |
| s1_min_pct=float(y.min()), |
| s1_max_pct=float(y.max()), |
| paired_gain_vs_doob_mean_pp=float(delta.mean()), |
| paired_gain_vs_doob_min_pp=float(delta.min()), |
| paired_gain_vs_doob_max_pp=float(delta.max()), |
| paired_vs_p25_mean_pp=float(delta25.mean()), |
| paired_vs_p25_min_pp=float(delta25.min()), |
| paired_vs_p25_max_pp=float(delta25.max()), |
| n_clusterings_better_than_doob=int(np.sum(delta > 1e-10)), |
| distinct_per_path_mean=float( |
| np.mean([x["expected_distinct_states"] for x in by]) |
| ), |
| self_steps_mean=float( |
| np.mean([x["expected_self_transitions"] for x in by]) |
| ), |
| ) |
| archive = json.loads((ROOT / "data/mb_stats.json").read_text()) |
| comparison = { |
| name: dict( |
| archived_mean_pct=float(np.mean(archive[key])), |
| mean_exact_minus_archived_pp=float(np.mean(values[name] - archive[key])), |
| max_absolute_difference_pp=float( |
| np.max(np.abs(values[name] - archive[key])) |
| ), |
| ) |
| for name, key in [("doob", "raw"), ("p25", "fk")] |
| } |
| check = dict( |
| toy_exhaustive_passed=all( |
| r["validation"]["exhaustive_reference_check"] for r in results |
| ), |
| max_independent_recursion_error=max( |
| x["independent_recursion_error"] for x in records |
| ), |
| max_terminal_error=max(abs(x["terminal_probability"] - 1) for x in records), |
| max_outside_corridor_probability=max( |
| x["max_outside_corridor_probability"] for x in records |
| ), |
| max_outside_arrival_mask_probability=max( |
| x["max_outside_arrival_mask_probability"] for x in records |
| ), |
| endpoint_mask_exceptions=[ |
| dict( |
| seed=r["seed"], |
| source=r["source_states_outside_arrival_mask"], |
| target=r["target_states_outside_arrival_mask"], |
| ) |
| for r in results |
| if r["source_states_outside_arrival_mask"] |
| or r["target_states_outside_arrival_mask"] |
| ], |
| seed0_saved_figure_check=results[0]["seed0_saved_figure_check"], |
| ) |
| summary = dict( |
| settings=SETTINGS, |
| conditions=stats, |
| checks=check, |
| archive_comparison=comparison, |
| scope="Exact probabilities on refitted MSMs from fixed MD data. SD is clustering sensitivity, not MD uncertainty.", |
| completed_clusterings=len(results), |
| ) |
| dump(OUT / "summary.json", summary) |
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
|
|
| labels = ["No floor", "10th", "25th (original)", "50th", "75th"] |
| names = ["none", "p10", "p25", "p50", "p75"] |
| fig, axes = plt.subplots(1, 2, figsize=(10, 4.3)) |
| x = np.arange(5) |
| for s in range(20): |
| axes[0].plot( |
| x, [values[n][s] for n in names], color="#93b6cc", alpha=0.42, lw=0.8 |
| ) |
| axes[1].plot( |
| x, |
| [values[n][s] - values["doob"][s] for n in names], |
| color="#93b6cc", |
| alpha=0.42, |
| lw=0.8, |
| ) |
| for ax, metric in zip(axes, ["s1_mean_pct", "paired_gain_vs_doob_mean_pp"]): |
| ax.plot( |
| x, |
| [stats[n][metric] for n in names], |
| "o-", |
| color="#164e70", |
| lw=2, |
| label="Mean of 20 clusterings", |
| ) |
| ax.set_xticks(x, labels, rotation=18) |
| ax.set_xlabel("Density-floor percentile") |
| ax.spines[["top", "right"]].set_visible(False) |
| ax.grid(axis="y", alpha=0.2) |
| axes[0].axhline( |
| stats["doob"]["s1_mean_pct"], ls="--", color="#777777", label="Pure-Doob mean" |
| ) |
| axes[1].axhline(0, ls="--", color="#777777") |
| axes[0].set_ylabel("Exact probability of S1 visitation (%)") |
| axes[1].set_ylabel("Paired gain over pure Doob (percentage points)") |
| axes[0].legend(frameon=False, fontsize=8) |
| fig.suptitle("MB floor sensitivity at fixed density strength and horizon") |
| fig.tight_layout() |
| (ROOT / "figures").mkdir(parents=True, exist_ok=True) |
| for ext in ["pdf", "png", "svg"]: |
| fig.savefig( |
| ROOT / "figures" / f"floor_sensitivity.{ext}", dpi=180, bbox_inches="tight" |
| ) |
| lines = [ |
| "MB FLOOR SENSITIVITY — completed", |
| "", |
| summary["scope"], |
| "", |
| "Condition S1 mean +/- SD (%) gain vs Doob (pp) better /20", |
| ] |
| for name, _ in CONDITIONS: |
| z = stats[name] |
| lines.append( |
| f"{name:12s} {z['s1_mean_pct']:7.3f} +/- {z['s1_sd_pct']:6.3f} " |
| f" {z['paired_gain_vs_doob_mean_pp']:+7.3f} " |
| f"{z['n_clusterings_better_than_doob']}" |
| ) |
| lines += [ |
| "", |
| "Checks:", |
| json.dumps(check, indent=2), |
| "", |
| "The original 25th-percentile setting is retained. No optimum was selected.", |
| "Use per-clustering paired differences in all_results.csv; original results were not overwritten.", |
| "Exact evaluation has no rollout Monte Carlo error. Refitting may differ from archived finite samples.", |
| "Aib9 and the p53 held-out density experiment used no density floor; this test concerns MB.", |
| ] |
| (OUT / "RESULTS.txt").write_text("\n".join(lines) + "\n") |
| print("\n".join(lines), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser() |
| group = p.add_mutually_exclusive_group(required=True) |
| group.add_argument("--seed", type=int, choices=range(20)) |
| group.add_argument("--summarize", action="store_true") |
| group.add_argument("--validate", action="store_true") |
| args = p.parse_args() |
| if args.validate: |
| print(json.dumps(validate_toy(), indent=2)) |
| elif args.summarize: |
| summarize() |
| else: |
| run(args.seed) |
|
|
