1import sys, json, time
 2from pathlib import Path
 3import numpy as np
 4import torch
 5import torch.nn as nn
 6sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
 7from bench import get_dataset, train_model, evaluate, sweep_baseline, make_report
 8from poisson_symmetry_track import get_dataset as get_pde, META
 9
10SEEDS = tuple(range(8))
11LRS = [1e-3, 3e-3, 1e-2]
12EPOCHS = 16
13BATCH = 64
14N = 8
15D = N*N
16
17def laplacian(n=N):
18    A = torch.zeros(D, D)
19    for i in range(n):
20        for j in range(n):
21            q=i*n+j; A[q,q]=4
22            for ii,jj in [((i+1)%n,j),((i-1)%n,j),(i,(j+1)%n),(i,(j-1)%n)]:
23                A[q,ii*n+jj] -= 1
24    return A
25
26L = laplacian()
27K = L + 1e-5*torch.eye(D)
28LAM = (4 - 2*torch.cos(2*torch.pi*torch.arange(N)/N)[:,None]
29       - 2*torch.cos(2*torch.pi*torch.arange(N)/N)[None,:]).float()
30
31class PoissonNet(nn.Module):
32    def __init__(self, idea=False):
33        super().__init__(); self.idea=idea
34        self.trunk=nn.Sequential(nn.Linear(D,64),nn.Tanh(),nn.Linear(64,64),nn.Tanh(),nn.Linear(64,D))
35    def solve(self, b):
36        if not self.idea:
37            return torch.linalg.solve(K.to(device=b.device, dtype=b.dtype), b.unsqueeze(-1)).squeeze(-1)
38        bh=torch.fft.fft2(b.reshape(-1,N,N), dim=(-2,-1))
39        uh=bh/(LAM.to(device=b.device, dtype=b.dtype)+1e-5); uh[...,0,0]=0
40        return torch.fft.ifft2(uh, dim=(-2,-1)).real.reshape(-1,D)
41    def forward(self,x): return self.solve(self.trunk(x))
42
43def ds(seed):
44    a=get_pde(seed,400,200)
45    return {k: torch.as_tensor(v, dtype=torch.float32) if k in ('xtr','ytr','xte','yte') else v for k,v in a.items()}
46
47def factory(idea, cfg):
48    def fn(seed):
49        torch.manual_seed(seed); np.random.seed(seed)
50        net=PoissonNet(idea)
51        _, metric, _=train_model(net, ds(seed), epochs=EPOCHS, lr=cfg['lr'], batch=BATCH, log=lambda *_:None)
52        return float(metric)
53    return fn
54
55def main():
56    # All idea learning rates are included in the baseline sweep (search-space parity).
57    base=sweep_baseline(lambda cfg: factory(False, cfg), [{'lr':lr} for lr in LRS], seeds=SEEDS)
58    runs=[]
59    for lr in LRS:
60        vals=[factory(True, {'lr':lr})(s) for s in SEEDS]
61        runs.append({'lr':lr,'mean':float(np.mean(vals)),'std':float(np.std(vals)),'per_seed':vals,'n':len(vals)})
62    chosen=min(runs,key=lambda z:z['mean'])
63    idea=dict(chosen); idea['sweep']=runs; idea['cfg']={'lr':chosen['lr'],'nearby_tested':LRS}
64    # Core trained-model signature: compare observed residual and equivariance of its output.
65    torch.manual_seed(0); net=PoissonNet(True)
66    trained, metric, _=train_model(net, ds(0), epochs=EPOCHS, lr=chosen['lr'], batch=BATCH, log=lambda *_:None)
67    with torch.no_grad():
68        dev=next(trained.parameters()).device; x=ds(0)['xte'][:32].to(dev); pred=trained(x)
69        Kd=K.to(device=dev, dtype=pred.dtype)
70        residual=((pred@Kd.T)-trained.trunk(x)).norm(dim=1)/(trained.trunk(x).norm(dim=1)+1e-8)
71        # Translation action is a mesh automorphism; test output equivariance on the trained net.
72        xx=x.reshape(-1,N,N); shifted=torch.roll(xx,1,dims=1).reshape(-1,D)
73        equiv=(trained(shifted)-torch.roll(pred.reshape(-1,N,N),1,dims=1).reshape(-1,D)).norm(dim=1)/(pred.norm(dim=1)+1e-8)
74        comm=float(torch.linalg.norm(L@torch.roll(torch.eye(D),1,dims=0)-torch.roll(torch.eye(D),1,dims=0)@L))
75        off=float((torch.fft.fft2(L.reshape(N,N,N,N).diagonal(dim1=1,dim2=3),dim=(-2,-1)) if False else torch.tensor(0.)).item())
76    sig={'prediction':'symmetry-basis implicit solve preserves the discretized solution and commutes with translations',
77         'trained_model_observed_max_relative_solve_residual':float(residual.max()),
78         'trained_model_observed_max_translation_equivariance_error':float(equiv.max()),
79         'operator_commutator_frobenius':comm,
80         'predicted_residual_and_equivariance':0.0,
81         'confirmed':bool(float(residual.max())<2e-4 and float(equiv.max())<2e-4 and comm<1e-6),
82         'note':'signature is measured from the trained idea model on held-out benchmark inputs'}
83    report=make_report('symmetric_periodic_poisson','matched_mlp_implicit_poisson',base,idea,sig)
84    report['custom_track']={'name':META['name'],'file':'poisson_symmetry_track.py','domain':META['domain']}
85    report['stage2_config']={'epochs':EPOCHS,'batch':BATCH,'seeds':list(SEEDS),'lr_union':LRS,'track_reason':'PDE inverse problem with periodic Poisson operator; no built-in track had this structure'}
86    Path('bench_report.json').write_text(json.dumps(report,indent=2))
87    print(json.dumps(report,indent=2))
88if __name__=='__main__': main()