geoarches

Machine learning framework for geospatial data, mainly climate and weather.

Documentation

What is geoarches?

geoarches is a machine learning library for training, running and evaluating models on weather and climate data, developed by Guillaume Couairon and Renu Singh in the ARCHES team at INRIA (Paris, France).

geoarches's building blocks can easily be integrated into any research ML pipelines. It can also be used to run the ArchesWeather and ArchesWeatherGen weather models.

geoarches builds on PyTorch, PyTorch Lightning and Hydra. Once installed, you can use its modules inside your own project, or use the main training and evaluating workflows. To develop your own models or modify existing ones, the intended usage is to work in your own working directory and create your own configurations files and Lightning modules. See the User Guide for the full details.

Code Overview

geoarches is meant to jumpstart your ML pipeline with building blocks for data handling, model training, and evaluation. This is an effort to share engineering tools and research knowledge across projects.

Data

• download/: Parallelized dataset download scripts with support for chunking to speed up read access.
• dataloaders/: PyTorch datasets for loading and preprocessing NetCDF files into ML-ready tensors.

Model training

• backbones/: Network architectures that plug into Lightning modules.
• lightning_modules/: Training and inference wrappers that are agnostic to the backbone but specific to the ML task — handle losses, optimizers, and metrics.

Evaluation

• metrics/: Tested suite of efficient, memory-friendly metrics.
• evaluation/: End-to-end scripts to benchmark model predictions and generate plots.

Pipeline

• main_hydra.py: Entry point for training or inference using Hydra configurations.
• docs/archesweather/: Quickstart code for training and inference.

Installation

Clone the repository and install the package with Poetry:

git clone git@github.com:INRIA/geoarches.git
cd geoarches
poetry install

For full installation instructions, see the documentation.

Contributing

The project welcomes contributions and suggestions. If you want to contribute, please read the Contributing Guide.

External Resources

Many thanks to the authors of WeatherLearn for adapting the Pangu-Weather pseudocode to PyTorch. The code for our model is mostly based on their codebase.

• WeatherBench
• WeatherLearn

License

geoarches is available under the BSD 3-Clause License.