boussinesq-fno-weights

Boussinesq FNO weights — variable amplitude

Weights for a Fourier Neural Operator (FNO) trained on the 2D Boussinesq system.

Live demo: https://akrammoustafa.github.io/boussinesq-demo/ Interactive Space: https://huggingface.co/spaces/moust1am/boussinesq-fno-space

What's in the checkpoint

final.pt stores (in a single torch.load dict):

• model_state_dict — the FNO weights.
• arch — the architecture that produced the state dict (n_modes, hidden, layers).
• stats — input normalization: mean, std of the two input channels (ω₀/ε, θ₀/ε), plus le_mean, le_std for the log(ε/0.1) channel.
• s_osc, s_dwb — stored log–log tables for the two output scales as a function of ε.
• epoch — training epoch at save time.

No other hyper-parameters are needed to reconstruct the full operator.

Input / output

• Inputs Two fields ω₀, θ₀ on a 128×128 grid over [0, 2π]², plus a scalar ε. Initial data are normalized so that ‖u₀‖_H² = ‖θ₀‖_H² = ε/2 (with u₀ obtained from ω₀ by Biot–Savart; Nyquist modes of ω₀ removed; mean of θ₀ removed).
• Trained ε range [0.021, 0.476]. Outside this range the network extrapolates.
• Random fields Covariance (τ² − Δ)^(−α), with ω₀: α = 2.5, τ = 7, θ₀: α = 3.5, τ = 7.
• Output The network outputs two channels. The vorticity at T = 100 is reconstructed as ω(T) = ⟨ω₀⟩_{x₁} + s_dwb(ε) · (⟨ch₂⟩_{x₁}) + s_osc(ε) · (ch₁ − ⟨ch₁⟩_{x₁}). Surrogate.omega_T in sc_common.py does this.
• Conventions Array axis -2 is x₁ (horizontal), axis -1 is x₂ (vertical). ν = η = 1, T = 100.

Quick start

Clone the repo (or download final.pt, sc_common.py, and the neuralop/ folder), then from that folder:

python predict_example.py --eps 0.1
python predict_example.py --eps 0.3
python predict_example.py --check     # compare with solver on multieps_test_interp.pt

Minimal code:

import sys, torch
sys.path.insert(0, "<this folder>")
import sc_common as sc

net = sc.Surrogate("final.pt", torch.device("cpu"))
eps = torch.tensor([0.1], dtype=torch.float64)
# w_raw, th_raw: (1, 128, 128) GRF samples
w0, th0, _, _ = sc.prepare_initial(w_raw, th_raw, eps)
with torch.no_grad():
    wT = net.omega_T(w0, th0, eps)   # (1, 128, 128)

Batches work the same way, with eps of shape (B,).

Important

Use the neuralop/ folder included here (put it first on the Python path); do not pip install neuralop. The state-dict keys are tied to this fork of the library.

Requires python 3.10+, torch, numpy.

What changed vs. the old (single-ε) model

1. ε is now an input. The network has three input channels: ω₀/ε, θ₀/ε (standardized), and a constant log(ε/0.1) channel (standardized).
2. Normalization. Old: ‖ω₀‖_H² + ‖θ₀‖_H² = 0.2. New: ‖u₀‖_H² = ‖θ₀‖_H² = ε/2.
3. Output head. Two channels + stored s_osc(ε), s_dwb(ε) scales (log–log interpolation in ε).

See README_for_Akram.txt for the full write-up.