1import json, math, random
  2import numpy as np
  3import torch
  4import torch.nn.functional as F
  5
  6SEED = 314
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(4)
  9DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 10try:
 11    if DEVICE == "cuda":
 12        torch.cuda.empty_cache()
 13except Exception:
 14    DEVICE = "cpu"
 15
 16N = 32
 17Y, X = torch.meshgrid(torch.arange(N, device=DEVICE), torch.arange(N, device=DEVICE), indexing="ij")
 18
 19def phase_field(centers):
 20    # A periodic-grid engineering toy: use two well-separated phase singularities.
 21    # The opposite signs make the far field approximately constant.
 22    ph = torch.zeros((N, N), device=DEVICE)
 23    for cy, cx, charge in centers:
 24        dy = Y.float() - cy
 25        dx = X.float() - cx
 26        ph = ph + charge * torch.atan2(dy, dx)
 27    return ph
 28
 29def unit_from_phase(ph):
 30    return torch.stack((torch.cos(ph), torch.sin(ph)), dim=0)
 31
 32def wrapped(d):
 33    return torch.atan2(torch.sin(d), torch.cos(d))
 34
 35def winding(u):
 36    ph = torch.atan2(u[1], u[0])
 37    # plaquette circulation with forward edges, periodic boundaries
 38    dx = wrapped(torch.roll(ph, -1, 1) - ph)
 39    dy = wrapped(torch.roll(ph, -1, 0) - ph)
 40    circulation = dx + torch.roll(dy, -1, 1) - torch.roll(dx, -1, 0) - dy
 41    return circulation / (2.0 * math.pi)
 42
 43def frac_energy(u, s=0.5):
 44    # torch FFT normalization: divide by N^2 so energy is resolution-comparable.
 45    fy = torch.fft.fftfreq(N, device=u.device) * N
 46    fx = torch.fft.fftfreq(N, device=u.device) * N
 47    ky, kx = torch.meshgrid(fy, fx, indexing="ij")
 48    rad = torch.sqrt(kx*kx + ky*ky)
 49    weight = torch.where(rad > 0, rad ** (1.0 + s), torch.zeros_like(rad))
 50    z = torch.fft.fft2(u)
 51    return (weight[None] * z.abs().square()).sum() / (N*N)**2
 52
 53def tv_energy(u):
 54    return (u.roll(-1, 0) - u).abs().sum() + (u.roll(-1, 1) - u).abs().sum()
 55
 56def task_loss(u, target):
 57    return (u-target).square().mean()
 58
 59def defect_stats(w, threshold=0.25):
 60    # Winding lives on plaquettes. Report all nonzero plaquettes, not vortex centers.
 61    nz = torch.nonzero(w.abs() >= threshold, as_tuple=False)
 62    vals = [(int(y), int(x), float(w[y, x])) for y, x in nz]
 63    return {"count": len(vals), "plaquettes": vals}
 64
 65# ---- Core numerical checks, before optimization ----
 66s = 0.5
 67ks = [1, 2, 4, 8]
 68mode_energies = []
 69for k in ks:
 70    ph = 0.7 * torch.sin(2*math.pi*k*X/N)
 71    u = unit_from_phase(ph)
 72    mode_energies.append(float(frac_energy(u, s)))
 73# A small phase amplitude makes the expected single-frequency power law especially clean.
 74log_slope = float(np.polyfit(np.log(ks), np.log(np.maximum(mode_energies, 1e-30)), 1)[0])
 75centers = [(N//2-5, N//2-5, 1), (N//2+5, N//2+5, -1)]
 76target_phase = phase_field(centers)
 77target = unit_from_phase(target_phase)
 78target_w = winding(target).detach()
 79initial_phase = target_phase + 0.65 * torch.randn((N,N), device=DEVICE)
 80initial = unit_from_phase(initial_phase)
 81initial_w = winding(initial).detach()
 82math_check = {
 83    "frequency_ks": ks,
 84    "fractional_energies": mode_energies,
 85    "loglog_slope_expected_1_plus_s": 1+s,
 86    "loglog_slope_observed": log_slope,
 87    "target_winding": defect_stats(target_w),
 88    "initial_winding": defect_stats(initial_w),
 89    "target_integer_defect_count": int((target_w.abs() > .5).sum().item())
 90}
 91
 92# Direct field fitting isolates the regularizer/topology behavior and has matched parameters.
 93def optimize(kind, steps=700):
 94    # inverse stereographic-style unconstrained channels, normalized in the loss
 95    p = torch.nn.Parameter(torch.logit(((initial[0]+1)/2).clamp(.01,.99)).detach().clone())
 96    q = torch.nn.Parameter(torch.logit(((initial[1]+1)/2).clamp(.01,.99)).detach().clone())
 97    # A direct angle parameter is better conditioned and still exactly S1-valued.
 98    angle = torch.nn.Parameter(torch.atan2(initial[1], initial[0]).detach().clone())
 99    opt = torch.optim.Adam([angle], lr=0.08)
100    for step in range(steps):
101        u = unit_from_phase(angle)
102        fit = task_loss(u, target)
103        if kind == "tv":
104            reg = tv_energy(u) / (2*N*N)
105            loss = fit + 0.018 * reg
106        else:
107            reg = frac_energy(u, s)
108            # Squared winding mismatch; target is the discrete current.
109            jac = winding(u)
110            jacloss = (jac-target_w).square().mean()
111            loss = fit + 0.004 * reg + 10.0 * jacloss
112        opt.zero_grad(); loss.backward(); opt.step()
113    with torch.no_grad():
114        u = unit_from_phase(angle)
115        w = winding(u)
116        # high-frequency fraction: energy above quarter Nyquist
117        fy = torch.fft.fftfreq(N, device=DEVICE)*N; fx=fy
118        ky,kx=torch.meshgrid(fy,fx,indexing='ij'); rad=torch.sqrt(kx*kx+ky*ky)
119        z=torch.fft.fft2(u); spec=z.abs().square().sum(0)/(N*N)**2
120        hf=float(spec[rad > N/4].sum().item())
121        return {
122            "task_mse": float(task_loss(u,target).item()),
123            "fractional_energy": float(frac_energy(u,s).item()),
124            "tv_energy": float((tv_energy(u)/(2*N*N)).item()),
125            "winding_mse": float((w-target_w).square().mean().item()),
126            "winding_defects": defect_stats(w),
127            "integer_defect_count": int((w.abs()>.5).sum().item()),
128            "high_frequency_spectrum": hf,
129            "final_winding_range": [float(w.min()), float(w.max())]
130        }
131
132# Wrapped phase circulation is piecewise constant. At a smooth constant field,
133# its gradient is generally exactly zero, so this engineering surrogate cannot
134# nucleate a defect from a topology-free initialization.
135probe = torch.nn.Parameter(torch.zeros((N, N), device=DEVICE))
136probe_u = unit_from_phase(probe)
137probe_loss = (winding(probe_u) - target_w).square().mean()
138probe_loss.backward()
139math_check["jacobian_loss_at_constant_field"] = float(probe_loss.item())
140math_check["jacobian_gradient_norm_at_constant_field"] = float(probe.grad.norm().item())
141
142results = {"device": DEVICE, "N": N, "math_check": math_check,
143           "baseline_tv": optimize("tv"), "idea_fractional_jacobian": optimize("idea")}
144print(json.dumps(results, indent=2))
145with open("results.json", "w") as f: json.dump(results, f, indent=2)