"""FFT weak-residual engine and numerical verification. Run with: python3 fft_weak_engine.py All transforms use numpy's convention xhat=fft(x), so Parseval is mean(|x|^2)=sum(|xhat|^2)/N^2. """ import json import time import numpy as np SEED = 2085 rng = np.random.default_rng(SEED) def integer_modes(n): return np.fft.fftfreq(n) * n def parseval_error(n): x = rng.normal(size=n) + 1j * rng.normal(size=n) xhat = np.fft.fft(x) a = np.mean(np.abs(x) ** 2) b = np.sum(np.abs(xhat) ** 2) / n**2 return abs(a - b) / a def derivative_error(n, ks): x = np.arange(n) / n errors = [] for k in ks: phi = np.exp(2j * np.pi * k * x) dphi = np.fft.ifft(2j * np.pi * integer_modes(n) * np.fft.fft(phi)) exact = 2j * np.pi * k * phi errors.append(np.linalg.norm(dphi - exact) / np.linalg.norm(exact)) return max(errors) def bandlimited_signal(n, kmax): """Real random signal with Fourier support |k|<=kmax on an n grid.""" h = np.zeros(n, dtype=complex) ks = integer_modes(n).astype(int) keep = np.abs(ks) <= kmax h[keep] = rng.normal(size=keep.sum()) + 1j * rng.normal(size=keep.sum()) # Hermitian symmetrization gives a real signal while retaining support. h = (h + np.conj(h[(-np.arange(n)) % n])) / 2 return np.fft.ifft(h).real def product_spectrum(u, n, dealias): """Spectrum of u^2/2, optionally using 3/2 padding.""" if not dealias: return np.fft.fft(0.5 * u * u) # Pad the Fourier coefficients to M=3N/2, multiply in physical space, # then crop the central N Fourier modes. This removes quadratic aliases # from the retained 2/3 band. m = 3 * n // 2 uh = np.fft.fft(u) up = np.zeros(m, dtype=complex) low = n // 2 up[:low] = uh[:low] up[-low:] = uh[-low:] u_pad = np.fft.ifft(up) * (m / n) fh_pad = np.fft.fft(0.5 * u_pad * u_pad) out = np.zeros(n, dtype=complex) out[:low] = fh_pad[:low] * (n / m) out[-low:] = fh_pad[-low:] * (n / m) return out def alias_sweep(n=96): """Compare raw and 3/2-dealiased coefficients to a 4N reference.""" rows = [] for frac in (0.10, 0.20, 0.30, 0.333, 0.40, 0.45): kmax = max(1, int(frac * n)) u = bandlimited_signal(n, kmax) # Use identical continuous Fourier coefficients on a 4N grid by # interpolating via the Fourier representation. nr = 4 * n uh = np.fft.fft(u) up = np.zeros(nr, dtype=complex) low = n // 2 up[:low] = uh[:low] up[-low:] = uh[-low:] ur = np.fft.ifft(up) * (nr / n) ref = np.fft.fft(0.5 * ur.real**2) # Compare only modes representable on the original grid. refc = np.zeros(n, dtype=complex) refc[:low] = ref[:low] * (n / nr) refc[-low:] = ref[-low:] * (n / nr) raw = product_spectrum(u, n, False) deal = product_spectrum(u, n, True) # Relevant retained modes are the standard 2/3 cutoff. ks = np.abs(integer_modes(n)) <= n / 3 denom = np.linalg.norm(refc[ks]) + 1e-30 rows.append({ "kmax_over_N": frac, "raw_rel_error_2/3": float(np.linalg.norm(raw[ks]-refc[ks])/denom), "dealiased_rel_error_2/3": float(np.linalg.norm(deal[ks]-refc[ks])/denom), }) return rows def retained_energy_sweep(n=256): """Parseval prediction: retained Fourier energy equals retained physical energy.""" x = np.arange(n) / n # Smooth function with known decreasing spectrum. u = np.sin(2*np.pi*3*x) + .5*np.sin(2*np.pi*11*x) + .2*np.sin(2*np.pi*31*x) h = np.fft.fft(u) total = np.sum(np.abs(h)**2) rows = [] for cutoff in (2, 4, 8, 16, 32, 64): keep = np.abs(integer_modes(n)) <= cutoff spectral_fraction = np.sum(np.abs(h[keep])**2) / total low = np.fft.ifft(h * keep).real physical_fraction = np.mean(low**2) / np.mean(u**2) rows.append({"cutoff": cutoff, "spectral_fraction": float(spectral_fraction), "physical_projection_fraction": float(physical_fraction), "abs_gap": float(abs(spectral_fraction-physical_fraction))}) return rows def fft_vs_direct(n=256, batch=64, repeats=30, modes=32): """Projection timing: FFT all modes versus explicit trigonometric bank.""" x = np.arange(n) / n k = np.arange(-modes, modes + 1) bank = np.exp(-2j*np.pi*np.outer(k, x)) a = rng.normal(size=(batch, n)) # warmup np.fft.fft(a, axis=1) bank @ a[0] t0 = time.perf_counter() for _ in range(repeats): z = np.fft.fft(a, axis=1)[:, np.mod(k, n)] fft_time = (time.perf_counter()-t0)/repeats t0 = time.perf_counter() for _ in range(repeats): z2 = a @ bank.T direct_time = (time.perf_counter()-t0)/repeats # Verify the selected coefficients agree exactly up to roundoff. check = np.max(np.abs(z - z2)) / (np.max(np.abs(z2)) + 1e-30) return {"fft_ms": 1000*fft_time, "direct_ms": 1000*direct_time, "direct_over_fft": direct_time/fft_time, "projection_rel_check": float(check)} def weak_loss(u, ut, n_modes=None): """Periodic Burgers weak residual at one time slice. Rhat_k = -(ut_hat + 2 pi i k * flux_hat), with flux=u^2/2. """ n = u.shape[-1] ks = integer_modes(n) uh = np.fft.fft(ut) fh = np.fft.fft(.5*u*u) r = -(uh + 2j*np.pi*ks*fh) if n_modes is not None: r = r[np.abs(ks) <= n_modes] return float(np.mean(np.abs(r)**2) / n**2) def torch_training_benchmark(): """Small fixed-data benchmark, with CPU fallback if CUDA is unavailable.""" try: import torch device = "cuda" if torch.cuda.is_available() else "cpu" torch.manual_seed(SEED) if device == "cuda": try: torch.cuda.empty_cache() except Exception: device = "cpu" n, nt, width = 96, 8, 32 x = torch.arange(n, device=device).float()/n t = torch.linspace(0, .15, nt, device=device) # Smooth exact-ish Burgers snapshots from a low-amplitude traveling wave. xx = x[None, :] - .7*t[:, None] target = .35*torch.sin(2*torch.pi*xx) + .12*torch.sin(4*torch.pi*xx) inp = torch.stack(torch.meshgrid(t, x, indexing='ij'), -1).reshape(-1,2) def run(kind): torch.manual_seed(SEED+kind) net = torch.nn.Sequential(torch.nn.Linear(2,width), torch.nn.Tanh(), torch.nn.Linear(width,width), torch.nn.Tanh(), torch.nn.Linear(width,1)).to(device) opt = torch.optim.Adam(net.parameters(), lr=2e-3) vals=[]; t0=time.perf_counter() for _ in range(80): pred=net(inp).reshape(nt,n) fit=((pred-target)**2).mean() if kind == 0: # FFT weak residual using finite time differences. dt=t[1]-t[0] ut=(pred[1:]-pred[:-1])/dt uu=pred[:-1] fh=torch.fft.fft(.5*uu*uu, dim=1) uh=torch.fft.fft(ut, dim=1) kk=torch.fft.fftfreq(n, device=device)*n rr=uh+2j*torch.pi*kk[None,:]*fh phys=(rr.abs()**2).mean()/n**2 else: # strong pointwise spatial derivative by autograd. q=net(inp).reshape(nt,n) grads=torch.autograd.grad(q.sum(), inp, create_graph=True)[0] ux=grads[:,1].reshape(nt,n) ut=torch.autograd.grad(q.sum(), inp, create_graph=True)[0][:,0].reshape(nt,n) phys=((ut+q*ux)**2).mean() loss=fit+0.02*phys opt.zero_grad(); loss.backward(); opt.step(); vals.append(float(loss.detach().cpu())) return {"final_loss": vals[-1], "seconds": time.perf_counter()-t0, "l2": float(torch.sqrt(((net(inp).reshape(nt,n)-target)**2).mean()).detach().cpu())} return {"device":device, "fft_weak":run(0), "strong_autodiff":run(1)} except Exception as e: return {"error": type(e).__name__+": "+str(e)} def main(): result = { "seed": SEED, "predictions": { "parseval_identity": {"predicted_relative_error": "machine precision", "observed": parseval_error(257)}, "exact_test_derivative": {"predicted_relative_error": "machine precision for resolved modes", "observed": derivative_error(256, [1,7,31,63])}, "dealiasing": {"prediction": "quadratic products are accurate in retained 2/3 band when kmax<=N/3; raw aliasing rises beyond it", "observed": alias_sweep()}, "frequency_truncation": {"prediction": "retained spectral energy equals Parseval physical projection energy", "observed": retained_energy_sweep()}, }, "speed_benchmark": fft_vs_direct(), "burgers_mini_experiment": torch_training_benchmark(), } with open("results.json", "w") as f: json.dump(result, f, indent=2) print(json.dumps(result, indent=2)) if __name__ == "__main__": main()