1import json, math, random, time
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn as nn
  6
  7SEED = 619
  8random.seed(SEED)
  9np.random.seed(SEED)
 10torch.manual_seed(SEED)
 11torch.set_num_threads(4)
 12device = "cuda" if torch.cuda.is_available() else "cpu"
 13try:
 14    if device == "cuda":
 15        torch.cuda.reset_peak_memory_stats()
 16except Exception:
 17    device = "cpu"
 18
 19
 20def make_accretive(d, delta=0.35, skew=1.8, seed=1):
 21    g = np.random.default_rng(seed)
 22    L = g.normal(size=(d, max(2, d // 2))) / math.sqrt(d)
 23    S = L @ L.T + delta * np.eye(d)
 24    U = g.normal(size=(d, d)) / math.sqrt(d)
 25    K = skew * (U - U.T) / 2
 26    return S + K
 27
 28
 29def exact_fractional(A, r):
 30    vals, vecs = np.linalg.eig(A.astype(np.complex128))
 31    return (vecs @ np.diag(vals ** r) @ np.linalg.inv(vecs)).real
 32
 33
 34def resolvent_fractional(A, r=0.5, m=10, eps=0.02, lo=1e-3, hi=100.0):
 35    # Log-grid midpoint quadrature for the stated integral after the
 36    # cancellation-safe identity lambda^-1 (lambda I+A)^-1 A.
 37    x = np.linspace(math.log(lo), math.log(hi), m + 1)
 38    lam = np.exp((x[:-1] + x[1:]) / 2)
 39    weights = lam * (x[1:] - x[:-1])
 40    Ae = A + eps * np.eye(A.shape[0])
 41    out = np.zeros_like(Ae, dtype=float)
 42    I = np.eye(A.shape[0])
 43    for l, w in zip(lam, weights):
 44        out += w * (l ** (r - 1.0)) * np.linalg.solve(l * I + Ae, Ae)
 45    return math.sin(math.pi * r) / math.pi * out
 46
 47
 48def math_check():
 49    rows = []
 50    for skew in [0.0, 1.8, 4.0]:
 51        A = make_accretive(8, skew=skew, seed=10)
 52        S = (A + A.T) / 2
 53        min_sym = np.linalg.eigvalsh(S).min()
 54        ref = exact_fractional(A, 0.5)
 55        approx = resolvent_fractional(A, m=10, eps=0.0, lo=1e-4, hi=300)
 56        rel = np.linalg.norm(approx-ref) / np.linalg.norm(ref)
 57        # Positive shifts should remain well-conditioned under accretivity.
 58        conds = [np.linalg.cond(l*np.eye(8)+A) for l in np.geomspace(1e-3, 100, 10)]
 59        rows.append({"skew": skew, "min_sym_eigenvalue": float(min_sym),
 60                     "relative_matrix_error": float(rel),
 61                     "max_shifted_condition": float(max(conds))})
 62    # Scalar claim: same quadrature approximates t^r over the sampled range.
 63    ts = np.geomspace(0.35, 30.0, 400)
 64    xs = np.linspace(math.log(1e-4), math.log(300), 11)
 65    lam = np.exp((xs[:-1]+xs[1:])/2)
 66    weights = lam*(xs[1:]-xs[:-1])
 67    pred = np.zeros_like(ts)
 68    for l,w in zip(lam,weights): pred += w*l**(-0.5)*ts/(l+ts)
 69    pred *= 1/math.pi
 70    scalar_rel = float(np.max(np.abs(pred-np.sqrt(ts))/np.sqrt(ts)))
 71    return {"matrices": rows, "scalar_max_relative_error": scalar_rel}
 72
 73
 74class ResolventBlock(nn.Module):
 75    def __init__(self, d, r=0.5, m=6, eps=0.03, lo=0.08, hi=8.0):
 76        super().__init__()
 77        self.d, self.r, self.m = d, r, m
 78        self.eps, self.lo, self.hi = eps, lo, hi
 79        self.L = nn.Parameter(torch.randn(d, d//2) / math.sqrt(d))
 80        self.U = nn.Parameter(torch.randn(d, d) / math.sqrt(d))
 81        self.gain = nn.Parameter(torch.tensor(0.15))
 82        self.register_buffer("eye", torch.eye(d))
 83
 84    def operator(self):
 85        S = self.L @ self.L.T + 0.20*self.eye
 86        K = (self.U-self.U.T)/2
 87        return S + K + self.eps*self.eye
 88
 89    def forward(self, x):
 90        A = self.operator()
 91        xx = x.T
 92        grid = torch.linspace(math.log(self.lo), math.log(self.hi), self.m+1, device=x.device, dtype=x.dtype)
 93        lam = torch.exp((grid[:-1]+grid[1:])/2)
 94        weights = lam*(grid[1:]-grid[:-1])
 95        z = torch.zeros_like(xx)
 96        for l,w in zip(lam, weights):
 97            rhs = A @ xx
 98            z = z + w*l**(self.r-1)*torch.linalg.solve(l*self.eye.to(x)+A, rhs)
 99        z = math.sin(math.pi*self.r)/math.pi*z
100        return x + self.gain*z.T
101
102
103class LinearBlock(nn.Module):
104    def __init__(self, d):
105        super().__init__()
106        self.W = nn.Parameter(torch.randn(d,d)/math.sqrt(d))
107        self.gain = nn.Parameter(torch.tensor(0.15))
108    def forward(self,x):
109        return x + self.gain*(x @ self.W.T)
110
111
112class Net(nn.Module):
113    def __init__(self, block, d=12, classes=3):
114        super().__init__()
115        self.inp=nn.Linear(2,d)
116        self.block=block
117        self.out=nn.Linear(d,classes)
118    def forward(self,x): return self.out(torch.tanh(self.block(torch.tanh(self.inp(x)))))
119
120
121def training_check():
122    g=torch.Generator().manual_seed(SEED)
123    n=360
124    centers=torch.tensor([[1.5,0.0],[-0.75,1.3],[-0.75,-1.3]])
125    y=torch.randint(0,3,(n,),generator=g)
126    x=centers[y]+0.42*torch.randn(n,2,generator=g)
127    perm=torch.randperm(n,generator=g); tr,va=perm[:270],perm[270:]
128    results={}
129    for name, factory in [("baseline",lambda:LinearBlock(12)),("resolvent",lambda:ResolventBlock(12))]:
130        torch.manual_seed(SEED+ (0 if name=='baseline' else 1))
131        net=Net(factory()).to(device)
132        opt=torch.optim.Adam(net.parameters(),lr=0.025)
133        losses=[]; gradvars=[]; t0=time.time()
134        for step in range(180):
135            opt.zero_grad(set_to_none=True)
136            logits=net(x[tr].to(device)); loss=nn.functional.cross_entropy(logits,y[tr].to(device))
137            loss.backward()
138            gn=torch.sqrt(sum((p.grad.detach()**2).sum() for p in net.parameters() if p.grad is not None))
139            gradvars.append(float(gn)); opt.step(); losses.append(float(loss))
140        with torch.no_grad():
141            pred=net(x[va].to(device)).argmax(1).cpu(); acc=float((pred==y[va]).float().mean())
142        results[name]={"final_train_loss":losses[-1],"loss_at_60":losses[59],"validation_accuracy":acc,
143                       "gradient_norm_variance":float(np.var(gradvars)),"seconds":time.time()-t0,
144                       "nan_or_inf":bool(not np.isfinite(losses).all())}
145    return results
146
147
148if __name__ == "__main__":
149    out={"device":device,"math":math_check(),"training":training_check()}
150    Path("results.json").write_text(json.dumps(out,indent=2))
151    print(json.dumps(out,indent=2))