import os, sys, json, random, math from pathlib import Path import numpy as np import torch from torch import nn import torch.nn.functional as F sys.path.insert(0, '/home/maxwelhelp/all/math2nn') from bench import get_dataset, evaluate, sweep_baseline, make_report SEEDS = tuple(range(8)) SWEEP_SEEDS = tuple(range(4)) DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' HIDDEN = 16 OMEGA_MIN = 1e-3 def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): try: torch.cuda.manual_seed_all(seed) except Exception: pass def math_check(): # Directed 3-node path: q(s)=(s-a)(s-b), Disc(q)=(a-b)^2. eps = np.logspace(-4, -1, 7) gaps, discs, conds = [], [], [] for e in eps: a, b = 1.0 + e, 1.0 - e L = np.array([[0.,0.,0.],[-a,a,0.],[0.,-b,b]]) vals, V = np.linalg.eig(L) nz = vals[np.abs(vals) > 1e-8] gaps.append(abs(nz[0]-nz[1])/(1+abs(nz[0])+abs(nz[1]))) discs.append((a-b)**2) conds.append(np.linalg.cond(V)) gs = np.polyfit(np.log(eps), np.log(gaps), 1)[0] ds = np.polyfit(np.log(eps), np.log(discs), 1)[0] L = np.array([[0.,0.,0.],[-1.,1.,0.],[0.,-1.,1.]]) nullity = 3 - np.linalg.matrix_rank(L-np.eye(3), tol=1e-10) return {'gap_log_slope': float(gs), 'discriminant_log_slope': float(ds), 'condition_times_epsilon': [float(c*e) for c,e in zip(conds,eps)], 'exact_collision_disc': 0.0, 'collision_defective': bool(nullity < 2), 'passed': bool(abs(gs-1)<.03 and abs(ds-2)<.03 and nullity < 2)} class DirectedRNN(nn.Module): """Small recurrent model for the dynamics track; identical architecture both sides.""" def __init__(self, idea=False, hidden=HIDDEN, omega_min=OMEGA_MIN): super().__init__(); self.idea = idea; self.hidden = hidden; self.omega_min = omega_min self.inp = nn.Linear(3, hidden); self.head = nn.Linear(hidden, 1) self.edge = nn.Parameter(torch.empty(hidden, hidden)) mask = (1 - torch.eye(hidden)).float(); self.register_buffer('mask', mask) nn.init.normal_(self.edge, 0, .08) def adjacency(self): if self.idea: return (self.omega_min + F.softplus(self.edge)) * self.mask return self.edge * self.mask def laplacian_stats(self): A = self.adjacency(); L = torch.diag(A.sum(dim=1)) - A # eig is differentiable away from collisions; barrier is recomputed in loop. ev, V = torch.linalg.eig(L) evr = ev.real nz = evr[evr > 1e-5] if nz.numel() < 2: return torch.tensor(0., device=A.device), torch.tensor(1e6, device=A.device), 0. dif = (nz[:,None]-nz[None,:]).abs(); den = 1+nz[:,None].abs()+nz[None,:].abs() dif = dif + torch.eye(nz.numel(), device=A.device)*1e6 gap = (dif/den).min() cond = torch.linalg.cond(V) return gap, cond, float(gap.detach().cpu()) def forward(self, x): seq = x.view(x.shape[0], -1, 3) A = self.adjacency(); L = torch.diag(A.sum(dim=1)) - A # stable residual Laplacian flow; same recurrence for baseline and idea. M = torch.eye(self.hidden, device=x.device) - .08 * L h = torch.zeros(x.shape[0], self.hidden, device=x.device) for t in range(seq.shape[1]): h = torch.tanh(self.inp(seq[:,t]) + h @ M) return self.head(h) def run_one(seed, lr, idea, gmin=.01, beta=20., epochs=12, n_train=800, n_test=300, collect=False): seed_all(seed); d = get_dataset('dynamics', seed, n_train=n_train, n_test=n_test) model = DirectedRNN(idea=idea).to(DEVICE) xtr,ytr,xte,yte = [d[k].to(DEVICE) for k in ('xtr','ytr','xte','yte')] opt = torch.optim.Adam(model.parameters(), lr=lr) bs=128 for ep in range(epochs): model.train(); order=torch.randperm(len(xtr), device=DEVICE) for ix in order.split(bs): pred=model(xtr[ix]); task=F.mse_loss(pred,ytr[ix]); loss=task if idea: gap, cond, _ = model.laplacian_stats() loss = loss + beta*F.relu(torch.as_tensor(gmin,device=DEVICE)-gap)**2 opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 2.0); opt.step() model.eval() with torch.no_grad(): metric=float(F.mse_loss(model(xte),yte).cpu()) if collect: gap, cond, g = model.laplacian_stats() # measured trained-model signature, not an identity: modal amplification proxy return metric, {'test_mse':metric, 'gap':float(gap.detach().cpu()), 'condition':float(cond.detach().cpu()), 'gap_times_condition':float((gap*cond).detach().cpu())} return metric def main(): math = math_check(); assert math['passed'], math # Union parity: all idea learning rates are also baseline-evaluated. lrs=[0.001,0.003,0.006] grid=[{'lr':lr} for lr in lrs] def base_fn(cfg): return lambda s: run_one(s,cfg['lr'],False) base=sweep_baseline(base_fn, grid, seeds=SWEEP_SEEDS) # Explicitly run idea at all three settings; same epochs and data budget. idea_trials=[] for lr in lrs: r=evaluate(lambda s,lr=lr: run_one(s,lr,True,gmin=.01,beta=20.), seeds=SEEDS) idea_trials.append({'cfg':{'lr':lr,'gmin':.01,'beta':20.}, **r}) best=min(idea_trials,key=lambda z:z['mean']) # Signature on the paired trained systems at selected config. sig=[] for s in SEEDS: bm,bs=run_one(s,base['best_cfg']['lr'],False,collect=True) im,ins=run_one(s,best['cfg']['lr'],True,gmin=.01,beta=20.,collect=True) sig.append({'seed':s,'baseline':bs,'idea':ins}) idea_res={k:best[k] for k in ('mean','std','per_seed','n')}; idea_res['best_cfg']=best['cfg']; idea_res['trials']=idea_trials report=make_report('dynamics','directed_rnn_laplacian',base,idea_res,extra={ 'prediction':'smaller normalized spectral gap predicts larger eigenvector conditioning', 'trained_model_observations':sig, 'baseline_mean_gap':float(np.mean([z['baseline']['gap'] for z in sig])), 'idea_mean_gap':float(np.mean([z['idea']['gap'] for z in sig])), 'baseline_mean_condition':float(np.mean([z['baseline']['condition'] for z in sig])), 'idea_mean_condition':float(np.mean([z['idea']['condition'] for z in sig])), 'confirmed': bool(np.mean([z['idea']['gap'] for z in sig]) > np.mean([z['baseline']['gap'] for z in sig]) and np.mean([z['idea']['condition'] for z in sig]) < np.mean([z['baseline']['condition'] for z in sig]))}) out={'math_check':math,'bench_report':report,'device':DEVICE,'settings':{'epochs':12,'n_train':800,'n_test':300}} Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2)) if __name__=='__main__': main()