1import os, sys, json, random
  2import numpy as np
  3import torch
  4from torch import nn
  5import torch.nn.functional as F
  6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  7from bench import train_model, sweep_baseline, evaluate, make_report
  8
  9META = {
 10    'name': 'singular_value_constitutive_energy',
 11    'domain': 'constitutive mechanics',
 12    'description': 'Regression of isotropic nonconvex deformation energy from 2x2 deformation gradients; target depends on positive singular values.'
 13}
 14
 15def get_dataset(seed, n_train=400, n_test=400):
 16    rng = np.random.default_rng(int(seed))
 17    xtr = rng.normal(0, 0.65, (n_train, 2, 2)).astype(np.float32)
 18    xte = rng.normal(0, 0.65, (n_test, 2, 2)).astype(np.float32)
 19    def target(a):
 20        s = np.linalg.svd(a, compute_uv=False)
 21        y = (0.35*(s[:,0]-0.9)**2 + 0.25*(s[:,1]-1.15)**2
 22             + 0.12*np.sin(4*s[:,0])*np.sin(3*s[:,1]) + 0.25)
 23        return y[:,None].astype(np.float32)
 24    return {'xtr':xtr, 'ytr':target(xtr), 'xte':xte, 'yte':target(xte),
 25            'task':'regression', 'metric':'mse', 'input_shape':(2,2), 'out_dim':1}
 26
 27def seed_all(seed):
 28    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 29    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 30
 31def sv(x):
 32    return torch.linalg.svdvals(x).clamp_min(1e-5)
 33
 34class InvariantMLP(nn.Module):
 35    def __init__(self, width=16):
 36        super().__init__()
 37        self.net = nn.Sequential(nn.Linear(2,width), nn.Tanh(), nn.Linear(width,width), nn.Tanh(), nn.Linear(width,1))
 38    def forward(self, Fm): return self.net(sv(Fm)).squeeze(-1)
 39
 40class ICNN(nn.Module):
 41    def __init__(self, width=16):
 42        super().__init__()
 43        self.w0 = nn.Parameter(torch.full((width,2), -2.0) + 0.05*torch.randn(width,2))
 44        self.b0 = nn.Parameter(torch.zeros(width))
 45        self.ar = nn.Parameter(torch.full((width,width), -2.0) + 0.05*torch.randn(width,width))
 46        self.br = nn.Parameter(torch.full((width,2), -2.0) + 0.05*torch.randn(width,2))
 47        self.b1 = nn.Parameter(torch.zeros(width))
 48        self.cr = nn.Parameter(torch.full((width,), -2.0) + 0.05*torch.randn(width))
 49        self.dr = nn.Parameter(torch.full((2,), -2.0) + 0.05*torch.randn(2))
 50        self.outb = nn.Parameter(torch.tensor(0.15))
 51    def pos(self, x): return F.softplus(x) + 1e-4
 52    def forward_x(self, x):
 53        z = F.softplus(x @ self.pos(self.w0).T + self.b0)
 54        z = F.softplus(z @ self.pos(self.ar).T + x @ self.pos(self.br).T + self.b1)
 55        return z @ self.pos(self.cr) + x @ self.pos(self.dr) + self.outb
 56    def forward(self, Fm): return self.forward_x(sv(Fm))
 57
 58def baseline_train(seed, lr, epochs=18):
 59    seed_all(seed); ds=get_dataset(seed)
 60    model=InvariantMLP().float()
 61    ds={**ds, **{k: torch.from_numpy(ds[k]) for k in ('xtr','ytr','xte','yte')}}
 62    ds['ytr']=ds['ytr'].squeeze(-1); ds['yte']=ds['yte'].squeeze(-1)
 63    model, metric, _=train_model(model, ds, epochs=epochs, lr=lr, batch=64, weight_decay=0.0, log=lambda *_:None)
 64    return float(metric)
 65
 66def idea_train(seed, lr, lower_weight=4.0, epochs=18, return_model=False):
 67    seed_all(seed); ds=get_dataset(seed)
 68    dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 69    model=ICNN().float().to(dev)
 70    x=torch.from_numpy(ds['xtr']).to(dev); y=torch.from_numpy(ds['ytr'][:,0]).to(dev)
 71    opt=torch.optim.Adam(model.parameters(), lr=lr)
 72    for ep in range(epochs):
 73        order=torch.randperm(len(x),device=dev)
 74        for ix in order.split(64):
 75            p=model(x[ix]); err=p-y[ix]
 76            loss=err.square().mean()+lower_weight*F.relu(err).square().mean()
 77            opt.zero_grad(); loss.backward(); opt.step()
 78    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())
 79    if return_model: return metric, model, ds
 80    return metric
 81
 82def behaviour_signature():
 83    # Measure the proposed structural mechanism on actually trained systems.
 84    seed_all(991)
 85    ds = get_dataset(991)
 86    bds={**ds, **{k: torch.from_numpy(ds[k]) for k in ('xtr','ytr','xte','yte')}}
 87    bds['ytr']=bds['ytr'].squeeze(-1); bds['yte']=bds['yte'].squeeze(-1)
 88    bm=InvariantMLP().float()
 89    bm,_,_=train_model(bm,bds,epochs=18,lr=0.003,batch=64,weight_decay=0.0,log=lambda *_:None)
 90    im=ICNN().float()
 91    dev=torch.device('cuda' if torch.cuda.is_available() else 'cpu'); im=im.to(dev)
 92    x=torch.from_numpy(ds['xtr']).to(dev); y=torch.from_numpy(ds['ytr'][:,0]).to(dev)
 93    opt=torch.optim.Adam(im.parameters(),lr=0.003)
 94    for _ in range(18):
 95        for ix in torch.randperm(len(x),device=dev).split(64):
 96            err=im(x[ix])-y[ix]
 97            loss=err.square().mean()+4.0*F.relu(err).square().mean()
 98            opt.zero_grad(); loss.backward(); opt.step()
 99    bm.eval(); im.eval()
100    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
101    with torch.no_grad():
102        bp=lambda q: bm.net(q).squeeze(-1)
103        ip=lambda q: im.forward_x(q)
104        bj=(bp(z)>t[:,0]*bp(xb)+(1-t[:,0])*bp(yb)+1e-6).float().mean().item()
105        ij=(ip(z)>t[:,0]*ip(xb)+(1-t[:,0])*ip(yb)+1e-6).float().mean().item()
106        delta=torch.zeros_like(xb); delta[:,0]=0.05
107        bmon=(bp(xb+delta)<bp(xb)-1e-6).float().mean().item(); imon=(ip(xb+delta)<ip(xb)-1e-6).float().mean().item()
108        q=0.35*(xb[:,0]-0.9)**2+0.25*(xb[:,1]-1.15)**2+0.12*torch.sin(4*xb[:,0])*torch.sin(3*xb[:,1])+0.25
109        blow=(bp(xb)>q+1e-6).float().mean().item(); ilow=(ip(xb)>q+1e-6).float().mean().item()
110    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)}
111
112def main():
113    os.makedirs('results',exist_ok=True)
114    # Union of all learning rates is evaluated for baseline and idea.
115    grid=[{'lr':0.001},{'lr':0.003},{'lr':0.006}]
116    def make_base(cfg): return lambda seed: baseline_train(seed,cfg['lr'])
117    base=sweep_baseline(make_base, grid, seeds=(0,1,2,3))
118    # Full paired idea results at all three settings; select on the same four-seed protocol.
119    idea_trials=[]
120    for cfg in grid:
121        r=evaluate(lambda seed: idea_train(seed,cfg['lr']), seeds=range(4))
122        idea_trials.append({'cfg':cfg,'mean':r['mean']})
123    best_cfg=min(idea_trials,key=lambda x:x['mean'])['cfg']
124    idea=evaluate(lambda seed: idea_train(seed,best_cfg['lr']), seeds=range(8))
125    sig=behaviour_signature()
126    rep=make_report('singular_value_constitutive_energy','shared_invariant_mlp_width16',base,idea,{
127        'custom_track':{'name':META['name'],'file':'stage2_bench.py','domain':META['domain']},
128        'idea_sweep':idea_trials,
129        'mechanism_signature':sig})
130    with open('results/bench_report.json','w') as f: json.dump(rep,f,indent=2)
131    print(json.dumps(rep,indent=2))
132
133if __name__=='__main__':
134    try: main()
135    except Exception as e:
136        print('CUDA/benchmark failure, retrying on CPU:',repr(e),file=sys.stderr)
137        torch.cuda.is_available=lambda:False
138        main()