1import json
  2import math
  3from pathlib import Path
  4import numpy as np
  5
  6
  7def make_chain(q: float, n_states: int = 10) -> np.ndarray:
  8    P = np.zeros((n_states, n_states), dtype=float)
  9    for z in range(n_states - 1):
 10        P[z, z] = 1.0 - q
 11        P[z, z + 1] = q
 12    P[-1, -1] = 1.0
 13    return P
 14
 15
 16def backward_messages(P: np.ndarray, horizon: int, target: int) -> np.ndarray:
 17    h = np.zeros((horizon + 1, P.shape[0]), dtype=float)
 18    h[horizon, target] = 1.0
 19    for t in range(horizon - 1, -1, -1):
 20        h[t] = P @ h[t + 1]
 21    return h
 22
 23
 24def doob_kernel(P: np.ndarray, h: np.ndarray, t: int, z: int) -> np.ndarray:
 25    d = h[t, z]
 26    if d <= 0:
 27        raise ValueError("zero-feasibility state")
 28    raw = P[z] * h[t + 1] / d
 29    return raw / raw.sum()
 30
 31
 32def sample_doob(P, h, horizon, start, rng):
 33    z = int(start)
 34    path = [z]
 35    for t in range(horizon):
 36        z = int(rng.choice(P.shape[0], p=doob_kernel(P, h, t, z)))
 37        path.append(z)
 38    return np.asarray(path, dtype=int)
 39
 40
 41def conditioned_path_probability(P, h, path):
 42    prob = 1.0
 43    for t, (z, zp) in enumerate(zip(path[:-1], path[1:])):
 44        prob *= doob_kernel(P, h, t, int(z))[int(zp)]
 45    return float(prob)
 46
 47
 48def transition_kl(p, q):
 49    mask = p > 0
 50    return float(np.sum(p[mask] * (np.log(p[mask]) - np.log(np.maximum(q[mask], 1e-300)))))
 51
 52
 53def expected_conditioned_kl(P, h, horizon, start, noise):
 54    """Expected KL of exact conditioned paths vs a log-message-perturbed sampler."""
 55    occ = np.zeros(P.shape[0]); occ[start] = 1.0
 56    total = 0.0
 57    for t in range(horizon):
 58        nxt_occ = np.zeros_like(occ)
 59        for z in range(P.shape[0]):
 60            if occ[z] == 0 or h[t, z] <= 0:
 61                continue
 62            exact = doob_kernel(P, h, t, z)
 63            vals = P[z] * h[t + 1] * np.exp(noise)
 64            approx = vals / vals.sum()
 65            total += occ[z] * transition_kl(exact, approx)
 66            nxt_occ += occ[z] * exact
 67        occ = nxt_occ
 68    return float(total)
 69
 70
 71def rejection_cost(event_probability, n=20000, seed=0):
 72    """Monte Carlo mean rejection cost using geometric draws, not a huge loop."""
 73    rng = np.random.default_rng(seed)
 74    # Number of Bernoulli trials through each success is geometric(p).
 75    return float(rng.geometric(event_probability, size=n).mean())
 76
 77
 78def main():
 79    rng = np.random.default_rng(2460)
 80    n, target, start = 10, 9, 0
 81    # Verification 1: backward recursion gives exact event probability and kernels normalize.
 82    prob_rows = []
 83    for horizon in [9, 12, 20, 30, 50]:
 84        q = 0.5
 85        P = make_chain(q, n)
 86        h = backward_messages(P, horizon, target)
 87        max_resid = max(abs(doob_kernel(P, h, t, z).sum() - 1)
 88                         for t in range(horizon) for z in range(target + 1) if h[t, z] > 0)
 89        prob_rows.append({"T": horizon, "event_probability": float(h[0, start]), "kernel_max_normalization_residual": float(max_resid), "predicted_probability": float(sum(math.comb(horizon, k) * q**k * (1-q)**(horizon-k) for k in range(target, horizon+1)))})
 90
 91    # Verification 2: exact Doob trajectories always hit target at the horizon.
 92    violation_rows = []
 93    for horizon in [9, 12, 20, 30, 50]:
 94        P = make_chain(0.5, n); h = backward_messages(P, horizon, target)
 95        paths = np.array([sample_doob(P, h, horizon, start, rng) for _ in range(3000)])
 96        violation_rows.append({"T": horizon, "violations": int(np.sum(paths[:, -1] != target)), "mean_final_state": float(paths[:, -1].mean())})
 97
 98    # Verification 3: exact path law agrees with the conditional original law.
 99    # Enumerate all binary advance/stay strings for T=12 and compare empirical
100    # Doob frequencies with original path probabilities divided by Pr(C).
101    P = make_chain(0.5, n); h = backward_messages(P, 12, target)
102    Tlaw = 12; samples = [sample_doob(P, h, Tlaw, start, rng) for _ in range(20000)]
103    keys = [tuple(x.tolist()) for x in samples]
104    empirical = {}
105    for k in keys: empirical[k] = empirical.get(k, 0) + 1 / len(keys)
106    exact = {}
107    for bits in np.ndindex(*(2,) * Tlaw):
108        z = 0; path = [z]; prob = 1.0
109        for b in bits:
110            zp = z + int(b) if z < target else target
111            prob *= P[z, zp]; z = zp; path.append(z)
112        if z == target:
113            exact[tuple(path)] = prob / h[0, start]
114    support = set(empirical) | set(exact)
115    tv = 0.5 * sum(abs(empirical.get(k, 0.0) - exact.get(k, 0.0)) for k in support)
116    law_check = {"T": Tlaw, "sample_count": len(samples),
117                 "sampled_terminal_violations": int(sum(x[-1] != target for x in samples)),
118                 "exact_conditioned_paths": len(exact), "empirical_exact_path_TV": float(tv)}
119
120    # Verification 4: error grows quadratically with small log-message error.
121    P = make_chain(0.5, n); h = backward_messages(P, 20, target)
122    kl_rows = []
123    for eps in [0.0, 0.01, 0.02, 0.05, 0.10, 0.20]:
124        # deterministic bounded perturbation, applied only where message is positive
125        noise = np.zeros(n)
126        noise[:-1] = eps * np.linspace(-1, 1, n-1)
127        kl = expected_conditioned_kl(P, h, 20, start, noise)
128        kl_rows.append({"epsilon": eps, "path_KL": kl, "KL_over_epsilon2": (kl/(eps*eps) if eps else 0.0)})
129
130    # Rare-event signature: q^9 probability; Doob costs T transitions, rejection costs 1/p.
131    cost_rows = []
132    for q in [0.5, 0.3, 0.2, 0.1]:
133        T = 9; P = make_chain(q, n); h = backward_messages(P, T, target)
134        p = float(h[0, start])
135        doob_transitions = T
136        rej = rejection_cost(p, n=3000, seed=123)
137        cost_rows.append({"q": q, "event_probability": p, "doob_transitions_per_sample": doob_transitions, "rejection_draws_per_sample": rej, "predicted_rejection_1_over_p": 1/p})
138
139    out = {"probability_and_normalization": prob_rows, "terminal_constraint": violation_rows, "path_law_check": law_check, "message_error_scaling": kl_rows, "rare_event_cost": cost_rows}
140    Path("results.json").write_text(json.dumps(out, indent=2))
141    print(json.dumps(out, indent=2))
142
143
144if __name__ == "__main__":
145    main()