Anemoi Core

Anemoi Training is a framework designed for training machine learning models specifically for weather forecasting. It is built to be flexible and customizable, allowing users to adjust model architectures and training processes via configuration without modifying the underlying source code.

Supported model architectures include:

• Graph Neural Networks (GNNs)
• Graph Transformers
• Transformers with Flash Attention

Anemoi Core is a monorepo providing core training and modelling functionality for Anemoi. It is composed of three primary packages:

• anemoi-training: Miscellaneous tools for training data-driven weather forecasts.
• anemoi-models: Core model components used to build the architecture of data-driven NWP (Numerical Weather Prediction) models.
• anemoi-graphs: Tools for building graphs for data-driven forecasts.

The anemoi.graphs.inspect module provides tools to analyze and visualize graph structures. When using this module, you can perform the following inspection tasks:

• Attribute Analysis: Generate distribution plots for both node and edge attributes.
• Subgraph Visualization: Create interactive plots of individual subgraphs.
• Isolation Analysis: Create interactive plots of isolated nodes.
• Console Summaries: Obtain a text-based description of the graph structure directly in the console.

A task defines the temporal I/O structure of a training sample. It determines which time steps are loaded as model inputs and which are used as prediction targets. Tasks are independent of the model architecture and the training method. They are configured under the task key in your configuration and are defined in anemoi.training.tasks.

All tasks inherit from BaseTask, which provides the interface for the training loop via:

• get_inputs() / get_targets(): Slices the loaded batch into model inputs and targets.
• steps(): An iterable of per-step keyword dicts consumed by the training loop (one dict = one forward pass).
• advance_input(): Updates model inputs between rollout steps (relevant for multi-step tasks).

In Anemoi, there is a distinction between scores that use a closed neighbourhood and those that use an open neighbourhood:

1. Graph Scores (GraphEnergyScoreLoss, GraphEdgeEnergyScoreLoss, GraphVariogramScoreLoss): These use the edges as defined in your graph configuration. If you include self-edges (e.g., via KNNEdges), they are treated as part of the neighbourhood.
2. Edge Scores (GraphEdgeCRPSLoss, GraphEdgeEnergyScoreLoss, GraphVariogramScoreLoss): These use an open neighbourhood. They automatically remove self-edges once during the loss graph construction (before row normalization), allowing you to reuse a self-containing graph configuration.

For GraphEdgeCRPSLoss, use the alpha parameter to select the estimator: 1 for fair CRPS, 0 for standard CRPS, and intermediate values for almost-fair CRPS.

The MultiDataset class is the primary high-level interface used for training. It wraps one or more NativeGridDataset or TrajectoryDataset instances to provide:

• Synchronization: Aligning samples across multiple datasets that may use different grids.
• Distributed Loading: Managing data loading across workers and communication groups.
• Training Logic: Handling shuffling, batching, and grid sharding for distributed training.

Customizing Batch Formats: If you need to change the format of the batch input sent to the model, you should subclass MultiDataset and override either the __iter__ method or the get_sample method.

Accuracy results are pushed to the anemoi integration-test data server under a separate accuracy tree to prevent mixing with performance benchmarks.

• Default behavior: Results are written to <path>/<kind>/<test_case>/.
• Accuracy path: For the global accuracy test, this is typically .../training/accuracy/global/.
• Benchmark path: Performance benchmarks are typically written to .../training/benchmarks/<test_case>/.
• Custom test cases: When adding new accuracy test cases (e.g., a limited-area variant), provide a unique test_case string so results land in a sibling directory (e.g., .../training/accuracy/lam/).

Note: The push to the server only occurs when running from the main branch and is controlled by the track_accuracy_final_loss flag (using the track_accuracy_result helper in src/anemoi/training/diagnostics/benchmark_server.py).

Since anemoi-training=0.4, the framework supports training with multiple datasets simultaneously. This is useful for downscaling or improving model generalisation by leveraging diverse data sources.

Core Architecture:

• Encoders/Decoders: Each dataset has its own dedicated encoder and decoder, allowing for dataset-specific input and output representations (e.g., different variable sets).
• Latent Space: Encoded representations from all datasets are combined into a single shared latent space.
• Processor: A single shared processor operates on this combined latent representation.

Constraint: All datasets must share the same time resolution and forecast horizon (or target times).

Anemoi's training framework for Graph Neural Networks (GNNs) is built on three primary components:

1. BaseTrainingModule: An abstract base class (inheriting from pytorch_lightning.LightningModule) that encapsulates shared logic for training, evaluation, distributed execution, normalization, and masking. It is not intended to be instantiated directly.
2. Tasks: Task-specific subclasses (e.g., Forecaster, TemporalDownscaler, Autoencoder) that define the temporal I/O structure and how time steps are mapped to inputs and targets.
3. AnemoiTrainer: The orchestrator that wraps a PyTorch Lightning Trainer to manage training/validation loops, distributed execution, dataset loading, and graph creation.

The framework is organized into five primary functional modules:

• Data Module: Manages data loading, preprocessing, and routing by interfacing with Anemoi Datasets.
• Training Module: Orchestrates the core training loop, including loss calculation, optimization, and learning rate scheduling.
• Loss Module: Implements various loss functions and manages model optimization.
• Diagnostics Module: Handles experiment tracking, metric logging, and visualization of training progress.
• Strategy Module: Implements training strategies, such as distributed training and advanced training techniques.

The TimeAggregateLossWrapper applies a base loss function to time-aggregated versions of predictions and targets to encourage temporal coherence. This is particularly useful for temporal downscalers.

Supported aggregations:

• mean: Penalizes bias in the temporal average.
• max: Penalizes errors in peak values.
• min: Penalizes errors in minimum values.
• diff: Penalizes unrealistic temporal transitions (consecutive step-to-step differences: pred[:, 1:] - pred[:, :-1]).

Important: The time_steps scaler must be excluded from the inner loss_fn because temporal aggregation collapses the time dimension. Only spatial and variable scalers should be listed in the inner loss configuration.

Because Python objects are stored by reference in pickle files, moving or removing classes can prevent old checkpoints from loading. To fix these import errors, use a migrate_setup callback. This function runs before the checkpoint is loaded and uses a MigrationContext to redirect or delete attributes/modules.

To generate a script template that includes the setup callback, use the --with-setup flag.

MigrationContext Methods:

• context.move_attribute(start_path, end_path): Indicates an attribute moved from start_path to end_path.
• context.move_module(start_path, end_path): Indicates a module moved from start_path to end_path.
• context.delete_attribute(path): Removes an attribute. Supports the wildcard * to delete all attributes in a module.

anemoi-models migration create migration-name --with-setup