JAX-ALFA: JAX-powered Atmospheric LES For All
Overview
JAX-ALFA is a JAX-based large-eddy simulation framework for atmospheric boundary layer simulations. It leverages JAX’s CPU/GPU/TPU acceleration capabilities to provide highly efficient cross-platform simulations without any code changes.
Features
Incompressible flow solver
Spectral methods for horizontal derivatives
Finite difference methods for vertical derivatives
FFT-based direct Poisson solver
Dynamic SGS coefficient computation
Prognostic potential temperature and specific humidity scalars
Virtual potential temperature buoyancy coupling
JAX-accelerated computations for CPUs & GPUs
Either single or double precision computations
Download
The JAX-ALFA package can be downloaded from: https://github.com/Sukantabasu/jax-alfa
Requirements
Python 3.8+
JAX
NumPy
Getting Started
User Guidelines
Tutorial
Performance Benchmarks
Model Structure
Additional Resources
- API Reference
- Configuration Loader
- Derived Variables
- Import All Modules
- Initialization of Variables
- Preprocessing of Variables
- Main Program
- All Terms
- Advection Terms
- Buoyancy Terms
- Pressure Terms
- SGS Terms
- Time Advancement
- Scalar All Terms
- Scalar Advection Terms
- Scalar SGS Terms
- Scalar Time Advancement
- Strain Rates Computations
- SGS Stress Computations: Smagorinsky
- SGS Stress Computations: Wong-Lilly
- Scalar SGS Flux Computations: Smagorinsky
- Scalar SGS Flux Computations: Wong-Lilly
- Static SGS Modeling: Main Code
- Dynamic SGS Modeling: Main Code
- SGS Model: LASDD-SM (Momentum)
- SGS Model: LASDD-WL (Momentum)
- SGS Model: LASDD-SM (Scalar)
- SGS Model: LASDD-WL (Scalar)
- Surface Flux Parameterization
- FFT Computations
- Derivative Computations
- Dealiasing Computations
- Filtering Computations
- Compute Statistics
- Miscellaneous Terms
Ask a Question
Have a question or found a bug? Please open an issue on GitHub: https://github.com/Sukantabasu/jax-alfa/issues