Machine learning framework for geospatial data, mainly climate and weather.
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.
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.
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.
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.
metrics/: Tested suite of efficient, memory-friendly metrics.evaluation/: End-to-end scripts to benchmark model predictions and generate plots.
main_hydra.py: Entry point for training or inference using Hydra configurations.docs/archesweather/: Quickstart code for training and inference.
Clone the repository and install the package with Poetry:
git clone [email protected]:INRIA/geoarches.git
cd geoarches
poetry installNote
By default, Poetry will install geoarches in editable mode. This allows you to make changes to the package locally, meaning any local changes will automatically be reflected to the code in your environment.
For full installation instructions, see the documentation.
The project welcomes contributions and suggestions. If you want to contribute, please read the Contributing Guide.
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.
geoarches is available under the BSD 3-Clause License.