import json, math, time, random import numpy as np import torch import torch.nn as nn SEED=917 np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED) torch.set_num_threads(4) device = "cuda" if torch.cuda.is_available() else "cpu" # Unit-square homogeneous problem: -Delta u=f, u=0 on boundary. def exact_u(x): return torch.sin(math.pi*x[:,0])*torch.sin(math.pi*x[:,1]) def forcing(x): return 2*math.pi**2*exact_u(x) def bubble(x): return x[:,0]*(1-x[:,0])*x[:,1]*(1-x[:,1]) # Smooth H^2 tests, vanishing on boundary; Laplacian is obtained analytically by autograd once. def make_tests(x, modes=12): vals=[]; laps=[] xx=x.detach().clone().requires_grad_(True) for k in range(1, modes+1): l=(k % 3)+1 w=bubble(xx)*torch.sin(k*math.pi*xx[:,0])*torch.sin(l*math.pi*xx[:,1]) g=torch.autograd.grad(w.sum(),xx,create_graph=True,retain_graph=True)[0] lap=0 for d in range(2): lap=lap+torch.autograd.grad(g[:,d].sum(),xx,create_graph=False,retain_graph=(d==0))[0][:,d] vals.append(w.detach()); laps.append(lap.detach()) return torch.stack(vals),torch.stack(laps) def vwf_residual(u, lapw, w, f, area_weights): # g=0, exactly the residual specified in the idea return -torch.sum(area_weights[None,:]*u[None,:]*lapw,1) - torch.sum(area_weights[None,:]*f[None,:]*w,1) class FloatNet(nn.Module): def __init__(self, hidden=16): super().__init__(); self.net=nn.Sequential(nn.Linear(2,hidden),nn.Tanh(),nn.Linear(hidden,hidden),nn.Tanh(),nn.Linear(hidden,1)) def forward(self,x): return bubble(x)*self.net(x).squeeze(-1) class BinaryNet(nn.Module): def __init__(self, hidden=16, tau=1.0): super().__init__(); self.tau=tau self.layers=nn.ModuleList([nn.Linear(2,hidden),nn.Linear(hidden,hidden),nn.Linear(hidden,1)]) def bweight(self,w): # STE sign, with learned per-layer positive scale. s=w.abs().mean().detach().clamp_min(1e-3) q=s*torch.sign(w) return w+(q-w).detach() def forward(self,x,hard=False): z=x for i,lay in enumerate(self.layers): z=nn.functional.linear(z,self.bweight(lay.weight),self.bweight(lay.bias)) if i<2: z=(z>=0).float() if hard else torch.sigmoid(z/self.tau) return bubble(x)*z.squeeze(-1) def quant_penalty(self): p=0. for lay in self.layers: s=lay.weight.abs().mean().detach().clamp_min(1e-3) p=p+((lay.weight-s*torch.sign(lay.weight))**2).mean() sb=lay.bias.abs().mean().detach().clamp_min(1e-3) p=p+((lay.bias-sb*torch.sign(lay.bias))**2).mean() return p def loss_for(model,x,w,lapw,f,weights,hard=False): u=model(x,hard=hard) if isinstance(model,BinaryNet) else model(x) r=vwf_residual(u,lapw,w,f,weights) return (r*r).mean(),r,u def sanity(x,w,lapw,weights): # Quantitative VWF identity check for exact smooth solution: residual should converge with quadrature. rows=[] for n in [10,20,40,80]: xx=(torch.arange(n,dtype=torch.float32)+.5)/n X,Y=torch.meshgrid(xx,xx,indexing='ij'); z=torch.stack([X.flatten(),Y.flatten()],1) ww,ll=make_tests(z,3); ff=forcing(z); uu=exact_u(z); aw=torch.full((n*n,),1/(n*n)) rr=vwf_residual(uu,ll,ww,ff,aw) rows.append(float(rr.abs().max())) # Boundary prediction: bubble is identically zero on all four edges, including hard model. edge=torch.tensor([[0.,.2],[1.,.7],[.3,0.],[.8,1.]]) bn=BinaryNet(); bv=bn(edge,hard=True).abs().max().item() rates=[rows[i]/rows[i+1] for i in range(len(rows)-1)] 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'} def sweep_hardening(x,w,lapw,f,weights): # Fixed binary parameters: prediction is tau->0 sigmoid -> Heaviside except zero margins. 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) with torch.no_grad(): hard=model(x,hard=True); rh=vwf_residual(hard,lapw,w,f,weights) # Hidden-logit margins determine the exact sigmoid-vs-step mismatch. z=x; hard_states=[]; margins=[] for i,lay in enumerate(model.layers): z=nn.functional.linear(z,model.bweight(lay.weight),model.bweight(lay.bias)) if i<2: margins.append(z.abs().flatten()); hard_states.append((z>=0)) z=torch.sigmoid(z) margin=torch.cat(margins) out=[] for tau in [1.0,.5,.2,.1,.05,.02,.01,.005]: model.tau=tau; soft=model(x,hard=False); rs=vwf_residual(soft,lapw,w,f,weights) z=x; mismatch=[] for i,lay in enumerate(model.layers): z=nn.functional.linear(z,model.bweight(lay.weight),model.bweight(lay.bias)) if i<2: mismatch.append(((torch.sigmoid(z/tau)>=.5) != hard_states[i]).float().mean()) z=torch.sigmoid(z/tau) 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())}) 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'} def train(kind,x,w,lapw,f,weights,steps=700): model=FloatNet().to(device) if kind=='float' else BinaryNet().to(device) opt=torch.optim.Adam(model.parameters(),lr=kind=='float' and 2e-3 or 1e-3) t0=time.perf_counter(); hist=[] for i in range(steps): opt.zero_grad(); l,r,u=loss_for(model,x,w,lapw,f,weights) if kind!='float': l=l+0.01*model.quant_penalty() l.backward(); opt.step() if i in [0,steps//2,steps-1]: hist.append(float(l.detach())) elapsed=time.perf_counter()-t0 with torch.no_grad(): lsoft,rsoft,u=loss_for(model,x,w,lapw,f,weights) if kind=='binary': lhard,rhard,uh=loss_for(model,x,w,lapw,f,weights,hard=True) else: lhard,rhard,uh=lsoft,rsoft,u grid=torch.rand(5000,2,device=device); pred=model(grid,hard=True) if kind=='binary' else model(grid) rel=float(torch.linalg.vector_norm(pred-exact_u(grid))/torch.linalg.vector_norm(exact_u(grid))) # Warmed-up deployment-style forward timing (includes ordinary PyTorch linear kernels). for _ in range(10): _=model(grid,hard=True) if kind=='binary' else model(grid) if device=='cuda': torch.cuda.synchronize() t_inf=time.perf_counter() for _ in range(50): _=model(grid,hard=True) if kind=='binary' else model(grid) if device=='cuda': torch.cuda.synchronize() inference_ms=1000*(time.perf_counter()-t_inf)/50 params=sum(p.numel() for p in model.parameters()) # Float32 baseline versus one-bit weights plus negligible float scales/bias metadata. float_bytes=params*4 if kind=='binary': binary_bytes=sum(l.weight.numel()+l.bias.numel() for l in model.layers)/8 + len(model.layers)*8 else: binary_bytes=float_bytes 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} def main(): global device 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) w,lapw=make_tests(x,12); f=forcing(x) # CPU/GPU failure fallback as required. try: san=sanity(x.cpu(),*make_tests(x.cpu(),12)[0:2],torch.full((n*n,),1/(n*n))) sweep=sweep_hardening(x,w,lapw,f,weights) results={'device':device,'sanity':san,'hardening':sweep,'float':train('float',x,w,lapw,f,weights),'binary':train('binary',x,w,lapw,f,weights)} except Exception as e: if device!='cpu': torch.cuda.empty_cache(); device='cpu'; return main_cpu() raise with open('results.json','w') as fp: json.dump(results,fp,indent=2) print(json.dumps(results,indent=2)) def main_cpu(): 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)) if __name__=='__main__': main()