1import json, math
  2from pathlib import Path
  3import numpy as np
  4
  5SEED = 1729
  6EPS = 0.10
  7PERSISTENCE = 3
  8RIDGE = 1e-4
  9CANDIDATES = [1, 2, 4, 8, 16, 32]
 10
 11
 12def closure_rho(z, max_lag=100):
 13    """Ensemble/time normalized autocorrelation, avoiding sequence concatenation."""
 14    z = np.asarray(z, float)
 15    z = z - np.mean(z)
 16    den = np.mean(z * z)
 17    return np.array([1.0 if k == 0 else np.mean(z[:, :-k] * z[:, k:]) / den
 18                     for k in range(max_lag)])
 19
 20
 21def first_persistent_crossing(rho, eps=EPS, persistence=PERSISTENCE):
 22    for k in range(1, len(rho) - persistence + 1):
 23        if np.all(np.abs(rho[k:k + persistence]) <= eps):
 24            return k
 25    return len(rho) - persistence
 26
 27
 28def make_data(a, nseq=240, length=240, noise=0.65, seed=SEED):
 29    rng = np.random.default_rng(seed)
 30    z = rng.normal(size=(nseq, length))
 31    for t in range(1, length):
 32        z[:, t] = a * z[:, t - 1] + math.sqrt(1 - a * a) * z[:, t]
 33    # x is a resolved signal with unresolved closure injection z observed noisily.
 34    x = z + noise * rng.normal(size=z.shape)
 35    return x, z
 36
 37
 38def fit_history_predictor(x, m, frac=.65):
 39    cut = int(len(x) * frac)
 40
 41    def xy(arr):
 42        X, Y = [], []
 43        for row in arr:
 44            for t in range(m - 1, len(row) - 1):
 45                X.append(row[t - m + 1:t + 1][::-1])
 46                Y.append(row[t + 1])
 47        return np.asarray(X), np.asarray(Y)
 48
 49    X, Y = xy(x[:cut])
 50    Xt, Yt = xy(x[cut:])
 51    w = np.linalg.solve(X.T @ X + RIDGE * np.eye(m), X.T @ Y)
 52    return w, float(np.mean((Xt @ w - Yt) ** 2))
 53
 54
 55def rollout_error(x, w, m, horizon=30):
 56    vals = []
 57    for row in x:
 58        t0 = len(row) - horizon - 1
 59        hist = list(row[t0 - m + 1:t0 + 1])
 60        pred = []
 61        for _ in range(horizon):
 62            y = float(np.dot(w, np.asarray(hist[-m:][::-1])))
 63            pred.append(y)
 64            hist.append(y)
 65        vals.append(np.mean((np.asarray(pred) - row[t0 + 1:t0 + 1 + horizon]) ** 2))
 66    return float(np.mean(vals))
 67
 68
 69def theoretical_crossing(a):
 70    return int(math.ceil(math.log(EPS) / math.log(abs(a))))
 71
 72
 73def theoretical_tau(a):
 74    # Discrete analogue of integral |rho| for rho(k)=a^k, including k=0.
 75    return 1.0 / (1.0 - abs(a))
 76
 77
 78def run():
 79    rows = []
 80    for a in (.2, .5, .7, .8, .9):
 81        x, z = make_data(a)
 82        rho = closure_rho(z)
 83        observed = first_persistent_crossing(rho)
 84        pred = theoretical_crossing(a)
 85        tau_obs = float(np.sum(np.abs(rho)))
 86        tau_pred = theoretical_tau(a)
 87        results = []
 88        for m in CANDIDATES:
 89            w, one = fit_history_predictor(x, m)
 90            results.append({'m': m, 'one_step': one, 'rollout': rollout_error(x, w, m)})
 91        scheduled = max(1, min(32, observed))
 92        ws, one = fit_history_predictor(x, scheduled)
 93        best = min(results, key=lambda q: q['one_step'])
 94        # Smallest window retaining at least 95% of the improvement from m=1 to m=32.
 95        e1, emax = results[0]['one_step'], results[-1]['one_step']
 96        target = e1 - .95 * (e1 - emax)
 97        elbow = next(q['m'] for q in results if q['one_step'] <= target)
 98        rows.append({'a': a, 'predicted_crossing': pred, 'observed_crossing': observed,
 99                     'predicted_tau': tau_pred, 'observed_tau': tau_obs,
100                     'scheduled_m': scheduled, 'elbow_m_95pct': elbow, 'best_m': best['m'],
101                     'results': results, 'scheduled_one_step': one,
102                     'scheduled_rollout': rollout_error(x, ws, scheduled)})
103
104    cross_err = [abs(q['observed_crossing'] - q['predicted_crossing']) for q in rows]
105    tau_rel_err = [abs(q['observed_tau'] - q['predicted_tau']) / q['predicted_tau'] for q in rows]
106    # Compare adaptive scheduler to Markov baseline and full 32-token model.
107    baseline = float(np.mean([q['results'][0]['one_step'] for q in rows]))
108    idea = float(np.mean([q['scheduled_one_step'] for q in rows]))
109    full = float(np.mean([q['results'][-1]['one_step'] for q in rows]))
110    summary = {
111        'crossing_mae_steps': float(np.mean(cross_err)),
112        'crossing_max_error_steps': int(max(cross_err)),
113        'tau_relative_mae': float(np.mean(tau_rel_err)),
114        'mean_baseline_m1_loss': baseline,
115        'mean_scheduler_loss': idea,
116        'mean_full_m32_loss': full,
117        'scheduler_vs_m1_percent': 100 * (baseline - idea) / baseline,
118        'scheduler_vs_m32_percent': 100 * (idea - full) / full,
119        'mean_scheduler_m': float(np.mean([q['scheduled_m'] for q in rows])),
120        'mean_elbow_m': float(np.mean([q['elbow_m_95pct'] for q in rows]))
121    }
122    out = {'config': {'epsilon': EPS, 'persistence': PERSISTENCE, 'seed': SEED,
123                      'candidates': CANDIDATES}, 'summary': summary, 'rows': rows}
124    Path('results.json').write_text(json.dumps(out, indent=2))
125    print(json.dumps(out, indent=2))
126
127
128if __name__ == '__main__':
129    run()