FNO-PINNsFormer

Physics-Aware Transformer: A Comprehensive Study of Transformer Autoregressive Architectures for Solving Partial Differential Equations

Fourier Neural Operator meets Physics-Informed Neural Networks: A research codebase comparing FNO-Transformer against PINN and PINNsFormer baselines across multiple PDEs.

Authors

• Aimen Boukhari
• Aya Benali Khodja
• Mohamed Boulanouar
• Ahmed Thebat Mazouz
• Selssabil Kadid

---

Overview

This repository implements three neural PDE solver architectures and benchmarks them on 1D and 2D partial differential equations:

• FNO-Transformer (proposed): Combines Fourier Neural Operator with PINNsFormer-style Transformer for spatio-temporal PDE solving
• PINN: Physics-Informed Neural Network baseline
• PINNsFormer: Transformer-based PINN baseline (Zhao et al., ICLR 2024)

Tools & Technologies

PyTorch	NumPy	Matplotlib	Python

Supported PDEs

PDE	Spatial Dim	Method	Key Feature
Reaction	1D	Exact solution	Logistic growth
Convection	1D	Spectral shift	Advection
Wave	1D	Spectral + RK4	2-component state [u, v]
Burgers	1D	Pseudo-spectral RK4	Dealiasing
Navier-Stokes	2D	Vorticity-streamfunction	Pseudo-spectral + RK4

Overall Average (Across All Experiments, under 1500 training epochs on an RTX 6000 GPU)

Model	Average Error
PINN	0.4426
PINNsFormer	0.4271
FNO-Transformer	0.1326
FNO-Transformer (Sampled)	0.2275
FNO-Transformer (SoftAttention)	0.1817

Note : All experimental outputs and results are stored in the /experiments/output directory for reproducibility.

Repository Structure

fno-pinnsformer-research/
├── src/
│   └── fno_pinnsformer/
│       ├── __init__.py              # Package init
│       ├── models/                  # Neural network architectures
│       │   ├── activations.py       # WaveAct (learnable wavelet activation)
│       │   ├── embedding.py         # FourierEmbedding
│       │   ├── fno.py               # SpectralConv1d/2d, FNOLayer1d/2d
│       │   ├── transformer.py       # TransformerEncoderLayer, DecoderLayer
│       │   ├── operators.py         # FNOPINNsFormerOperator, FNOPINNsFormerOperator2D
│       │   └── baselines.py         # VanillaPINN, VanillaPINNsFormer
│       ├── pde/                     # PDE definitions
│       │   ├── residuals.py         # Autograd and FD residuals
│       │   └── solvers.py           # Analytical/spectral solvers
│       ├── data/                    # Data generation
│       │   ├── generators.py        # make_1d_dataset, make_ns2d_dataset
│       │   └── ic_registry.py       # Initial condition variants
│       ├── training/                # Training infrastructure
│       │   ├── trainers.py          # CollocationTrainer, TrajectoryTrainer
│       │   └── utils.py             # Early stopping, helpers
│       ├── viz/                     # Visualization
│       │   └── plots.py             # Heatmaps, curves, bar charts
│       └── benchmark/               # Benchmark runner
│           └── runner.py            # run_1d_benchmark
├── experiments/                     # Individual experiment scripts
│   ├── run_reaction.py
│   ├── run_convection.py
│   ├── run_wave.py
│   ├── run_navier_stokes.py
│   ├── run_burgers.py
│   └── run_all.py                   # Run all + dashboard
├── configs/
│   └── experiments/                 # JSON configs for each experiment + run_all
├── notebooks/
│   └── small_experiments.ipynb      # Lightweight notebook sections for quick tests
├── requirements.txt
├── pyproject.toml
└── README.md

Installation

# Clone the repository
git clone https://github.com/SchoolofAI-Algiers/physics-aware-transformer
cd physics-aware-transformer

# Install dependencies
pip install -r requirements.txt

# Install package in development mode
pip install -e .

Quick Start

Run a single experiment

cd ./experiments 
python ./run_reaction.py

Run a single experiment with a custom config

cd ./experiments
python ./run_reaction.py --config configs/experiments/reaction.json

Run all experiments

cd ./experiments
python ./run_all.py --config configs/experiments/run_all.json

Open the small-experiments notebook

Use notebooks/small_experiments.ipynb for quick, reduced-size Reaction and Burgers runs. It also includes a preview cell for the Navier-Stokes rollout GIF.

Generate the Navier-Stokes rollout GIF

cd ./experiments
python ./run_navier_stokes.py

This now saves output/ns2d_fno_prediction.gif alongside the existing NS2D plots.

Use as a library

from fno_pinnsformer.models import FNOPINNsFormerOperator
from fno_pinnsformer.pde import solve_burgers, residual_burgers
from fno_pinnsformer.benchmark import run_1d_benchmark

# Run benchmark
results = run_1d_benchmark(
    pde_name="Burgers",
    solver_fn=solve_burgers,
    solver_kwargs={"nu": 0.01},
    residual_fn=residual_burgers,
    pde_kwargs={"nu": 0.01},
    ic_variants=[...],
    results_dir="output",
)

Key Architecture Components

FNO-Transformer (Proposed)

Input u(t) -> Pseudo-sequence builder -> SpatioTemporalMixer
    -> Transformer Encoder (WaveAct self-attention)
    -> Transformer Decoder (WaveAct cross-attention)
    -> Project to FNO width -> Lifting P
    -> FNO Fourier Layers x depth -> Projection Q
    -> Residual: u(t+dt) = FNO_out + u(t)
    -> Autoregress for full trajectory

WaveAct (Learnable Activation)

sigma(x) = w1 * sin(x) + w2 * cos(x)

Learnable parameters w1, w2 updated via backpropagation. Captures multi-frequency oscillatory dynamics.

---

How to Cite This Project

If you use this codebase in your research, please cite our work using the following BibTeX entry:

@misc{soai2026physics,
  title={Physics-Aware Transformer: A Comprehensive Study of Transformer Autoregressive Architectures for Solving Partial Differential Equations},
  author={Aimen Boukhari and Aya Benali Khodja and Mohamed Boulanouar and Ahmed Thebat Mazouz and Selssabil Kadid},
  year={2026}
}

---

References

The following works have been instrumental in developing this research:

1. Zhao, Z., et al. (2024). "PINNsFormer: A Transformer-based Framework For Physics-Informed Neural Networks." In Proceedings of the International Conference on Learning Representations (ICLR).

2. Li, Z., et al. (2021). "Fourier Neural Operator for Parametric Partial Differential Equations." In Proceedings of the International Conference on Learning Representations (ICLR).

3. Raissi, M., Perdikaris, P., & Karniadakis, G. E. (2019). "Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations." Journal of Computational Physics, 378, 686-707.

4. Vaswani, A., et al. (2017). "Attention is All You Need." In Proceedings of Advances in Neural Information Processing Systems (NeurIPS).

---

Issues & Support

Reporting Issues

If you encounter any bugs, have feature requests, or face issues running the code, please:

1. Check existing issues at GitHub Issues to see if your problem has been addressed
2. Create a new issue with:
   • A clear description of the problem
   • Steps to reproduce the issue
   • Your environment details (Python version, PyTorch version, OS)
   • Any relevant error messages or logs

Contacting the Authors

For questions, collaborations, or feedback not related to specific bugs, feel free to reach out to the authors:

• Aimen Boukhari - email
• Aya Benali Khodja - email
• Mohamed Boulanouar - email
• Ahmed Thebat Mazouz - email
• Selssabil Kadid - email

Alternatively, you can open a discussion on the repository for general questions and feedback.

---

License

MIT