FARGO3D Pre-trained FNO Model Weights: Orzsag-Tang vortex

FARGO3D Pre-trained FNO Model Weights: Orzsag-Tang vortexPre-trained Fourier Neural Operator (FNO) weights for 2D magnetohydrodynamic turbulence surrogate described in:> Duarte, Nemmen & Lima, *Spectral Learning of Magnetized Plasma Dynamics: A Neural Operator Application*, [arXiv:2507.01388](https://arxiv.org/abs/2507.01388)Codebase: [github.com/black-hole-group/fno-vortex](https://github.com/black-hole-group/fno-vortex)---## ContentsOne trained model per physical field, each compressed as a separate archive:| Archive | Field | Physical quantity ||---------|-------|------------------|| `gasdens_weights.7z` | `gasdens` | Gas density ρ || `gasvy_weights.7z` | `gasvy` | Gas velocity y-component v_y || `gasvz_weights.7z` | `gasvz` | Gas velocity z-component v_z || `by_weights.7z` | `by` | Magnetic field y-component B_y || `bz_weights.7z` | `bz` | Magnetic field z-component B_z |Each archive extracts to:```/├── model/│ ├── model_64_30.pt # model weights (~2.8 GB)│ └── loss_64_30.npy # training/validation loss history````model_64_30.pt` is a plain PyTorch state dict saved with `torch.save(model.state_dict(), ...)`.---## ArchitectureThe weights are for `FNO3d(64, 64, 5, 30)`:| Hyperparameter | Value ||----------------|-------|| Fourier modes x | 64 || Fourier modes y | 64 || Fourier modes t | 5 || Hidden width | 30 || Fourier layers | 5 || Input channels | 7 (5 snapshots + ν + μ) || Input shape | `(batch, 128, 128, 10, 7)` || Output shape | `(batch, 128, 128, 10)` — 10 predicted frames |---## Loading weights standalone```pythonimport torchfrom architecture import FNO3d # from src/architecture.py in the repomodel = FNO3d(64, 64, 5, 30).cuda()model.load_state_dict( torch.load('gasdens/model/model_64_30.pt', map_location='cuda', weights_only=True))model.eval()```Or without the repo, copy the `FNO3d` and `SpectralConv3d` class definitions from `src/architecture.py` directly into your script.---## Running inference with the codebaseUse the **`main` branch at commit `283b524`**:```bashgit clone https://github.com/black-hole-group/fno-vortex.gitcd fno-vortexgit checkout 283b524```The inference script expects weights at `experiments//checkpoints/model_64_30.pt`. After extracting the archives, create that structure with the appropriate field name (see mapping below):```bash# Example for gas densitymkdir -p experiments/density/checkpointscp gasdens/model/model_64_30.pt experiments/density/checkpoints/# Magnetic field y-component (name already matches)mkdir -p experiments/by/checkpointscp by/model/model_64_30.pt experiments/by/checkpoints/```**Field name mapping** (archive name → `--param` argument):| Archive field | `--param` value ||---------------|----------------|| `gasdens` | `density` || `gasvy` | `vy` || `gasvz` | `vz` || `by` | `by` || `bz` | `bz` |Then run inference against the test data (see companion data repository):```bashcd srcpython inference.py --param density```Predictions are saved to `experiments//visualizations/pred_.npy`.---## Dependencies```bashpip install torch numpy matplotlib scipy```A CUDA-capable GPU is required.---## Notes- These weights are compatible only with the `main` branch (`283b524`). The current development branches use a different data format and architecture configuration.- The model predicts **10 future frames** per forward pass (not 20 as in the newer Idefix-based pipeline).- Training used the FARGO3D dataset (available in a companion Figshare repository).