1import json
 2from pathlib import Path
 3import numpy as np
 4import torch
 5
 6SEED = 7
 7np.random.seed(SEED)
 8torch.manual_seed(SEED)
 9
10
11def toy_check():
12    a = 0.8
13    ns = np.array([2, 4, 8, 16, 32, 64, 128, 256])
14    cumulative = np.array([(1-a**int(n))/(1-a) for n in ns])
15    mean_gap = cumulative / ns
16    slope = float(np.polyfit(np.log(ns[2:]), np.log(mean_gap[2:]), 1)[0])
17    return {'a': a, 'N': ns.tolist(), 'cumulative_gap': cumulative.tolist(),
18            'mean_gap': mean_gap.tolist(), 'loglog_slope': slope,
19            'theoretical_C': 1/(1-a)}
20
21
22def make_data(T=240, a=.82, b=.55, noise=.025):
23    rng = np.random.default_rng(SEED)
24    u = rng.normal(size=T).astype(np.float32)
25    x = np.zeros(T+1, dtype=np.float32); y = np.zeros(T, dtype=np.float32)
26    x[0] = .7
27    for k in range(T):
28        y[k] = x[k] + noise*rng.normal(); x[k+1] = a*x[k] + b*u[k]
29    return torch.tensor(u), torch.tensor(y)
30
31
32def rollout(theta, x0, u):
33    a, b = theta[0], theta[1]; x, out = x0, []
34    for uk in u:
35        out.append(x); x = a*x + b*uk
36    return torch.stack(out)
37
38
39def estimate_gap(theta, u, y, x0, horizon, inner_steps=20):
40    z = torch.nn.Parameter(x0.detach().clone())
41    inner = torch.optim.Adam([z], lr=.08)
42    uu, yy = u[:horizon], y[:horizon]
43    for _ in range(inner_steps):
44        inner.zero_grad(); loss = ((rollout(theta.detach(), z, uu)-yy)**2).mean()
45        loss.backward(); inner.step()
46    with torch.no_grad():
47        fixed, free = rollout(theta, x0, uu), rollout(theta, z, uu)
48        return (fixed-free).abs().sum().item(), ((free-yy)**2).mean().item(), float(z)
49
50
51def fit(method, u, y, steps=600, base_horizon=40, eps=.05, every=25):
52    theta = torch.nn.Parameter(torch.tensor([.35, .25])); opt = torch.optim.Adam([theta], lr=.018)
53    fixed_x0 = torch.tensor(0.0); chosen, chat = base_horizon, 0.0; hist = []
54    for step in range(steps):
55        if method == 'calibrated' and step % every == 0:
56            c, _, _ = estimate_gap(theta, u, y, fixed_x0, base_horizon)
57            chat = .8*chat + .2*c
58            chosen = min(base_horizon, max(4, int(np.ceil(chat/eps))))
59        n = base_horizon if method == 'full' else (8 if method == 'fixed_8' else chosen)
60        start = (3*step) % (len(u)-n+1)
61        loss = ((rollout(theta, fixed_x0, u[start:start+n])-y[start:start+n])**2).mean()
62        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_([theta], 5.0); opt.step()
63        if step % 50 == 0: hist.append([step, float(loss), int(n), float(chat)])
64    with torch.no_grad(): mse = ((rollout(theta, fixed_x0, u)-y)**2).mean().item()
65    gap, free_loss, free_x0 = estimate_gap(theta, u, y, fixed_x0, base_horizon)
66    return {'mse': mse, 'theta': theta.detach().numpy().tolist(), 'mean_fixed_free_gap': gap/base_horizon,
67            'free_loss': free_loss, 'free_x0': free_x0, 'selected_horizon': int(chosen), 'history': hist}
68
69
70def empirical_sweep(theta, u, y):
71    ns = [2, 4, 8, 16, 32, 40]
72    gaps = [estimate_gap(theta, u, y, torch.tensor(0.0), n)[0]/n for n in ns]
73    slope = float(np.polyfit(np.log(ns[1:]), np.log(gaps[1:]), 1)[0])
74    return {'N': ns, 'mean_gap': gaps, 'loglog_slope': slope}
75
76
77def main():
78    u, y = make_data()
79    full = fit('full', u, y)
80    fixed = fit('fixed_8', u, y)
81    calibrated = fit('calibrated', u, y, eps=.05)
82    relaxed = fit('calibrated', u, y, eps=.20)
83    out = {'toy_check': toy_check(), 'training': {'full_40': full, 'fixed_8': fixed,
84           'calibrated_eps_.05': calibrated, 'calibrated_eps_.20': relaxed},
85           'empirical_gap_sweep_full_model': empirical_sweep(torch.tensor(full['theta']), u, y)}
86    Path('results.json').write_text(json.dumps(out, indent=2)); print(json.dumps(out, indent=2))
87
88if __name__ == '__main__': main()