1import json, math, random
 2from pathlib import Path
 3import numpy as np
 4
 5SEED = 2378
 6C = 0.587
 7SQRT2PI = math.sqrt(2.0 * math.pi)
 8random.seed(SEED); np.random.seed(SEED)
 9
10def normal_cdf(x):
11    return 0.5 * (1.0 + np.vectorize(math.erf)(x / math.sqrt(2.0)))
12
13def run_math_check():
14    rng = np.random.default_rng(SEED); rows = []
15    for d in (1, 2, 8):
16        x = rng.normal(size=(200000, d))
17        jac = 2.0 / SQRT2PI * np.exp(-0.5 * x*x)
18        empirical = float(np.max(jac)); directional = C * 2.0
19        global_bound = C * 2.0 * math.sqrt(d)
20        rows.append({'d': d, 'empirical_max_jacobian': empirical,
21                     'exact_supremum': math.sqrt(2.0/math.pi),
22                     'directional_bound': directional, 'global_bound': global_bound,
23                     'directional_holds': empirical <= directional,
24                     'global_holds': empirical <= global_bound})
25    return rows
26
27def train_comparison(d, steps=400, batch=128):
28    try:
29        import torch
30        torch.manual_seed(SEED+d)
31        device = 'cuda' if torch.cuda.is_available() else 'cpu'
32        try: _ = torch.zeros(1, device=device) + 1
33        except Exception: device = 'cpu'
34    except Exception as e: return {'d': d, 'error': str(e)}
35    rng = np.random.default_rng(SEED+d)
36    x_np = rng.normal(size=(4096,d)).astype('float32')
37    target_np = (2*normal_cdf(x_np)-1).astype('float32')
38    x = torch.tensor(x_np, device=device); target = torch.tensor(target_np, device=device)
39    L = C*2*math.sqrt(d); out = {'d': d, 'device': device, 'L_star': L, 'initial_scale': 4.0}
40    for name, capped in [('baseline_uncapped',False), ('certificate_projected',True)]:
41        torch.manual_seed(SEED+d+(1000 if capped else 0))
42        scale = torch.nn.Parameter(torch.tensor(4.0,device=device)); opt = torch.optim.Adam([scale],lr=.08)
43        post_peak = 0.; violations = 0; losses=[]
44        for _ in range(steps):
45            idx=torch.randint(0,x.shape[0],(batch,),device=device)
46            loss=((scale*torch.tanh(x[idx])-target[idx])**2).mean()
47            opt.zero_grad(); loss.backward(); opt.step()
48            if capped:
49                with torch.no_grad(): scale.clamp_(0.,L)
50            with torch.no_grad():
51                a=float(scale.cpu()); post_peak=max(post_peak,a); violations += int(a>L+1e-6)
52                losses.append(float(loss.cpu()))
53        with torch.no_grad():
54            a=float(scale.cpu()); mse=float(((a*torch.tanh(x)-target)**2).mean().cpu())
55        out[name]={'final_scale':a,'post_update_peak_jacobian':post_peak,
56                   'cap_violations_after_updates':violations,'final_mse':mse,
57                   'mean_last50_mse':float(np.mean(losses[-50:]))}
58    return out
59
60def main():
61    out={'math_check':run_math_check(), 'experiments':[train_comparison(d) for d in (1,2,8)]}
62    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
63if __name__=='__main__': main()