Data

This module is used to initialise datasets (constructed using anemoi-datasets) and load data into the model. It performs validation checks, such as ensuring that the training dataset end date is before the start date of the validation dataset.

Each training, validation, and test dataset is a map-style dataset: an index selects one sample, which __getitem__ reads from all data readers. The datamodule pairs each dataset with a sampler that decides which samples a rank loads and in which order, and with a dataloader that records its position so that training can resume in the middle of an epoch.

Dataset Architecture

The data module provides two types of dataset readers that wrap anemoi-datasets data:

Native Grid Dataset

The NativeGridDataset class is used for standard atmospheric data on a native grid. It provides a simple interface for reading data samples at specified time indices.

Trajectory Dataset

The TrajectoryDataset class extends NativeGridDataset to support trajectory-based sampling, where data is organized into temporal trajectories. This is useful for tracking atmospheric features over time or for specialized training strategies that require trajectory awareness.

Trajectories are defined by:

• Trajectory start: The reference datetime from which trajectories begin

• Trajectory length: The number of time steps in each trajectory

Each sample in the dataset is associated with a trajectory ID, ensuring that samples are correctly grouped and that trajectory boundaries are respected during training.

Multi-Dataset

The MultiDataset class provides a higher-level wrapper that can synchronize and combine multiple datasets (either NativeGridDataset or TrajectoryDataset instances). This is the primary interface used for training and supports:

• Synchronizing samples across multiple datasets with different grids

• Handling grid sharding for distributed training

Shuffling and the split of samples between groups of ranks are done by the sampler that the datamodule builds, not by the dataset.

Note

Users wishing to change the format of the batch input into the model should sub-class MultiDataset and override the get_sample method.

Data Loading Policy

The DataLoadingPolicy pairs the dataset construction above with the batch sampler that orders, shards and batches its samples, so the two always stay consistent with each other. It is selected via the dataloader.loading_policy config entry. The default, SynchronizedLoadingPolicy, builds one synchronized MultiDataset and splits the samples of a distributed sampler into batches with torchdata’s BatchSampler, which remembers its position so that training can resume in the middle of an epoch.

API Reference

Dataset Readers

class anemoi.training.data.data_reader.BaseAnemoiReader(dataset: str | dict | None = None, dataset_config: str | dict | None = None, start: datetime | int | None = None, end: datetime | int | None = None)

Bases: object

Anemoi data reader for native grid datasets.

A sample is addressed by a 2-D anchor (sequence, position). Relative time offsets are applied within a single sequence, so a sample can never span two sequences. Analysis datasets expose a single sequence (the whole time series); forecast datasets expose one sequence per initialisation (base date), with the forecast step as the position.

default_sampling: dict

Sampling config used by compute_anchors(). {"stride": 1} keeps every valid position; {"stride": None} uses stride = window size (non-overlapping).

property num_sequences: int

Number of independent sequences in the dataset.

sequence_length(sequence: int = 0) → int

Return the number of positions in sequence.

property missing_sequences: set[int]

Return sequences that are entirely missing and must not be sampled.

missing_positions(sequence: int = 0) → set[int]

Return positions within sequence that are missing.

compute_anchors(relative_indices: list[int] | ndarray, sampling: dict | None = None) → ndarray

Return the valid (sequence, position) anchors for a relative window.

Parameters:

• relative_indices (list[int] | np.ndarray) – Relative offsets (in positions) requested around each anchor.

• sampling (dict | None) – Sampling configuration with key "stride". {"stride": None} uses stride = window size (non-overlapping); {"stride": 1} keeps every valid position; {"stride": 6} steps anchors by 6. Defaults to default_sampling.

Returns:

Array of shape (n_anchors, 2) with (sequence, position) rows.

Return type:

np.ndarray

property dates: ndarray

Return dataset dates.

property grid_size: int

Return dataset grid size.

property statistics: dict

Return dataset statistics.

statistics_tendencies(timestep: int | str | timedelta | None = None) → dict | None

Return dataset tendency statistics.

property variables: list[str]

Return dataset variables.

property missing: set[int]

Return dataset missing values mask.

property metadata: dict

Return dataset metadata.

property frequency: timedelta

Return dataset frequency.

property supporting_arrays: dict

Return dataset supporting_arrays.

property name_to_index: dict[str, int]

Return dataset name-to-index mapping.

property resolution: str

Return dataset resolution.

property cutout_mask: ndarray

Return cutout mask.

property boundary_mask: ndarray

Return boundary mask, defined as the complement of the cutout mask.

get_sample(sequence: int, positions: slice | int | list[int] | ndarray, grid_shard_indices: ndarray | slice | None = None) → Tensor

Get a sample from the dataset.

For analysis datasets there is a single sequence, so sequence is ignored and positions index the time axis directly.

class anemoi.training.data.data_reader.NativeGridDataset(dataset: str | dict | None = None, dataset_config: str | dict | None = None, start: datetime | int | None = None, end: datetime | int | None = None, sampling: dict | None = None)

Bases: BaseAnemoiReader

Native grid (analysis) dataset.

A single sequence covering the whole time series; relative offsets index the time axis directly. This is the default analysis behaviour.

class anemoi.training.data.data_reader.TrajectoryDataset(dataset: str | dict | None = None, dataset_config: str | dict | None = None, start: datetime | int | None = None, end: datetime | int | None = None, sampling: dict | None = None)

Bases: BaseAnemoiReader

Trajectory dataset with an explicit lead-step axis.

Wraps a 5-D trajectories-layout dataset opened through anemoi.datasets.open_dataset() (on-disk shape (base_dates, variables, ensembles, steps, cells)). Each base date (forecast initialisation) is exposed as an independent sequence and the forecast step is the within-sequence position, so a training sample is always contained within a single forecast and never crosses initialisation boundaries.

Step subsetting (steps, step_start, step_end, step_frequency) and base-date subsetting (start/end on the valid-time envelope, or base_start/base_end) are handled by open_dataset via the dataset configuration.

property num_sequences: int

Number of forecast initialisations (base dates).

sequence_length(sequence: int = 0) → int

Return the number of forecast steps per initialisation.

property missing_sequences: set[int]

Return the base-date indices that are missing.

missing_positions(sequence: int = 0) → set[int]

Forecast datasets do not track per-step missing values.

property frequency: timedelta

Return the step frequency (spacing between consecutive forecast steps).

get_sample(sequence: int, positions: slice | int | list[int] | ndarray, grid_shard_indices: ndarray | slice | None = None) → Tensor

Load forecast steps positions of initialisation sequence.

anemoi.training.data.data_reader.create_dataset(dataset_config: dict, **_kwargs) → BaseAnemoiReader

Factory function to create a data reader based on the dataset configuration.

Multi-Dataset API

class anemoi.training.data.multidataset.MultiDataset(data_readers: dict[str, BaseAnemoiReader], relative_date_indices: dict[str, slice | int | list[int] | ndarray], fake_dataloading: bool = False)

Bases: Dataset

Multi-dataset wrapper that returns synchronized samples from multiple data readers.

Each index refers to one valid (sequence, position) anchor shared by all data readers. The dataloader’s sampler decides which indices a rank loads and in which order.

set_relative_date_indices(relative_date_indices: dict[str, slice | int | list[int] | ndarray]) → None

Set the time steps loaded for each sample, and the samples that provide them.

Call this before the dataloader is created: the number of samples changes with the time steps, and the sampler reads it when it is constructed.

property statistics: dict[str, dict]

Return combined statistics from all data readers.

property metadata: dict[str, dict]

Return combined metadata from all data readers.

property supporting_arrays: dict[str, dict]

Return combined supporting arrays from all data readers.

property variables: dict[str, list[str]]

Return combined variables from all data readers.

property data: dict

Return data from all data readers as dictionary.

property name_to_index: dict[str, dict]

Return combined name_to_index mapping from all data readers.

property resolution: dict[str, str]

Return combined resolution from all data readers.

property frequency: timedelta

Return combined frequency from all data readers.

set_reader_group_info(reader_group_rank: int, shard_sizes: dict[str, list[int] | None]) → None

Set the reader group information, called by the DDP strategies.

Parameters:

• reader_group_rank (int) – Reader group rank

• shard_sizes (dict[str, ShardSizes]) – Shard sizes for all datasets

Data Loading Policy API

class anemoi.training.data.loading.DataLoadingPolicy

Bases: ABC

Builds the dataset and batch sampler consumed by one dataloader stage.

A policy owns both the dataset layout (e.g. one synchronized dataset vs. several independently indexed ones) and the batch sampler that orders, shards and batches its samples, so the two always stay consistent with each other.

abstractmethod build_dataset(data_readers: dict[str, BaseAnemoiReader], relative_date_indices: dict[str, slice | int | list[int] | ndarray], *, fake_dataloading: bool = False) → Dataset

Build the dataset for one dataloader stage.

Parameters:

• data_readers (dict[str, BaseAnemoiReader]) – Dictionary mapping dataset names to their data readers.

• relative_date_indices (dict[str, TimeIndices]) – Precomputed relative date indices for each data reader.

• fake_dataloading (bool, optional) – Load one real sample and reuse it for subsequent accesses, by default False.

Returns:

Map-style dataset that returns one sample for each index.

Return type:

Dataset

abstractmethod build_batch_sampler(dataset: Dataset, *, num_replicas: int, rank: int, shuffle: bool, seed: int, batch_size: int, drop_last: bool) → BatchSampler

Build the batch sampler that orders, shards and batches the dataset’s samples.

The returned batch sampler must meet two conditions:

• Its sampler attribute has a set_epoch(epoch) method that selects the shuffle of an epoch. The datamodule calls it when it builds a dataloader, and Lightning calls it at the start of every epoch.

• Its batches depend only on the seed and the epoch. A resumed run restores the position saved by torchdata’s stateful BatchSampler, or, if the batch sampler keeps no state, replays the batches of the epoch and skips the ones already trained on. Both give the same remaining batches only if the order is reproducible.

Parameters:

• dataset (Dataset) – Dataset built by build_dataset() for the same stage.

• num_replicas (int) – Number of sample groups the dataset is split across.

• rank (int) – Sample group this batch sampler produces batches for.

• shuffle (bool) – Shuffle the samples.

• seed (int) – Seed for the shuffle, shared by every rank of a sample group.

• batch_size (int) – Number of samples per batch.

• drop_last (bool) – Drop the remaining samples if they don’t divide evenly across sample groups.

Returns:

Batch sampler that hands out lists of sample indices for this sample group.

Return type:

BatchSampler

class anemoi.training.data.loading.SynchronizedLoadingPolicy

Bases: DataLoadingPolicy

Default policy: one synchronized dataset and a distributed sampler split into batches.

Every sample contains data from every reader, and the samplers of a sample group all hand out the same indices in the same (shuffled) order.

build_dataset(data_readers: dict[str, BaseAnemoiReader], relative_date_indices: dict[str, slice | int | list[int] | ndarray], *, fake_dataloading: bool = False) → MultiDataset

Build the dataset for one dataloader stage.

Parameters:

• data_readers (dict[str, BaseAnemoiReader]) – Dictionary mapping dataset names to their data readers.

• relative_date_indices (dict[str, TimeIndices]) – Precomputed relative date indices for each data reader.

• fake_dataloading (bool, optional) – Load one real sample and reuse it for subsequent accesses, by default False.

Returns:

Map-style dataset that returns one sample for each index.

Return type:

Dataset

build_batch_sampler(dataset: Dataset, *, num_replicas: int, rank: int, shuffle: bool, seed: int, batch_size: int, drop_last: bool) → BatchSampler

Build the batch sampler that orders, shards and batches the dataset’s samples.

The returned batch sampler must meet two conditions:

• Its sampler attribute has a set_epoch(epoch) method that selects the shuffle of an epoch. The datamodule calls it when it builds a dataloader, and Lightning calls it at the start of every epoch.

• Its batches depend only on the seed and the epoch. A resumed run restores the position saved by torchdata’s stateful BatchSampler, or, if the batch sampler keeps no state, replays the batches of the epoch and skips the ones already trained on. Both give the same remaining batches only if the order is reproducible.

Parameters:

• dataset (Dataset) – Dataset built by build_dataset() for the same stage.

• num_replicas (int) – Number of sample groups the dataset is split across.

• rank (int) – Sample group this batch sampler produces batches for.

• shuffle (bool) – Shuffle the samples.

• seed (int) – Seed for the shuffle, shared by every rank of a sample group.

• batch_size (int) – Number of samples per batch.

• drop_last (bool) – Drop the remaining samples if they don’t divide evenly across sample groups.

Returns:

Batch sampler that hands out lists of sample indices for this sample group.

Return type:

BatchSampler