Sejibeji/agfno-darcy

AGF-NO: Adaptive Geometry-Aware Fourier Neural Operator

Adaptive Geometry-Aware Fourier Neural Operator — a geometry-aware variant of the Fourier Neural Operator (FNO) that injects obstacle geometry directly into the spectral layer, evaluated against the standard FNO on 2-D Darcy flow with random irregular polygonal obstacles.

Results (1,000-sample held-out test set, identical training budgets)

Metric	FNO baseline	AGF-NO (ours)	Improvement
Relative L2 (global)	0.1004	0.0541	−46%
Relative L2 (near-boundary ring)	0.2334	0.0627	−73%
Wall fidelity (mean |pred−target| inside obstacles)	0.0769	0.0129	5.9×
Zero-shot 2× super-resolution (48→96 grid)	0.4666	0.2295	−51%

Both models: width 64, 4 Fourier blocks, 12×12 retained modes, 300 epochs, identical optimizer schedules, trained on a NVIDIA Tesla T4 (Kaggle). Inference throughput on T4: FNO 1,774 samples/s, AGF-NO 925 samples/s.

Method

Standard FNO blocks assume periodic, uniform grids. AGF-NO modifies each Fourier block with two zero-gated geometry paths (so at initialization the network is exactly an FNO, and training learns how much geometry to use):

1. Geometry-aware spectral modulation. The signed-distance field (SDF) of the obstacles is pushed through its own spectral conv; the result multiplicatively modulates the primary spectral output, y = (W u) * (1 + g_spec * tanh(S(sdf))) — a position- and frequency-dependent effective kernel K_eff = W · S that adapts to boundary geometry while preserving O(N log N) FFT scaling.
2. Anti-forgetting injection. Raw SDF + coordinate Fourier features are re-injected into the pointwise MLP branch at every depth, so boundary information never has to survive a Markovian chain of global mixers (the "geometric forgetting" failure mode of deep operators).

Benchmark

2-D Darcy flow, -div(K ∇u) = 1 on the unit square with 1–3 random irregular polygonal obstacles (3–9 jittered vertices) under homogeneous Dirichlet walls, plus a sealed outer frame. K is a smooth log-normal random field (heterogeneous, as in Li et al. 2021). Ground truth comes from a real solver — preconditioned conjugate gradient with a batched shifted-Laplacian multigrid preconditioner, relative residual tolerance 1e-6 — fully batched on GPU. Exact polygon SDFs (winding number + point-segment distance, 4× oversampled) are evaluated identically at 48×48 and 96×96, enabling rigorous zero-shot super-resolution. 2,000 train / 200 val / 1,000 test samples, seed-pinned and byte-reproducible (verified by cross-machine evaluation).

Files

File	Description
fno_final.pt, agfno_final.pt	trained weights (also *_best.pt by val loss)
comparison.png	ground truth vs FNO vs AGF-NO with error maps
training_curves.png	training/validation histories
super_resolution.png	coarse vs 2× super-resolution error
results.json	full configs, metrics, timing, histories (Kaggle run)
reeval_results.json	independent local re-evaluation (cross-check)
train_log.json	raw Kaggle kernel log

Reproduction

The complete pipeline (data generation → training → evaluation → figures) is self-contained and runs on any CUDA GPU:

• Kaggle kernel (one-click, free T4): sehajrsingh/agfno-darcy-full-run
• Source package (Kaggle dataset): sehajrsingh/agfno-darcy-bench-code

Locally: python -m pytest (16 correctness tests: spectral mode truncation, zero-gate FNO equivalence, SDF exactness, solver convergence/mass balance).

Citation

If you use this work, please cite the foundational FNO paper:

@article{li2021fourier,
  title={Fourier Neural Operator for Parametric Partial Differential Equations},
  author={Li, Zongyi and Kovachki, Nikola and Azizzadenesheli, Kamyar and
          Liu, Burigede and Bhattacharya, Kaushik and Stuart, Andrew and
          Anandkumar, Anima},
  journal={ICLR},
  year={2021}
}