1import json, math, os, random
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 17
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8try:
  9    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 10except Exception:
 11    device = torch.device("cpu")
 12
 13
 14def harmonic_cycle(n=24):
 15    # D is the vertex-edge incidence matrix; ker(D) is the 1-cycle space.
 16    D = np.zeros((n, n), dtype=float)
 17    for e in range(n):
 18        D[e, e] = -1.0
 19        D[(e + 1) % n, e] = 1.0
 20    # Positive diagonal mass matrices, deliberately nonuniform.
 21    m = 0.5 + np.linspace(0.2, 1.4, n)
 22    M = np.diag(m)
 23    # There are no 2-simplices, hence the closed constraint is vacuous.
 24    # Coclosedness is D^T M h = 0. Its one-dimensional solution is
 25    # h proportional to M^{-1} times the constant circulation vector.
 26    h = np.linalg.solve(M, np.ones(n))
 27    h /= math.sqrt(h @ M @ h)
 28    H = h[:, None]
 29    return D, M, H
 30
 31
 32def math_check():
 33    D, M, H = harmonic_cycle()
 34    # This cycle has no 2-simplices, so B_2 is the empty (0 x n) matrix.
 35    # Thus B_2 H = 0 exactly; D is B_1 and tests the codifferential.
 36    B2 = np.zeros((0, len(H)))
 37    closed = np.linalg.norm(B2 @ H)
 38    coclosed = np.linalg.norm(D.T @ M @ H)
 39    ortho = float((H.T @ M @ H).item())
 40    rng = np.random.default_rng(SEED)
 41    x = rng.normal(size=(len(H), 50))
 42    c = H.T @ M @ x
 43    u = x - H @ c
 44    projection_error = np.max(np.abs(H.T @ M @ u))
 45    reconstruction = np.max(np.abs(x - (u + H @ c)))
 46    # Period is chosen as the linear circulation functional whose value on H is 1.
 47    P = (H.T @ M)
 48    period_consistency = np.max(np.abs(P @ H - 1.0))
 49    return {
 50        "closed_residual": float(closed), "coclosed_residual": float(coclosed),
 51        "weighted_orthonormality": ortho, "max_gauge_residual": float(projection_error),
 52        "max_reconstruction_residual": float(reconstruction),
 53        "period_basis_residual": float(period_consistency)
 54    }, D, M, H
 55
 56
 57class LocalGNN(nn.Module):
 58    def __init__(self, n, hidden=32, harmonic=False, H=None, M=None):
 59        super().__init__(); self.n=n; self.harmonic=harmonic
 60        self.inp=nn.Linear(1, hidden)
 61        self.layers=nn.ModuleList([nn.Linear(hidden, hidden) for _ in range(4)])
 62        self.local=nn.Linear(hidden*2, 1)
 63        if harmonic:
 64            self.coeff=nn.Sequential(nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 1))
 65            self.register_buffer("H", torch.tensor(H, dtype=torch.float32))
 66            self.register_buffer("M", torch.tensor(M, dtype=torch.float32))
 67    def forward(self, x):
 68        # x: batch,n node scalars. Ring message passing, no global operation.
 69        z=torch.tanh(self.inp(x.unsqueeze(-1)))
 70        for layer in self.layers:
 71            msg=(torch.roll(z,1,1)+torch.roll(z,-1,1))/2
 72            z=torch.tanh(layer(z+msg))
 73        edge_z=(z+torch.roll(z,-1,1))/2
 74        raw=self.local(torch.cat([z, edge_z], dim=-1)).squeeze(-1)
 75        if not self.harmonic:
 76            return raw
 77        # Global latent channel: pooled embedding predicts circulation coefficient.
 78        a=self.coeff(z.mean(dim=1)).squeeze(-1)
 79        # Gauge-fix local output using H^T M, then reconstruct y=u+Ha.
 80        c=raw @ (torch.diagonal(self.M) * self.H.squeeze(1))
 81        u=raw-c[:,None]*self.H.squeeze(1)
 82        return u+a[:,None]*self.H.squeeze(1)
 83
 84
 85def run_training(H, M, n=24, steps=700):
 86    # A source scalar is placed at one node. The target has global circulation
 87    # a on every edge; targets are randomized independently.
 88    rng=np.random.default_rng(SEED)
 89    Ntr,Nte=256,256
 90    xtr=np.zeros((Ntr,n),np.float32); xte=np.zeros((Nte,n),np.float32)
 91    atr=rng.uniform(-2,2,Ntr).astype(np.float32); ate=rng.uniform(-2,2,Nte).astype(np.float32)
 92    xtr[:,0]=atr; xte[:,0]=ate
 93    ytr=(atr[:,None]*H[:,0][None,:]).astype(np.float32); yte=(ate[:,None]*H[:,0][None,:]).astype(np.float32)
 94    xtr=torch.tensor(xtr,device=device); xte=torch.tensor(xte,device=device)
 95    ytr=torch.tensor(ytr,device=device); yte=torch.tensor(yte,device=device)
 96    out={}
 97    for name,is_h in [("baseline",False),("harmonic",True)]:
 98        torch.manual_seed(SEED)
 99        model=LocalGNN(n,hidden=24,harmonic=is_h,H=H,M=M).to(device)
100        opt=torch.optim.Adam(model.parameters(),lr=3e-3)
101        for step in range(steps):
102            # minibatches keep memory modest
103            ix=torch.randint(0,Ntr,(64,),device=device)
104            pred=model(xtr[ix]); loss=((pred-ytr[ix])**2).mean()
105            opt.zero_grad(); loss.backward(); opt.step()
106        with torch.no_grad():
107            pred=model(xte)
108            field_rmse=float(torch.sqrt(((pred-yte)**2).mean()).cpu())
109            P=torch.tensor((H.T@M).astype(np.float32),dtype=torch.float32,device=device)
110            periods=pred@P.T; target_period=yte@P.T
111            period_rmse=float(torch.sqrt(((periods-target_period)**2).mean()).cpu())
112            # Error separately on edges more than four hops from source.
113            far=list(range(5,n-4))
114            far_rmse=float(torch.sqrt(((pred[:,far]-yte[:,far])**2).mean()).cpu())
115        out[name]={"field_rmse":field_rmse,"period_rmse":period_rmse,"far_edge_rmse":far_rmse,
116                   "parameters":sum(p.numel() for p in model.parameters())}
117    return out
118
119if __name__ == "__main__":
120    checks,D,M,H=math_check()
121    results=run_training(H,M)
122    report={"device":str(device),"math_check":checks,"results":results,
123            "setup":"24-edge weighted cycle with no 2-simplices; source scalar at one node; randomized global harmonic circulation target; 4 message-passing layers; 700 Adam steps; harmonic head has 3122 vs baseline 2497 parameters"}
124    print(json.dumps(report,indent=2))
125    with open("results.json","w") as f: json.dump(report,f,indent=2)