1import json, math, random
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED = 1380
  8np.random.seed(SEED)
  9random.seed(SEED)
 10torch.manual_seed(SEED)
 11try:
 12    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 13except Exception:
 14    device = torch.device('cpu')
 15
 16# Scalar residual map h_next=(1+a)h. Its exact finite-time Lyapunov rate is log|1+a|.
 17def lyapunov_autodiff(a, L=12, dim=8):
 18    h = torch.zeros(dim, device=device, requires_grad=True)
 19    v = torch.ones(dim, device=device)
 20    v = v / torch.linalg.vector_norm(v)
 21    rates = []
 22    for _ in range(L):
 23        def f(x):
 24            return x + a * x
 25        _, jv = torch.func.jvp(f, (h,), (v,))
 26        gain = float(torch.linalg.vector_norm(jv).detach().cpu())
 27        rates.append(math.log(max(gain, 1e-30)))
 28        v = jv / max(gain, 1e-30)
 29        h = (h + a * h).detach().requires_grad_(True)
 30    return float(np.mean(rates))
 31
 32def controller(rhat, target=0.0, kappa=0.7, dt=0.1, qmax=2.0):
 33    return float(np.clip(2*kappa*max(rhat-target, 0.0)*dt, 0, qmax))
 34
 35def mechanism_sweeps():
 36    # A: rhat=log|1+a| and the boundary rhat=0 occurs at a=0.
 37    avals = np.array([-0.45, -0.35, -0.25, -0.15, -0.05, 0.0, 0.05, 0.15, 0.25, 0.35, 0.45, 0.55, 0.65])
 38    boundary = []
 39    for a in avals:
 40        pred = math.log(abs(1+a))
 41        obs = lyapunov_autodiff(float(a))
 42        boundary.append({'a': float(a), 'pred_r': pred, 'observed_r': obs,
 43                         'abs_error': abs(pred-obs)})
 44
 45    # B: q=2*kappa*dt*[r-r_target]_+, so fitted positive-side slope is predicted.
 46    kappa, dt, target = 0.7, 0.1, 0.0
 47    rs = np.linspace(-0.4, 0.8, 25)
 48    qrows = [{'r': float(r), 'q_pred': controller(float(r), target, kappa, dt)} for r in rs]
 49    positive = [x for x in qrows if x['r'] > target]
 50    slope = float(np.polyfit([x['r'] for x in positive], [x['q_pred'] for x in positive], 1)[0])
 51
 52    # C: directly sample gate noise; conditional variance should be q*u*(1-u/K).
 53    rng = np.random.default_rng(SEED)
 54    K, uvals = 1.0, [0.2, 0.5, 0.8]
 55    samples = []
 56    for r in [0.1, 0.3, 0.6]:
 57        q = controller(r, target, kappa, dt)
 58        for u in uvals:
 59            z = 1.0 + math.sqrt(q*u*(1-u/K))*rng.standard_normal(200000)
 60            samples.append({'r': r, 'u': u, 'q': q, 'emp_var': float(np.var(z, ddof=1)),
 61                            'pred_var': q*u*(1-u/K)})
 62    x = np.array([s['pred_var'] for s in samples])
 63    y = np.array([s['emp_var'] for s in samples])
 64    var_slope = float(np.dot(x, y) / np.dot(x, x))
 65    return {'boundary': boundary, 'q_sweep': qrows, 'q_slope': slope,
 66            'q_slope_pred': 2*kappa*dt, 'gate_variance': samples,
 67            'variance_slope': var_slope}
 68
 69class ToyNet(nn.Module):
 70    def __init__(self, adaptive=False, fixed_q=0.08, depth=8):
 71        super().__init__()
 72        self.adaptive, self.fixed_q, self.depth = adaptive, fixed_q, depth
 73        self.a = nn.Parameter(torch.tensor(0.18))
 74        self.readout = nn.Linear(1, 1)
 75    def forward(self, x, return_r=False):
 76        h = x
 77        rates = []
 78        for _ in range(self.depth):
 79            r = torch.log(torch.abs(1 + self.a) + 1e-8)
 80            rates.append(r)
 81            if self.adaptive:
 82                q = torch.clamp(2*0.7*torch.relu(r.detach())*0.1, 0, 0.5)
 83            else:
 84                q = torch.tensor(self.fixed_q, device=x.device)
 85            eps = torch.randn_like(h)
 86            z = 1 + torch.sqrt(q * 0.5 * (1-0.5)) * eps
 87            h = h + z * self.a * h
 88        out = self.readout(h)
 89        return (out, torch.stack(rates).mean()) if return_r else out
 90
 91def mini_training():
 92    g = torch.Generator(device=device); g.manual_seed(SEED)
 93    x = torch.linspace(-1, 1, 96, device=device).reshape(-1, 1)
 94    y = 0.7*x + 0.1*torch.sin(5*x)
 95    results = {}
 96    for name, adaptive in [('fixed_noise', False), ('lyapunov_noise', True)]:
 97        torch.manual_seed(SEED)
 98        model = ToyNet(adaptive=adaptive).to(device)
 99        opt = torch.optim.Adam(model.parameters(), lr=0.025)
100        losses, rates, qs = [], [], []
101        for step in range(180):
102            opt.zero_grad(set_to_none=True)
103            pred, r = model(x, return_r=True)
104            loss = ((pred-y)**2).mean()
105            loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
106            opt.step()
107            losses.append(float(loss.detach().cpu())); rates.append(float(r.detach().cpu()))
108            qs.append(controller(rates[-1]) if adaptive else 0.08)
109        results[name] = {'final_loss': losses[-1], 'best_loss': min(losses),
110                         'initial_loss': losses[0], 'mean_r': float(np.mean(rates)),
111                         'mean_q': float(np.mean(qs))}
112    return results
113
114def main():
115    out = {'device': str(device), 'mechanism': mechanism_sweeps(), 'training': mini_training()}
116    Path('results.json').write_text(json.dumps(out, indent=2))
117    print(json.dumps({'device': out['device'], 'q_slope': out['mechanism']['q_slope'],
118                      'q_slope_pred': out['mechanism']['q_slope_pred'],
119                      'variance_slope': out['mechanism']['variance_slope'],
120                      'training': out['training']}, indent=2))
121
122if __name__ == '__main__':
123    main()