1import json, math, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 7
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device = "cuda" if torch.cuda.is_available() else "cpu"
 10    if device == "cuda":
 11        torch.cuda.manual_seed_all(SEED)
 12except Exception:
 13    device = "cpu"
 14
 15
 16def tangent_project(x, q, lo=-1.0, hi=1.0):
 17    """Projection of q onto the tangent cone of [lo,hi]^d at x."""
 18    p = q.clone()
 19    p = torch.where(x <= lo, torch.clamp(p, min=0), p)
 20    p = torch.where(x >= hi, torch.clamp(p, max=0), p)
 21    return p
 22
 23
 24def numerical_math_check():
 25    # Same input-dependent forcing cancels in the difference.  The drift is -Mx,
 26    # M positive semidefinite, and the projected Euler flow stays in the box.
 27    rng = np.random.default_rng(SEED)
 28    d, trials, steps = 5, 1000, 2000
 29    R = rng.normal(size=(d, d)); M = R.T @ R / d + .15*np.eye(d)
 30    ratios = []; max_box = 0.; violations = 0
 31    h = 0.01
 32    for _ in range(trials):
 33        x = rng.uniform(-1, 1, d); z = rng.uniform(-1, 1, d)
 34        # common bounded forcing represents identical v(t)
 35        force = rng.normal(size=d) * .2
 36        initial = np.linalg.norm(x-z)
 37        for _ in range(steps):
 38            qx = -M @ x + force; qz = -M @ z + force
 39            px = qx.copy(); pz = qz.copy()
 40            px[x <= -1] = np.maximum(px[x <= -1], 0); px[x >= 1] = np.minimum(px[x >= 1], 0)
 41            pz[z <= -1] = np.maximum(pz[z <= -1], 0); pz[z >= 1] = np.minimum(pz[z >= 1], 0)
 42            x = np.clip(x + h*px, -1, 1); z = np.clip(z + h*pz, -1, 1)
 43            max_box = max(max_box, float(np.max(np.abs(x))), float(np.max(np.abs(z))))
 44        final = np.linalg.norm(x-z)
 45        ratios.append(final / (initial + 1e-12))
 46        violations += int(final > initial + 2e-7)
 47    ratios = np.asarray(ratios)
 48    return {"trials": trials, "max_distance_ratio": float(ratios.max()),
 49            "mean_distance_ratio": float(ratios.mean()), "violations": violations,
 50            "max_abs_state": max_box, "step": h, "duration": steps*h}
 51
 52
 53class VanillaResidual(nn.Module):
 54    def __init__(self, inp, hid):
 55        super().__init__(); self.hid = hid
 56        self.wx = nn.Linear(hid, hid, bias=False); self.wv = nn.Linear(inp, hid)
 57        nn.init.normal_(self.wx.weight, std=.35/math.sqrt(hid))
 58    def forward(self, seq):
 59        x = torch.zeros(seq.size(0), self.hid, device=seq.device)
 60        for t in range(seq.size(1)):
 61            x = x + .08*torch.tanh(self.wx(x) + self.wv(seq[:,t]))
 62        return x
 63
 64class ProjectedDissipative(nn.Module):
 65    def __init__(self, inp, hid):
 66        super().__init__(); self.hid=hid; self.h=.08
 67        self.L = nn.Parameter(torch.randn(hid, hid)*(.18/math.sqrt(hid)))
 68        self.B = nn.Linear(inp, hid); self.bias=nn.Parameter(torch.zeros(hid))
 69    def forward(self, seq):
 70        x=torch.zeros(seq.size(0), self.hid, device=seq.device)
 71        # positive lambda makes the drift strictly dissipative
 72        M=self.L.T@self.L + .12*torch.eye(self.hid, device=seq.device)
 73        for t in range(seq.size(1)):
 74            q=-(x@M.T) + self.B(seq[:,t]) + self.bias
 75            p=tangent_project(x,q)
 76            x=torch.clamp(x+self.h*p,-1,1)
 77        return x
 78
 79
 80def train_model(model, train, target, test, ytest, steps=300):
 81    model.to(device); opt=torch.optim.Adam(model.parameters(), lr=3e-3); lossfn=nn.BCEWithLogitsLoss()
 82    head=nn.Linear(model.hid,1).to(device); opt=torch.optim.Adam(list(model.parameters())+list(head.parameters()),lr=3e-3)
 83    xtr=train.to(device); y=target.to(device)
 84    losses=[]
 85    for k in range(steps):
 86        opt.zero_grad(); pred=head(model(xtr)).squeeze(-1); loss=lossfn(pred,y); loss.backward(); opt.step()
 87        if k in (0, steps-1): losses.append(float(loss.detach().cpu()))
 88    with torch.no_grad():
 89        acc=((torch.sigmoid(head(model(test.to(device))).squeeze(-1))>.5)==ytest.to(device)).float().mean().item()
 90        # response to a small initial-state perturbation, using identical sequence
 91        def run_with_initial(s, initial):
 92            x=initial
 93            if isinstance(model, VanillaResidual):
 94                for t in range(s.size(1)): x=x+.08*torch.tanh(model.wx(x)+model.wv(s[:,t]))
 95            else:
 96                M=model.L.T@model.L+.12*torch.eye(model.hid,device=device)
 97                for t in range(s.size(1)):
 98                    q=-(x@M.T)+model.B(s[:,t])+model.bias; x=torch.clamp(x+model.h*tangent_project(x,q),-1,1)
 99            return x
100        s=test[:32].to(device); a=torch.zeros(32,model.hid,device=device); b=a+1e-3*torch.randn_like(a)
101        ratio=(run_with_initial(s,b)-run_with_initial(s,a)).norm(dim=1).div((b-a).norm(dim=1)).mean().item()
102        final=model(test[:128].to(device)); state_norm=final.norm(dim=1).mean().item(); max_abs=final.abs().max().item()
103    return {"loss_start":losses[0],"loss_final":losses[-1],"accuracy":acc,"perturbation_ratio":ratio,"mean_final_norm":state_norm,"max_abs_final":max_abs,"parameters":sum(p.numel() for p in model.parameters())+sum(p.numel() for p in head.parameters())}
104
105
106def main():
107    math_result=numerical_math_check()
108    # Long-horizon sign-of-sum task; one scalar input per time step.
109    ntr, nte, T = 256, 128, 200
110    train=torch.randn(ntr,T,1); test=torch.randn(nte,T,1)
111    y=(train.sum(1).squeeze(1)>0).float(); yt=(test.sum(1).squeeze(1)>0).float()
112    baseline=train_model(VanillaResidual(1,16),train,y,test,yt)
113    idea=train_model(ProjectedDissipative(1,16),train,y,test,yt)
114    out={"device":device,"math_check":math_result,"baseline":baseline,"idea":idea}
115    print(json.dumps(out, indent=2))
116    with open("results.json","w") as f: json.dump(out,f,indent=2)
117
118if __name__ == "__main__":
119    try: main()
120    except Exception as e:
121        if device == "cuda":
122            device="cpu"; print("CUDA failed, rerun with CPU", repr(e)); main()
123        else: raise