1"""FFT weak-residual engine and numerical verification.
  2
  3Run with: python3 fft_weak_engine.py
  4All transforms use numpy's convention xhat=fft(x), so Parseval is
  5mean(|x|^2)=sum(|xhat|^2)/N^2.
  6"""
  7import json
  8import time
  9import numpy as np
 10
 11SEED = 2085
 12rng = np.random.default_rng(SEED)
 13
 14
 15def integer_modes(n):
 16    return np.fft.fftfreq(n) * n
 17
 18
 19def parseval_error(n):
 20    x = rng.normal(size=n) + 1j * rng.normal(size=n)
 21    xhat = np.fft.fft(x)
 22    a = np.mean(np.abs(x) ** 2)
 23    b = np.sum(np.abs(xhat) ** 2) / n**2
 24    return abs(a - b) / a
 25
 26
 27def derivative_error(n, ks):
 28    x = np.arange(n) / n
 29    errors = []
 30    for k in ks:
 31        phi = np.exp(2j * np.pi * k * x)
 32        dphi = np.fft.ifft(2j * np.pi * integer_modes(n) * np.fft.fft(phi))
 33        exact = 2j * np.pi * k * phi
 34        errors.append(np.linalg.norm(dphi - exact) / np.linalg.norm(exact))
 35    return max(errors)
 36
 37
 38def bandlimited_signal(n, kmax):
 39    """Real random signal with Fourier support |k|<=kmax on an n grid."""
 40    h = np.zeros(n, dtype=complex)
 41    ks = integer_modes(n).astype(int)
 42    keep = np.abs(ks) <= kmax
 43    h[keep] = rng.normal(size=keep.sum()) + 1j * rng.normal(size=keep.sum())
 44    # Hermitian symmetrization gives a real signal while retaining support.
 45    h = (h + np.conj(h[(-np.arange(n)) % n])) / 2
 46    return np.fft.ifft(h).real
 47
 48
 49def product_spectrum(u, n, dealias):
 50    """Spectrum of u^2/2, optionally using 3/2 padding."""
 51    if not dealias:
 52        return np.fft.fft(0.5 * u * u)
 53    # Pad the Fourier coefficients to M=3N/2, multiply in physical space,
 54    # then crop the central N Fourier modes. This removes quadratic aliases
 55    # from the retained 2/3 band.
 56    m = 3 * n // 2
 57    uh = np.fft.fft(u)
 58    up = np.zeros(m, dtype=complex)
 59    low = n // 2
 60    up[:low] = uh[:low]
 61    up[-low:] = uh[-low:]
 62    u_pad = np.fft.ifft(up) * (m / n)
 63    fh_pad = np.fft.fft(0.5 * u_pad * u_pad)
 64    out = np.zeros(n, dtype=complex)
 65    out[:low] = fh_pad[:low] * (n / m)
 66    out[-low:] = fh_pad[-low:] * (n / m)
 67    return out
 68
 69
 70def alias_sweep(n=96):
 71    """Compare raw and 3/2-dealiased coefficients to a 4N reference."""
 72    rows = []
 73    for frac in (0.10, 0.20, 0.30, 0.333, 0.40, 0.45):
 74        kmax = max(1, int(frac * n))
 75        u = bandlimited_signal(n, kmax)
 76        # Use identical continuous Fourier coefficients on a 4N grid by
 77        # interpolating via the Fourier representation.
 78        nr = 4 * n
 79        uh = np.fft.fft(u)
 80        up = np.zeros(nr, dtype=complex)
 81        low = n // 2
 82        up[:low] = uh[:low]
 83        up[-low:] = uh[-low:]
 84        ur = np.fft.ifft(up) * (nr / n)
 85        ref = np.fft.fft(0.5 * ur.real**2)
 86        # Compare only modes representable on the original grid.
 87        refc = np.zeros(n, dtype=complex)
 88        refc[:low] = ref[:low] * (n / nr)
 89        refc[-low:] = ref[-low:] * (n / nr)
 90        raw = product_spectrum(u, n, False)
 91        deal = product_spectrum(u, n, True)
 92        # Relevant retained modes are the standard 2/3 cutoff.
 93        ks = np.abs(integer_modes(n)) <= n / 3
 94        denom = np.linalg.norm(refc[ks]) + 1e-30
 95        rows.append({
 96            "kmax_over_N": frac,
 97            "raw_rel_error_2/3": float(np.linalg.norm(raw[ks]-refc[ks])/denom),
 98            "dealiased_rel_error_2/3": float(np.linalg.norm(deal[ks]-refc[ks])/denom),
 99        })
100    return rows
101
102
103def retained_energy_sweep(n=256):
104    """Parseval prediction: retained Fourier energy equals retained physical energy."""
105    x = np.arange(n) / n
106    # Smooth function with known decreasing spectrum.
107    u = np.sin(2*np.pi*3*x) + .5*np.sin(2*np.pi*11*x) + .2*np.sin(2*np.pi*31*x)
108    h = np.fft.fft(u)
109    total = np.sum(np.abs(h)**2)
110    rows = []
111    for cutoff in (2, 4, 8, 16, 32, 64):
112        keep = np.abs(integer_modes(n)) <= cutoff
113        spectral_fraction = np.sum(np.abs(h[keep])**2) / total
114        low = np.fft.ifft(h * keep).real
115        physical_fraction = np.mean(low**2) / np.mean(u**2)
116        rows.append({"cutoff": cutoff, "spectral_fraction": float(spectral_fraction),
117                     "physical_projection_fraction": float(physical_fraction),
118                     "abs_gap": float(abs(spectral_fraction-physical_fraction))})
119    return rows
120
121
122def fft_vs_direct(n=256, batch=64, repeats=30, modes=32):
123    """Projection timing: FFT all modes versus explicit trigonometric bank."""
124    x = np.arange(n) / n
125    k = np.arange(-modes, modes + 1)
126    bank = np.exp(-2j*np.pi*np.outer(k, x))
127    a = rng.normal(size=(batch, n))
128    # warmup
129    np.fft.fft(a, axis=1)
130    bank @ a[0]
131    t0 = time.perf_counter()
132    for _ in range(repeats):
133        z = np.fft.fft(a, axis=1)[:, np.mod(k, n)]
134    fft_time = (time.perf_counter()-t0)/repeats
135    t0 = time.perf_counter()
136    for _ in range(repeats):
137        z2 = a @ bank.T
138    direct_time = (time.perf_counter()-t0)/repeats
139    # Verify the selected coefficients agree exactly up to roundoff.
140    check = np.max(np.abs(z - z2)) / (np.max(np.abs(z2)) + 1e-30)
141    return {"fft_ms": 1000*fft_time, "direct_ms": 1000*direct_time,
142            "direct_over_fft": direct_time/fft_time, "projection_rel_check": float(check)}
143
144
145def weak_loss(u, ut, n_modes=None):
146    """Periodic Burgers weak residual at one time slice.
147    Rhat_k = -(ut_hat + 2 pi i k * flux_hat), with flux=u^2/2.
148    """
149    n = u.shape[-1]
150    ks = integer_modes(n)
151    uh = np.fft.fft(ut)
152    fh = np.fft.fft(.5*u*u)
153    r = -(uh + 2j*np.pi*ks*fh)
154    if n_modes is not None:
155        r = r[np.abs(ks) <= n_modes]
156    return float(np.mean(np.abs(r)**2) / n**2)
157
158
159def torch_training_benchmark():
160    """Small fixed-data benchmark, with CPU fallback if CUDA is unavailable."""
161    try:
162        import torch
163        device = "cuda" if torch.cuda.is_available() else "cpu"
164        torch.manual_seed(SEED)
165        if device == "cuda":
166            try: torch.cuda.empty_cache()
167            except Exception: device = "cpu"
168        n, nt, width = 96, 8, 32
169        x = torch.arange(n, device=device).float()/n
170        t = torch.linspace(0, .15, nt, device=device)
171        # Smooth exact-ish Burgers snapshots from a low-amplitude traveling wave.
172        xx = x[None, :] - .7*t[:, None]
173        target = .35*torch.sin(2*torch.pi*xx) + .12*torch.sin(4*torch.pi*xx)
174        inp = torch.stack(torch.meshgrid(t, x, indexing='ij'), -1).reshape(-1,2)
175        def run(kind):
176            torch.manual_seed(SEED+kind)
177            net = torch.nn.Sequential(torch.nn.Linear(2,width), torch.nn.Tanh(),
178                                      torch.nn.Linear(width,width), torch.nn.Tanh(),
179                                      torch.nn.Linear(width,1)).to(device)
180            opt = torch.optim.Adam(net.parameters(), lr=2e-3)
181            vals=[]; t0=time.perf_counter()
182            for _ in range(80):
183                pred=net(inp).reshape(nt,n)
184                fit=((pred-target)**2).mean()
185                if kind == 0: # FFT weak residual using finite time differences.
186                    dt=t[1]-t[0]
187                    ut=(pred[1:]-pred[:-1])/dt
188                    uu=pred[:-1]
189                    fh=torch.fft.fft(.5*uu*uu, dim=1)
190                    uh=torch.fft.fft(ut, dim=1)
191                    kk=torch.fft.fftfreq(n, device=device)*n
192                    rr=uh+2j*torch.pi*kk[None,:]*fh
193                    phys=(rr.abs()**2).mean()/n**2
194                else: # strong pointwise spatial derivative by autograd.
195                    q=net(inp).reshape(nt,n)
196                    grads=torch.autograd.grad(q.sum(), inp, create_graph=True)[0]
197                    ux=grads[:,1].reshape(nt,n)
198                    ut=torch.autograd.grad(q.sum(), inp, create_graph=True)[0][:,0].reshape(nt,n)
199                    phys=((ut+q*ux)**2).mean()
200                loss=fit+0.02*phys
201                opt.zero_grad(); loss.backward(); opt.step(); vals.append(float(loss.detach().cpu()))
202            return {"final_loss": vals[-1], "seconds": time.perf_counter()-t0,
203                    "l2": float(torch.sqrt(((net(inp).reshape(nt,n)-target)**2).mean()).detach().cpu())}
204        return {"device":device, "fft_weak":run(0), "strong_autodiff":run(1)}
205    except Exception as e:
206        return {"error": type(e).__name__+": "+str(e)}
207
208
209def main():
210    result = {
211      "seed": SEED,
212      "predictions": {
213        "parseval_identity": {"predicted_relative_error": "machine precision", "observed": parseval_error(257)},
214        "exact_test_derivative": {"predicted_relative_error": "machine precision for resolved modes", "observed": derivative_error(256, [1,7,31,63])},
215        "dealiasing": {"prediction": "quadratic products are accurate in retained 2/3 band when kmax<=N/3; raw aliasing rises beyond it", "observed": alias_sweep()},
216        "frequency_truncation": {"prediction": "retained spectral energy equals Parseval physical projection energy", "observed": retained_energy_sweep()},
217      },
218      "speed_benchmark": fft_vs_direct(),
219      "burgers_mini_experiment": torch_training_benchmark(),
220    }
221    with open("results.json", "w") as f: json.dump(result, f, indent=2)
222    print(json.dumps(result, indent=2))
223
224if __name__ == "__main__": main()