Anemoi Packages

Detailed below are brief descriptions of the packages which make up the Anemoi framework.

anemoi-datasets

anemoi-datasets provides the tools to build datasets which are optimised for machine learning training, with appropriate chunking and precomputed statistics for normalisation. These datasets can be built from a range of input sources, including MARS, Grib, NetCDF, Zarr and more. Additionally, a number of filters have been developed which can be used to support data transformations to aid model development. Finally, the package also supports parallel and incremental dataset creation (useful for large datasets).

anemoi-training

anemoi-training provides the functionality to train machine learning models, using pytorch-lightning and Hydra. The training is highly configurable and fully defined through configuration files (utilising anemoi-models to achieve this). The package also includes profiling evaluation, plotting and logging of defined model and system metrics. anemoi-training is designed to work with datasets created using anemoi-datasets.

anemoi-graphs

anemoi-graphs provides the functionality to create complex global or local area graphs. The graph contains nodes and edges that indicate how the grid is structured and how the input, latent and output states are connected. The graph can then be used by the model to identify the neighbouring points of each node in the encoder, processor and decoder communications.

anemoi-models

anemoi-models provides implementations for various type of models. These models are based on a graph encoder-processor-decoder approach and are implemented using the PyTorch library. The model contains the weights of the different layers applied to each node of a graph.

anemoi-inference

anemoi-inference provides the tools to take a trained model and perform the inference/rollout given some initial conditions. Inference makes full use of the metadata stored in a checkpoint to facilitate simple execution without requiring large amounts of boilerplate code.

anemoi-transform

anemoi-transform provides tools for consistent data transformations across training and inference workflows. This package includes common functionality, such as filters, that can be shared and reused across anemoi-datasets, anemoi-training, and anemoi-inference to promote consistency and reduce duplication.

Other anemoi packages

The packages described above also use some functionality implemented in other anemoi subpackages:

anemoi-registry: contains the functionality to save a dataset, a model or an experiment to the anemoi catalogue so that it can be easily shared with others.
anemoi-utils: contains miscellaneous utility functions which are used across the other packages.

Dependencies between packages

The relationship between the various Anemoi packages is illustrated in the figure below. Solid lines indicate mandatory dependencies between packages and dashed lines indicate optional dependencies.