1import json, math, time, random
  2import numpy as np
  3import torch
  4import torch.nn as nn
  5
  6SEED=917
  7np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9device = "cuda" if torch.cuda.is_available() else "cpu"
 10
 11# Unit-square homogeneous problem: -Delta u=f, u=0 on boundary.
 12def exact_u(x): return torch.sin(math.pi*x[:,0])*torch.sin(math.pi*x[:,1])
 13def forcing(x): return 2*math.pi**2*exact_u(x)
 14def bubble(x): return x[:,0]*(1-x[:,0])*x[:,1]*(1-x[:,1])
 15
 16# Smooth H^2 tests, vanishing on boundary; Laplacian is obtained analytically by autograd once.
 17def make_tests(x, modes=12):
 18    vals=[]; laps=[]
 19    xx=x.detach().clone().requires_grad_(True)
 20    for k in range(1, modes+1):
 21        l=(k % 3)+1
 22        w=bubble(xx)*torch.sin(k*math.pi*xx[:,0])*torch.sin(l*math.pi*xx[:,1])
 23        g=torch.autograd.grad(w.sum(),xx,create_graph=True,retain_graph=True)[0]
 24        lap=0
 25        for d in range(2):
 26            lap=lap+torch.autograd.grad(g[:,d].sum(),xx,create_graph=False,retain_graph=(d==0))[0][:,d]
 27        vals.append(w.detach()); laps.append(lap.detach())
 28    return torch.stack(vals),torch.stack(laps)
 29
 30def vwf_residual(u, lapw, w, f, area_weights):
 31    # g=0, exactly the residual specified in the idea
 32    return -torch.sum(area_weights[None,:]*u[None,:]*lapw,1) - torch.sum(area_weights[None,:]*f[None,:]*w,1)
 33
 34class FloatNet(nn.Module):
 35    def __init__(self, hidden=16):
 36        super().__init__(); self.net=nn.Sequential(nn.Linear(2,hidden),nn.Tanh(),nn.Linear(hidden,hidden),nn.Tanh(),nn.Linear(hidden,1))
 37    def forward(self,x): return bubble(x)*self.net(x).squeeze(-1)
 38
 39class BinaryNet(nn.Module):
 40    def __init__(self, hidden=16, tau=1.0):
 41        super().__init__(); self.tau=tau
 42        self.layers=nn.ModuleList([nn.Linear(2,hidden),nn.Linear(hidden,hidden),nn.Linear(hidden,1)])
 43    def bweight(self,w):
 44        # STE sign, with learned per-layer positive scale.
 45        s=w.abs().mean().detach().clamp_min(1e-3)
 46        q=s*torch.sign(w)
 47        return w+(q-w).detach()
 48    def forward(self,x,hard=False):
 49        z=x
 50        for i,lay in enumerate(self.layers):
 51            z=nn.functional.linear(z,self.bweight(lay.weight),self.bweight(lay.bias))
 52            if i<2: z=(z>=0).float() if hard else torch.sigmoid(z/self.tau)
 53        return bubble(x)*z.squeeze(-1)
 54    def quant_penalty(self):
 55        p=0.
 56        for lay in self.layers:
 57            s=lay.weight.abs().mean().detach().clamp_min(1e-3)
 58            p=p+((lay.weight-s*torch.sign(lay.weight))**2).mean()
 59            sb=lay.bias.abs().mean().detach().clamp_min(1e-3)
 60            p=p+((lay.bias-sb*torch.sign(lay.bias))**2).mean()
 61        return p
 62
 63def loss_for(model,x,w,lapw,f,weights,hard=False):
 64    u=model(x,hard=hard) if isinstance(model,BinaryNet) else model(x)
 65    r=vwf_residual(u,lapw,w,f,weights)
 66    return (r*r).mean(),r,u
 67
 68def sanity(x,w,lapw,weights):
 69    # Quantitative VWF identity check for exact smooth solution: residual should converge with quadrature.
 70    rows=[]
 71    for n in [10,20,40,80]:
 72        xx=(torch.arange(n,dtype=torch.float32)+.5)/n
 73        X,Y=torch.meshgrid(xx,xx,indexing='ij'); z=torch.stack([X.flatten(),Y.flatten()],1)
 74        ww,ll=make_tests(z,3); ff=forcing(z); uu=exact_u(z); aw=torch.full((n*n,),1/(n*n))
 75        rr=vwf_residual(uu,ll,ww,ff,aw)
 76        rows.append(float(rr.abs().max()))
 77    # Boundary prediction: bubble is identically zero on all four edges, including hard model.
 78    edge=torch.tensor([[0.,.2],[1.,.7],[.3,0.],[.8,1.]])
 79    bn=BinaryNet(); bv=bn(edge,hard=True).abs().max().item()
 80    rates=[rows[i]/rows[i+1] for i in range(len(rows)-1)]
 81    return {'identity_max_abs_residual_n10_20_40_80':rows,'observed_refinement_ratios':rates,'prediction':'midpoint quadrature for smooth integrands has O(h^2) error, ratio 4 under grid doubling','boundary_max_abs_hard_output':bv,'boundary_prediction':'bubble factor enforces exactly zero on all edges'}
 82
 83def sweep_hardening(x,w,lapw,f,weights):
 84    # Fixed binary parameters: prediction is tau->0 sigmoid -> Heaviside except zero margins.
 85    model=BinaryNet(hidden=16,tau=1.0).to(device); x=x.to(device); w=w.to(device); lapw=lapw.to(device); f=f.to(device); weights=weights.to(device)
 86    with torch.no_grad():
 87        hard=model(x,hard=True); rh=vwf_residual(hard,lapw,w,f,weights)
 88        # Hidden-logit margins determine the exact sigmoid-vs-step mismatch.
 89        z=x; hard_states=[]; margins=[]
 90        for i,lay in enumerate(model.layers):
 91            z=nn.functional.linear(z,model.bweight(lay.weight),model.bweight(lay.bias))
 92            if i<2:
 93                margins.append(z.abs().flatten()); hard_states.append((z>=0))
 94                z=torch.sigmoid(z)
 95        margin=torch.cat(margins)
 96        out=[]
 97        for tau in [1.0,.5,.2,.1,.05,.02,.01,.005]:
 98            model.tau=tau; soft=model(x,hard=False); rs=vwf_residual(soft,lapw,w,f,weights)
 99            z=x; mismatch=[]
100            for i,lay in enumerate(model.layers):
101                z=nn.functional.linear(z,model.bweight(lay.weight),model.bweight(lay.bias))
102                if i<2:
103                    mismatch.append(((torch.sigmoid(z/tau)>=.5) != hard_states[i]).float().mean())
104                    z=torch.sigmoid(z/tau)
105            out.append({'tau':tau,'mean_output_gap':float((soft-hard).abs().mean()),'residual_l2_gap':float(torch.sqrt(((rs-rh)**2).mean())),'activation_mismatch':float(torch.stack(mismatch).mean())})
106    return {'sweep':out,'median_logit_margin':float(margin.median()),'prediction':'for nonzero margins, sigmoid(z/tau)>=0.5 equals H(z), so mismatch is predicted 0 for every tau'}
107
108def train(kind,x,w,lapw,f,weights,steps=700):
109    model=FloatNet().to(device) if kind=='float' else BinaryNet().to(device)
110    opt=torch.optim.Adam(model.parameters(),lr=kind=='float' and 2e-3 or 1e-3)
111    t0=time.perf_counter(); hist=[]
112    for i in range(steps):
113        opt.zero_grad(); l,r,u=loss_for(model,x,w,lapw,f,weights)
114        if kind!='float': l=l+0.01*model.quant_penalty()
115        l.backward(); opt.step()
116        if i in [0,steps//2,steps-1]: hist.append(float(l.detach()))
117    elapsed=time.perf_counter()-t0
118    with torch.no_grad():
119        lsoft,rsoft,u=loss_for(model,x,w,lapw,f,weights)
120        if kind=='binary': lhard,rhard,uh=loss_for(model,x,w,lapw,f,weights,hard=True)
121        else: lhard,rhard,uh=lsoft,rsoft,u
122        grid=torch.rand(5000,2,device=device); pred=model(grid,hard=True) if kind=='binary' else model(grid)
123        rel=float(torch.linalg.vector_norm(pred-exact_u(grid))/torch.linalg.vector_norm(exact_u(grid)))
124        # Warmed-up deployment-style forward timing (includes ordinary PyTorch linear kernels).
125        for _ in range(10):
126            _=model(grid,hard=True) if kind=='binary' else model(grid)
127        if device=='cuda': torch.cuda.synchronize()
128        t_inf=time.perf_counter()
129        for _ in range(50): _=model(grid,hard=True) if kind=='binary' else model(grid)
130        if device=='cuda': torch.cuda.synchronize()
131        inference_ms=1000*(time.perf_counter()-t_inf)/50
132    params=sum(p.numel() for p in model.parameters())
133    # Float32 baseline versus one-bit weights plus negligible float scales/bias metadata.
134    float_bytes=params*4
135    if kind=='binary': binary_bytes=sum(l.weight.numel()+l.bias.numel() for l in model.layers)/8 + len(model.layers)*8
136    else: binary_bytes=float_bytes
137    return {'loss_soft':float(lsoft),'weak_residual_rms_soft':float(torch.sqrt((rsoft*rsoft).mean())),'hard_loss':float(lhard),'hard_residual_rms':float(torch.sqrt((rhard*rhard).mean())),'relative_L2':rel,'params':params,'float32_bytes':float_bytes,'binary_estimated_bytes':binary_bytes,'storage_ratio':float_bytes/binary_bytes,'train_seconds':elapsed,'inference_ms_batch5000':inference_ms,'loss_trace':hist}
138
139def main():
140    global device
141    n=28; q=(torch.arange(n,dtype=torch.float32)+.5)/n; X,Y=torch.meshgrid(q,q,indexing='ij'); x=torch.stack([X.flatten(),Y.flatten()],1).to(device); weights=torch.full((n*n,),1/(n*n),device=device)
142    w,lapw=make_tests(x,12); f=forcing(x)
143    # CPU/GPU failure fallback as required.
144    try:
145        san=sanity(x.cpu(),*make_tests(x.cpu(),12)[0:2],torch.full((n*n,),1/(n*n)))
146        sweep=sweep_hardening(x,w,lapw,f,weights)
147        results={'device':device,'sanity':san,'hardening':sweep,'float':train('float',x,w,lapw,f,weights),'binary':train('binary',x,w,lapw,f,weights)}
148    except Exception as e:
149        if device!='cpu':
150            torch.cuda.empty_cache(); device='cpu'; return main_cpu()
151        raise
152    with open('results.json','w') as fp: json.dump(results,fp,indent=2)
153    print(json.dumps(results,indent=2))
154
155def main_cpu():
156    global device; device='cpu'; n=28; q=(torch.arange(n)+.5)/n; X,Y=torch.meshgrid(q,q,indexing='ij'); x=torch.stack([X.flatten(),Y.flatten()],1); weights=torch.full((n*n,),1/(n*n)); w,lapw=make_tests(x,12); f=forcing(x); results={'device':'cpu','sanity':sanity(x,w,lapw,weights),'hardening':sweep_hardening(x,w,lapw,f,weights),'float':train('float',x,w,lapw,f,weights),'binary':train('binary',x,w,lapw,f,weights)}; open('results.json','w').write(json.dumps(results,indent=2)); print(json.dumps(results,indent=2))
157if __name__=='__main__': main()