1import json, time
 2import torch
 3import torch.nn as nn
 4
 5SEED = 2370
 6torch.manual_seed(SEED)
 7torch.set_num_threads(4)
 8DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 9try:
10    if DEVICE == "cuda":
11        torch.cuda.set_device(0)
12        torch.empty(1, device="cuda")
13except Exception:
14    DEVICE = "cpu"
15
16def dx(a):
17    return (torch.roll(a, -1, dims=-1) - torch.roll(a, 1, dims=-1)) * 0.5
18
19def dy(a):
20    return (torch.roll(a, -1, dims=-2) - torch.roll(a, 1, dims=-2)) * 0.5
21
22def curl2(a):
23    return torch.stack((dy(a), -dx(a)), dim=1)
24
25def divergence(b):
26    return dx(b[:, 0]) + dy(b[:, 1])
27
28class TinyNet(nn.Module):
29    def __init__(self, out_channels):
30        super().__init__()
31        self.net = nn.Sequential(
32            nn.Conv2d(1, 16, 3, padding=1, padding_mode="circular"), nn.Tanh(),
33            nn.Conv2d(16, 16, 3, padding=1, padding_mode="circular"), nn.Tanh(),
34            nn.Conv2d(16, out_channels, 1))
35    def forward(self, x): return self.net(x)
36
37def make_data(n, h=16, w=16):
38    kx = torch.fft.fftfreq(w).reshape(1, 1, 1, w)
39    ky = torch.fft.fftfreq(h).reshape(1, 1, h, 1)
40    x = torch.randn(n, 1, h, w)
41    f = torch.fft.fft2(x)
42    filt = torch.exp(-((kx / .22) ** 2 + (ky / .22) ** 2))
43    a = torch.fft.ifft2(f * filt).real
44    a = a / (a.std(dim=(-2, -1), keepdim=True) + 1e-7)
45    return a, curl2(a[:, 0])
46
47def train(kind, xtr, ytr, xte, yte, epochs=100):
48    torch.manual_seed(SEED + (1 if kind == "curl" else 2))
49    net = TinyNet(1 if kind == "curl" else 2).to(DEVICE)
50    opt = torch.optim.Adam(net.parameters(), lr=3e-3)
51    xtr, ytr, xte, yte = [z.to(DEVICE) for z in (xtr, ytr, xte, yte)]
52    for _ in range(epochs):
53        pred = net(xtr)
54        if kind == "curl":
55            field = curl2(pred[:, 0])
56            loss = ((field - ytr) ** 2).mean()
57        else:
58            field = pred
59            loss = ((field - ytr) ** 2).mean() + (divergence(field) ** 2).mean()
60        opt.zero_grad(); loss.backward(); opt.step()
61    with torch.no_grad():
62        pred = net(xte)
63        field = curl2(pred[:, 0]) if kind == "curl" else pred
64        div = divergence(field)
65        return {"mse": ((field - yte) ** 2).mean().item(),
66                "max_abs_divergence": div.abs().max().item(),
67                "rms_divergence": torch.sqrt((div ** 2).mean()).item(),
68                "parameters": sum(p.numel() for p in net.parameters())}
69
70def run_models(xtr, ytr, xte, yte):
71    return train("direct", xtr, ytr, xte, yte), train("curl", xtr, ytr, xte, yte)
72
73def main():
74    global DEVICE
75    torch.manual_seed(SEED)
76    a = torch.randn(3, 11, 13, dtype=torch.float64)
77    exact_div = dx(dy(a)) - dy(dx(a))
78    generic = torch.randn(3, 2, 11, 13, dtype=torch.float64)
79    gd = divergence(generic)
80    identity = {"max_abs_DCA": exact_div.abs().max().item(),
81                "rms_DCA": torch.sqrt((exact_div ** 2).mean()).item(),
82                "generic_random_rms_divergence": torch.sqrt((gd ** 2).mean()).item()}
83    xtr, ytr = make_data(192); xte, yte = make_data(64)
84    t0 = time.time()
85    try:
86        baseline, idea = run_models(xtr, ytr, xte, yte)
87    except Exception as exc:
88        if DEVICE != "cuda": raise
89        DEVICE = "cpu"
90        torch.manual_seed(SEED)
91        baseline, idea = run_models(xtr, ytr, xte, yte)
92        identity["cuda_fallback"] = type(exc).__name__ + ": " + str(exc).split("\n")[0]
93    result = {"device": DEVICE, "seed": SEED, "identity_check": identity,
94              "baseline_direct_plus_penalty": baseline, "idea_exact_curl": idea,
95              "elapsed_sec": time.time() - t0}
96    print(json.dumps(result, indent=2))
97    with open("results.json", "w") as f: json.dump(result, f, indent=2)
98
99if __name__ == "__main__": main()