Skip to content

Examples

Every example listed here runs successfully and produces the documented output. Start with Getting Started if you're new, or jump to the category that matches your problem.

Prerequisites

source ./activate.sh
python -c "import opifex; print('Opifex imported successfully')"

Getting Started

New to Opifex? Start here with minimal, self-contained examples.

Example Time Description
Your First Neural Operator 5 min Train an FNO on Darcy flow in ~50 lines
Your First PINN 5 min Solve a Poisson equation with PINNs in ~40 lines

Neural Operators

Data-driven operator learning: map input functions to output functions.

Example Architecture Dataset Description
FNO on Darcy Flow FNO + GridEmbedding2D Darcy Flow 64x64 Standard FNO benchmark
FNO on Burgers FNO Burgers 1D Time-dependent PDE
SFNO Climate (Simple) Spherical FNO Shallow Water 32x32 Quick start for spherical data
SFNO Climate (Full) Spherical FNO + Conservation Shallow Water 64x64 Conservation-aware training
U-FNO on Turbulence U-FNO + Energy Loss 2D Burgers 64x64 Multi-scale architecture
UNO on Darcy Flow UNO + Super-Resolution Darcy Flow 32x32 Zero-shot super-resolution
TFNO on Darcy Flow Tensorized FNO Darcy Flow 64x64 Memory-efficient decomposition
Local FNO on Darcy Local FNO Darcy Flow 64x64 Local + global frequency mixing
GNO on Darcy Flow Graph Neural Operator Darcy Flow 32x32 Irregular mesh support
DeepONet on Darcy DeepONet Darcy Flow 64x64 Branch-trunk architecture
DeepONet Antiderivative DeepONet Antiderivative Classic DeepONet benchmark
PINO on Burgers Physics-Informed NO Burgers 1D Hybrid data + physics loss
Operator Comparison Tour All architectures Multiple Overview of all 26 operators 
python examples/neural-operators/fno_darcy.py

PINNs

Solve PDEs from governing equations using physics-informed neural networks.

Example PDE Type Description
Poisson Equation Elliptic Classic Laplace equation benchmark
Heat Equation Parabolic 2D diffusion with time evolution
Burgers Equation Nonlinear Shock formation and viscosity
Wave Equation Hyperbolic 1D wave propagation
Helmholtz Equation Oscillatory Frequency-domain wave equation
Advection Equation Hyperbolic Transport phenomena
Allen-Cahn Equation Reaction-Diffusion Phase-field dynamics
Diffusion-Reaction Coupled Multi-physics systems
Navier-Stokes Fluid Dynamics Kovasznay flow benchmark
Euler Beam Structural 4th-order ODE for beam deflection
Inverse Diffusion Inverse Problem Parameter discovery from data 
python examples/pinns/poisson.py

Domain Decomposition

Parallel and decomposed PINNs for large-scale problems.

Example Method Description
FBPINN on Poisson Finite Basis PINN Overlapping subdomains with window functions
XPINN on Helmholtz Extended PINN Non-overlapping domain decomposition
CPINN on Advection-Diffusion Conservative PINN Flux conservation at interfaces 
python examples/domain-decomposition/fbpinn_poisson.py

Advanced Training

Techniques for improving PINN training dynamics and convergence.

Example Technique Description
NTK Analysis Neural Tangent Kernel Eigenvalue spectrum and spectral bias detection
GradNorm Adaptive Loss Balancing Automatic loss weight adjustment
Adaptive Sampling RAR-D Refinement Residual-based collocation point refinement 
python examples/advanced-training/ntk_analysis.py

Optimization

Learn-to-optimize and meta-learning for PDE solving.

Example Method Description
Learn to Optimize L2O Learned optimizers for parametric PDEs
Meta-Optimization MAML/Reptile Fast adaptation across PDE families 
python examples/optimization/learn_to_optimize.py

Uncertainty Quantification

Calibration, Bayesian inference, and uncertainty-aware operators.

Example Method Description
Calibration Methods Platt, Isotonic, Conformal Post-hoc calibration and conformal prediction
UQNO on Darcy Bayesian Spectral Conv Uncertainty-aware neural operator
Bayesian FNO Variational Framework Amortized variational inference 
python examples/uncertainty/calibration.py

Quantum Chemistry

Neural approaches to density functional theory.

Example Method Description
Neural DFT Neural SCF Solver H2 molecule ground state
Neural XC Functional Learned Exchange-Correlation Training on LDA reference 
python examples/quantum-chemistry/neural_dft.py

Layers & Components

Individual neural operator building blocks demonstrated in isolation.

Example Component Description
DISCO Convolutions DiscoConv2D Discrete-continuous convolutions with 6x+ speedup
Grid Embeddings GridEmbedding2D Coordinate injection and positional encoding
Fourier Continuation FourierContinuation Boundary handling for non-periodic domains
Spectral Normalization SpectralNormalization Training stability for deep operators 
python examples/layers/disco_convolutions_example.py

Equation Discovery

Recover governing equations directly from data.

Example Method Description
SINDy on Lorenz System SINDy Sparse identification of Lorenz equations from trajectories 
python examples/discovery/sindy_lorenz.py

Fields

JAX-native field abstractions for differential operators and fluid simulation.

Example Focus Description
Field Operations Differential Operators Gradient, Laplacian, advection, pressure projection on grids 
python examples/fields/field_operations.py

Data & Analysis

Explore and validate the synthetic datasets used by neural operator examples.

Example Focus
Darcy Flow Analysis FNO prediction validation, error analysis
Spectral Analysis Power spectrum and mode analysis
PDEBench Data Loading Loading PDEBench HDF5 datasets 
python examples/data/darcy_flow_analysis.py

Distributed Computing

Multi-device data-parallel training for large-scale problems.

Example Method Description
Distributed PDE Training Data Parallel Multi-GPU training with JAX sharding 
python examples/distributed/distributed_pde.py

Benchmarking

Performance comparisons and GPU optimization guides.

Example Focus Runtime
Neural Operator Benchmark UNO, FNO, SFNO across resolutions ~15 min
GPU Profiling & Optimization Memory pools, mixed precision, JIT analysis ~5 min 
python examples/benchmarking/operator_benchmark.py

Troubleshooting

Import errors: Ensure the environment is activated with source ./activate.sh.

GPU availability: Check with python -c "import jax; print(jax.default_backend(), jax.devices())".

Memory issues: Reduce BATCH_SIZE or N_TRAIN in the configuration section of the example.

Additional Resources