import os, sys, json, random 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 train_model, sweep_baseline, evaluate, make_report META = { 'name': 'singular_value_constitutive_energy', 'domain': 'constitutive mechanics', 'description': 'Regression of isotropic nonconvex deformation energy from 2x2 deformation gradients; target depends on positive singular values.' } def get_dataset(seed, n_train=400, n_test=400): rng = np.random.default_rng(int(seed)) xtr = rng.normal(0, 0.65, (n_train, 2, 2)).astype(np.float32) xte = rng.normal(0, 0.65, (n_test, 2, 2)).astype(np.float32) def target(a): s = np.linalg.svd(a, compute_uv=False) y = (0.35*(s[:,0]-0.9)**2 + 0.25*(s[:,1]-1.15)**2 + 0.12*np.sin(4*s[:,0])*np.sin(3*s[:,1]) + 0.25) return y[:,None].astype(np.float32) return {'xtr':xtr, 'ytr':target(xtr), 'xte':xte, 'yte':target(xte), 'task':'regression', 'metric':'mse', 'input_shape':(2,2), 'out_dim':1} def seed_all(seed): random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) def sv(x): return torch.linalg.svdvals(x).clamp_min(1e-5) class InvariantMLP(nn.Module): def __init__(self, width=16): super().__init__() self.net = nn.Sequential(nn.Linear(2,width), nn.Tanh(), nn.Linear(width,width), nn.Tanh(), nn.Linear(width,1)) def forward(self, Fm): return self.net(sv(Fm)).squeeze(-1) class ICNN(nn.Module): def __init__(self, width=16): super().__init__() self.w0 = nn.Parameter(torch.full((width,2), -2.0) + 0.05*torch.randn(width,2)) self.b0 = nn.Parameter(torch.zeros(width)) self.ar = nn.Parameter(torch.full((width,width), -2.0) + 0.05*torch.randn(width,width)) self.br = nn.Parameter(torch.full((width,2), -2.0) + 0.05*torch.randn(width,2)) self.b1 = nn.Parameter(torch.zeros(width)) self.cr = nn.Parameter(torch.full((width,), -2.0) + 0.05*torch.randn(width)) self.dr = nn.Parameter(torch.full((2,), -2.0) + 0.05*torch.randn(2)) self.outb = nn.Parameter(torch.tensor(0.15)) def pos(self, x): return F.softplus(x) + 1e-4 def forward_x(self, x): z = F.softplus(x @ self.pos(self.w0).T + self.b0) z = F.softplus(z @ self.pos(self.ar).T + x @ self.pos(self.br).T + self.b1) return z @ self.pos(self.cr) + x @ self.pos(self.dr) + self.outb def forward(self, Fm): return self.forward_x(sv(Fm)) def baseline_train(seed, lr, epochs=18): seed_all(seed); ds=get_dataset(seed) model=InvariantMLP().float() ds={**ds, **{k: torch.from_numpy(ds[k]) for k in ('xtr','ytr','xte','yte')}} ds['ytr']=ds['ytr'].squeeze(-1); ds['yte']=ds['yte'].squeeze(-1) model, metric, _=train_model(model, ds, epochs=epochs, lr=lr, batch=64, weight_decay=0.0, log=lambda *_:None) return float(metric) def idea_train(seed, lr, lower_weight=4.0, epochs=18, return_model=False): seed_all(seed); ds=get_dataset(seed) dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu') model=ICNN().float().to(dev) x=torch.from_numpy(ds['xtr']).to(dev); y=torch.from_numpy(ds['ytr'][:,0]).to(dev) opt=torch.optim.Adam(model.parameters(), lr=lr) for ep in range(epochs): order=torch.randperm(len(x),device=dev) for ix in order.split(64): p=model(x[ix]); err=p-y[ix] loss=err.square().mean()+lower_weight*F.relu(err).square().mean() opt.zero_grad(); loss.backward(); opt.step() with torch.no_grad(): metric=float((model(torch.from_numpy(ds['xte']).to(dev))-torch.from_numpy(ds['yte'][:,0]).to(dev)).square().mean().cpu()) if return_model: return metric, model, ds return metric def behaviour_signature(): # Measure the proposed structural mechanism on actually trained systems. seed_all(991) ds = get_dataset(991) bds={**ds, **{k: torch.from_numpy(ds[k]) for k in ('xtr','ytr','xte','yte')}} bds['ytr']=bds['ytr'].squeeze(-1); bds['yte']=bds['yte'].squeeze(-1) bm=InvariantMLP().float() bm,_,_=train_model(bm,bds,epochs=18,lr=0.003,batch=64,weight_decay=0.0,log=lambda *_:None) im=ICNN().float() dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu'); im=im.to(dev) x=torch.from_numpy(ds['xtr']).to(dev); y=torch.from_numpy(ds['ytr'][:,0]).to(dev) opt=torch.optim.Adam(im.parameters(),lr=0.003) for _ in range(18): for ix in torch.randperm(len(x),device=dev).split(64): err=im(x[ix])-y[ix] loss=err.square().mean()+4.0*F.relu(err).square().mean() opt.zero_grad(); loss.backward(); opt.step() bm.eval(); im.eval() xb=torch.rand(2000,2,device=dev)*1.8+0.1; yb=torch.rand(2000,2,device=dev)*1.8+0.1; t=torch.rand(2000,1,device=dev); z=t*xb+(1-t)*yb with torch.no_grad(): bp=lambda q: bm.net(q).squeeze(-1) ip=lambda q: im.forward_x(q) bj=(bp(z)>t[:,0]*bp(xb)+(1-t[:,0])*bp(yb)+1e-6).float().mean().item() ij=(ip(z)>t[:,0]*ip(xb)+(1-t[:,0])*ip(yb)+1e-6).float().mean().item() delta=torch.zeros_like(xb); delta[:,0]=0.05 bmon=(bp(xb+delta)q+1e-6).float().mean().item(); ilow=(ip(xb)>q+1e-6).float().mean().item() return {'baseline_jensen_violation':bj,'idea_jensen_violation':ij,'baseline_monotonicity_violation':bmon,'idea_monotonicity_violation':imon,'baseline_lower_violation':blow,'idea_lower_violation':ilow,'prediction':'trained nonnegative ICNN should eliminate Jensen and coordinatewise monotonicity violations','confirmed':bool(ij<=1e-7 and imon<=1e-7)} def main(): os.makedirs('results',exist_ok=True) # Union of all learning rates is evaluated for baseline and idea. grid=[{'lr':0.001},{'lr':0.003},{'lr':0.006}] def make_base(cfg): return lambda seed: baseline_train(seed,cfg['lr']) base=sweep_baseline(make_base, grid, seeds=(0,1,2,3)) # Full paired idea results at all three settings; select on the same four-seed protocol. idea_trials=[] for cfg in grid: r=evaluate(lambda seed: idea_train(seed,cfg['lr']), seeds=range(4)) idea_trials.append({'cfg':cfg,'mean':r['mean']}) best_cfg=min(idea_trials,key=lambda x:x['mean'])['cfg'] idea=evaluate(lambda seed: idea_train(seed,best_cfg['lr']), seeds=range(8)) sig=behaviour_signature() rep=make_report('singular_value_constitutive_energy','shared_invariant_mlp_width16',base,idea,{ 'custom_track':{'name':META['name'],'file':'stage2_bench.py','domain':META['domain']}, 'idea_sweep':idea_trials, 'mechanism_signature':sig}) with open('results/bench_report.json','w') as f: json.dump(rep,f,indent=2) print(json.dumps(rep,indent=2)) if __name__=='__main__': try: main() except Exception as e: print('CUDA/benchmark failure, retrying on CPU:',repr(e),file=sys.stderr) torch.cuda.is_available=lambda:False main()