mllam/neural-lam

What it solves

Neural-LAM provides a framework for neural weather prediction specifically designed for Limited Area Modeling (LAM). It allows researchers to implement and train graph-based neural networks to emulate atmospheric fluid flow over time for specific geographic regions, offering a modular alternative to traditional numerical weather prediction models.

How it works

The project uses Graph Neural Networks (GNNs) to model weather patterns. It is built on PyTorch and PyTorch Lightning, with GNN implementations handled by PyG (PyTorch Geometric). The system is modular, separating the model architecture, the graph structure (which defines how message-passing occurs), and the data loading process.

Data is handled via "datastores" (like MDPDatastore for zarr files or NpyFilesDatastoreMEPS for numpy files), which load weather data and flatten spatial coordinates into a grid index. A WeatherDataset class then samples this data in time for training and evaluation. The framework supports several model architectures, including versions of GraphCast and hierarchical GNNs.

Who it’s for

It is designed for meteorologists, climate scientists, and AI researchers focusing on regional weather forecasting and atmospheric modeling.

Highlights

• Multiple Model Architectures: Includes implementations of GraphCast, hierarchical models, and other graph-based prediction models.
• Modular Design: Allows independent swapping of models, graphs, and datastores to fit different geographic areas and data formats.
• Flexible Data Pipeline: Supports various input formats through a customizable datastore interface and integrates with mllam-data-prep for optimized training data.
• Scientific Grounding: Directly implements models from peer-reviewed publications in the field of neural weather prediction.