1import os, sys, json, random, math
  2from pathlib import Path
  3import numpy as np
  4import torch
  5from torch import nn
  6import torch.nn.functional as F
  7
  8sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  9from bench import get_dataset, evaluate, sweep_baseline, make_report
 10
 11SEEDS = tuple(range(8))
 12SWEEP_SEEDS = tuple(range(4))
 13DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 14HIDDEN = 16
 15OMEGA_MIN = 1e-3
 16
 17
 18def seed_all(seed):
 19    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 20    if torch.cuda.is_available():
 21        try: torch.cuda.manual_seed_all(seed)
 22        except Exception: pass
 23
 24
 25def math_check():
 26    # Directed 3-node path: q(s)=(s-a)(s-b), Disc(q)=(a-b)^2.
 27    eps = np.logspace(-4, -1, 7)
 28    gaps, discs, conds = [], [], []
 29    for e in eps:
 30        a, b = 1.0 + e, 1.0 - e
 31        L = np.array([[0.,0.,0.],[-a,a,0.],[0.,-b,b]])
 32        vals, V = np.linalg.eig(L)
 33        nz = vals[np.abs(vals) > 1e-8]
 34        gaps.append(abs(nz[0]-nz[1])/(1+abs(nz[0])+abs(nz[1])))
 35        discs.append((a-b)**2)
 36        conds.append(np.linalg.cond(V))
 37    gs = np.polyfit(np.log(eps), np.log(gaps), 1)[0]
 38    ds = np.polyfit(np.log(eps), np.log(discs), 1)[0]
 39    L = np.array([[0.,0.,0.],[-1.,1.,0.],[0.,-1.,1.]])
 40    nullity = 3 - np.linalg.matrix_rank(L-np.eye(3), tol=1e-10)
 41    return {'gap_log_slope': float(gs), 'discriminant_log_slope': float(ds),
 42            'condition_times_epsilon': [float(c*e) for c,e in zip(conds,eps)],
 43            'exact_collision_disc': 0.0, 'collision_defective': bool(nullity < 2),
 44            'passed': bool(abs(gs-1)<.03 and abs(ds-2)<.03 and nullity < 2)}
 45
 46
 47class DirectedRNN(nn.Module):
 48    """Small recurrent model for the dynamics track; identical architecture both sides."""
 49    def __init__(self, idea=False, hidden=HIDDEN, omega_min=OMEGA_MIN):
 50        super().__init__(); self.idea = idea; self.hidden = hidden; self.omega_min = omega_min
 51        self.inp = nn.Linear(3, hidden); self.head = nn.Linear(hidden, 1)
 52        self.edge = nn.Parameter(torch.empty(hidden, hidden))
 53        mask = (1 - torch.eye(hidden)).float(); self.register_buffer('mask', mask)
 54        nn.init.normal_(self.edge, 0, .08)
 55
 56    def adjacency(self):
 57        if self.idea: return (self.omega_min + F.softplus(self.edge)) * self.mask
 58        return self.edge * self.mask
 59
 60    def laplacian_stats(self):
 61        A = self.adjacency(); L = torch.diag(A.sum(dim=1)) - A
 62        # eig is differentiable away from collisions; barrier is recomputed in loop.
 63        ev, V = torch.linalg.eig(L)
 64        evr = ev.real
 65        nz = evr[evr > 1e-5]
 66        if nz.numel() < 2: return torch.tensor(0., device=A.device), torch.tensor(1e6, device=A.device), 0.
 67        dif = (nz[:,None]-nz[None,:]).abs(); den = 1+nz[:,None].abs()+nz[None,:].abs()
 68        dif = dif + torch.eye(nz.numel(), device=A.device)*1e6
 69        gap = (dif/den).min()
 70        cond = torch.linalg.cond(V)
 71        return gap, cond, float(gap.detach().cpu())
 72
 73    def forward(self, x):
 74        seq = x.view(x.shape[0], -1, 3)
 75        A = self.adjacency(); L = torch.diag(A.sum(dim=1)) - A
 76        # stable residual Laplacian flow; same recurrence for baseline and idea.
 77        M = torch.eye(self.hidden, device=x.device) - .08 * L
 78        h = torch.zeros(x.shape[0], self.hidden, device=x.device)
 79        for t in range(seq.shape[1]): h = torch.tanh(self.inp(seq[:,t]) + h @ M)
 80        return self.head(h)
 81
 82
 83def run_one(seed, lr, idea, gmin=.01, beta=20., epochs=12, n_train=800, n_test=300, collect=False):
 84    seed_all(seed); d = get_dataset('dynamics', seed, n_train=n_train, n_test=n_test)
 85    model = DirectedRNN(idea=idea).to(DEVICE)
 86    xtr,ytr,xte,yte = [d[k].to(DEVICE) for k in ('xtr','ytr','xte','yte')]
 87    opt = torch.optim.Adam(model.parameters(), lr=lr)
 88    bs=128
 89    for ep in range(epochs):
 90        model.train(); order=torch.randperm(len(xtr), device=DEVICE)
 91        for ix in order.split(bs):
 92            pred=model(xtr[ix]); task=F.mse_loss(pred,ytr[ix]); loss=task
 93            if idea:
 94                gap, cond, _ = model.laplacian_stats()
 95                loss = loss + beta*F.relu(torch.as_tensor(gmin,device=DEVICE)-gap)**2
 96            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 2.0); opt.step()
 97    model.eval()
 98    with torch.no_grad(): metric=float(F.mse_loss(model(xte),yte).cpu())
 99    if collect:
100        gap, cond, g = model.laplacian_stats()
101        # measured trained-model signature, not an identity: modal amplification proxy
102        return metric, {'test_mse':metric, 'gap':float(gap.detach().cpu()), 'condition':float(cond.detach().cpu()),
103                        'gap_times_condition':float((gap*cond).detach().cpu())}
104    return metric
105
106
107def main():
108    math = math_check(); assert math['passed'], math
109    # Union parity: all idea learning rates are also baseline-evaluated.
110    lrs=[0.001,0.003,0.006]
111    grid=[{'lr':lr} for lr in lrs]
112    def base_fn(cfg): return lambda s: run_one(s,cfg['lr'],False)
113    base=sweep_baseline(base_fn, grid, seeds=SWEEP_SEEDS)
114    # Explicitly run idea at all three settings; same epochs and data budget.
115    idea_trials=[]
116    for lr in lrs:
117        r=evaluate(lambda s,lr=lr: run_one(s,lr,True,gmin=.01,beta=20.), seeds=SEEDS)
118        idea_trials.append({'cfg':{'lr':lr,'gmin':.01,'beta':20.}, **r})
119    best=min(idea_trials,key=lambda z:z['mean'])
120    # Signature on the paired trained systems at selected config.
121    sig=[]
122    for s in SEEDS:
123        bm,bs=run_one(s,base['best_cfg']['lr'],False,collect=True)
124        im,ins=run_one(s,best['cfg']['lr'],True,gmin=.01,beta=20.,collect=True)
125        sig.append({'seed':s,'baseline':bs,'idea':ins})
126    idea_res={k:best[k] for k in ('mean','std','per_seed','n')}; idea_res['best_cfg']=best['cfg']; idea_res['trials']=idea_trials
127    report=make_report('dynamics','directed_rnn_laplacian',base,idea_res,extra={
128        'prediction':'smaller normalized spectral gap predicts larger eigenvector conditioning',
129        'trained_model_observations':sig,
130        'baseline_mean_gap':float(np.mean([z['baseline']['gap'] for z in sig])),
131        'idea_mean_gap':float(np.mean([z['idea']['gap'] for z in sig])),
132        'baseline_mean_condition':float(np.mean([z['baseline']['condition'] for z in sig])),
133        'idea_mean_condition':float(np.mean([z['idea']['condition'] for z in sig])),
134        'confirmed': bool(np.mean([z['idea']['gap'] for z in sig]) > np.mean([z['baseline']['gap'] for z in sig]) and
135                          np.mean([z['idea']['condition'] for z in sig]) < np.mean([z['baseline']['condition'] for z in sig]))})
136    out={'math_check':math,'bench_report':report,'device':DEVICE,'settings':{'epochs':12,'n_train':800,'n_test':300}}
137    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
138
139if __name__=='__main__': main()