1import json, math, os, random, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 320
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_default_dtype(torch.float32)
  9try:
 10    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 11except Exception:
 12    device = torch.device('cpu')
 13
 14# A deliberately nonconvex/nonmonotone synthetic energy on positive singular values.
 15def target(x):
 16    # Positive baseline plus localized oscillatory wells; the ICNN is trained as a lower fit.
 17    a, b = x[:, 0], x[:, 1]
 18    return 0.35*(a-1.05)**2 + 0.28*(b-1.25)**2 + 0.16*torch.sin(5.0*a)*torch.sin(4.0*b) + 0.18
 19
 20class MLP(nn.Module):
 21    def __init__(self, width=24):
 22        super().__init__()
 23        self.net = nn.Sequential(nn.Linear(2,width), nn.Tanh(), nn.Linear(width,width), nn.Tanh(), nn.Linear(width,1))
 24    def forward(self,x): return self.net(x).squeeze(-1)
 25
 26class ICNN(nn.Module):
 27    """z_l=softplus(A_l z+B_l x+b), with A,B >= 0; positive output skips preserve monotonicity."""
 28    def __init__(self, width=20):
 29        super().__init__()
 30        self.w0 = nn.Parameter(torch.full((width,2), -4.0) + torch.randn(width,2)*0.05)
 31        self.b0 = nn.Parameter(torch.zeros(width))
 32        self.Araw = nn.ParameterList([nn.Parameter(torch.full((width,width), -4.0) + torch.randn(width,width)*0.05)])
 33        self.Braw = nn.ParameterList([nn.Parameter(torch.full((width,2), -4.0) + torch.randn(width,2)*0.05)])
 34        self.b1 = nn.Parameter(torch.zeros(width))
 35        self.craw = nn.Parameter(torch.full((width,), -4.0) + torch.randn(width)*0.05)
 36        self.draw = nn.Parameter(torch.full((2,), -4.0) + torch.randn(2)*0.05)
 37        self.outb = nn.Parameter(torch.tensor(0.15))
 38    @staticmethod
 39    def pos(v): return torch.nn.functional.softplus(v) + 1e-4
 40    def forward(self,x):
 41        # First layer has nonnegative direct input weights.
 42        z = torch.nn.functional.softplus(x @ self.pos(self.w0).T + self.b0)
 43        A, B = self.pos(self.Araw[0]), self.pos(self.Braw[0])
 44        z = torch.nn.functional.softplus(z @ A.T + x @ B.T + self.b1)
 45        return z @ self.pos(self.craw) + x @ self.pos(self.draw) + self.outb
 46
 47def energy_from_F(model, F, eps=1e-6):
 48    s = torch.linalg.svdvals(F).clamp_min(eps)
 49    return model(s), s
 50
 51def violation_rates(model, n=4000):
 52    x = 0.2 + 1.8*torch.rand(n,2,device=device)
 53    y = 0.2 + 1.8*torch.rand(n,2,device=device)
 54    t = torch.rand(n,1,device=device)
 55    z = t*x + (1-t)*y
 56    with torch.no_grad():
 57        convex_v = (model(z) > t.squeeze()*model(x)+(1-t.squeeze())*model(y)+1e-5).float().mean().item()
 58        mono_v = []
 59        for j in range(2):
 60            d = torch.zeros_like(x); d[:,j] = 0.05
 61            mono_v.append((model(x+d) < model(x)-1e-5).float().mean().item())
 62        q = target(x)
 63        lower_v = (model(x) > q+1e-4).float().mean().item()
 64    return convex_v, float(np.mean(mono_v)), lower_v
 65
 66def train(model, lower=False, steps=1800):
 67    model.to(device); opt = torch.optim.Adam(model.parameters(), lr=1e-3)
 68    for step in range(steps):
 69        x = 0.2 + 1.8*torch.rand(128,2,device=device)
 70        q = target(x); p = model(x)
 71        # Exact requested lower-envelope penalty, with a modest value term.
 72        if lower:
 73            loss = ((p-q)**2).mean() + 4.0*torch.relu(p-q).square().mean()
 74        else:
 75            loss = ((p-q)**2).mean()
 76        opt.zero_grad(); loss.backward(); opt.step()
 77    return float(loss.detach().cpu())
 78
 79def main():
 80    # Stage-1 math sanity: random ICNN parameters should satisfy the architectural claims.
 81    torch.manual_seed(SEED)
 82    icnn = ICNN().to(device)
 83    sanity = violation_rates(icnn, 12000)
 84    # Fit baseline and proposed model under identical sampled training protocol.
 85    torch.manual_seed(SEED+1); base = MLP().to(device)
 86    t0=time.perf_counter(); train(base, False, steps=1800); base_time=time.perf_counter()-t0
 87    torch.manual_seed(SEED+2); idea = ICNN().to(device)
 88    t0=time.perf_counter(); train(idea, True, steps=1800); idea_time=time.perf_counter()-t0
 89    base_v = violation_rates(base); idea_v = violation_rates(idea)
 90    grid1=torch.linspace(.2,2.0,100,device=device); grid2=torch.linspace(.2,2.0,100,device=device)
 91    xx,yy=torch.meshgrid(grid1,grid2,indexing='ij'); grid=torch.stack([xx.ravel(),yy.ravel()],1)
 92    with torch.no_grad():
 93        qb=target(grid); pb=base(grid); pi=idea(grid)
 94        metrics={
 95          'baseline_rmse': float(torch.sqrt(((pb-qb)**2).mean()).cpu()),
 96          'idea_rmse': float(torch.sqrt(((pi-qb)**2).mean()).cpu()),
 97          'baseline_grid_lower_violation': float((pb>qb+1e-4).float().mean().cpu()),
 98          'idea_grid_lower_violation': float((pi>qb+1e-4).float().mean().cpu()),
 99          'baseline_grid_mean_signed_error': float((pb-qb).mean().cpu()),
100          'idea_grid_mean_signed_error': float((pi-qb).mean().cpu()),
101          'baseline_seconds':base_time, 'idea_seconds':idea_time,
102          'baseline_parameters':sum(p.numel() for p in base.parameters()),
103          'idea_parameters':sum(p.numel() for p in idea.parameters()),
104          'sanity_jensen_violation':sanity[0], 'sanity_monotonicity_violation':sanity[1],
105          'sanity_lower_not_applicable':sanity[2],
106          'baseline_random_jensen_violation':base_v[0], 'idea_random_jensen_violation':idea_v[0],
107          'baseline_random_monotonicity_violation':base_v[1], 'idea_random_monotonicity_violation':idea_v[1],
108          'baseline_random_lower_violation':base_v[2], 'idea_random_lower_violation':idea_v[2],
109        }
110    # Verify SVD wrapper shape and permutation invariance for a matrix energy.
111    F=torch.randn(16,2,2,device=device); e,s=energy_from_F(idea,F)
112    metrics['svd_energy_shape_ok']=bool(e.shape==(16,) and s.shape==(16,2))
113    metrics['device']=str(device)
114    os.makedirs('results',exist_ok=True)
115    with open('results/metrics.json','w') as f: json.dump(metrics,f,indent=2)
116    print(json.dumps(metrics,indent=2))
117
118if __name__=='__main__':
119    try: main()
120    except Exception as exc:
121        if str(device)=='cuda':
122            print('CUDA failed, rerun on CPU:',repr(exc)); device=torch.device('cpu'); main()
123        else: raise