Tasks

The anemoi.training.tasks module defines the temporal structure of a training sample independently of the Lightning module. Each task specifies which time steps are loaded as inputs and which are used as targets, referred as offsets as they are defined as relative positions in time compared to a reference point, and provides helpers for mapping those offsets to batch positions.

This separation lets plotting callbacks, dataloaders, and the trainer remain task-agnostic — they call task.get_batch_input_indices and task.get_batch_output_indices without knowing anything about the specific workflow.

BaseTask

BaseTask is the abstract root of the hierarchy. It is constructed from two lists of timedelta objects:

input_offsets — time offsets of the model inputs relative to the analysis time, e.g. [-6H, 0H] for a two-step input.
output_offsets — time offsets of the model targets, e.g. [6H] for a single-step forecast.

The union of these two lists (task._offsets) tells the datamodule which time steps to load for each sample.

Key properties and methods:

num_input_timesteps / num_output_timesteps — lengths of the offset lists.
steps — iterable of per-step dicts (e.g. {"rollout_step": 0}); the training loop iterates over these.
get_batch_input_indices(**kwargs) — positions of input offsets within the full batch tensor’s time dimension.
get_batch_output_indices(**kwargs) — positions of output offsets within the full batch tensor’s time dimension.
get_inputs / get_targets — extract and index-select the appropriate slices from a batch dict.
get_input_offsets() / get_output_offsets() / get_offsets() — return the input, output, and full list of time offsets, respectively. Used by the datamodule.

class anemoi.training.tasks.base.BaseTask(input_offsets: list[timedelta], output_offsets: list[timedelta])

Bases: ABC

Base class for all tasks.

Tasks define the temporal structure of a training sample by specifying input and output time offsets as timedelta objects. The base class provides:

An _offsets property that is the union of input and output offsets, used by the datamodule to determine which time steps to load.
get_inputs / get_targets to split a loaded batch into model inputs and targets based on position within _offsets.

name

Name of the task.

Type:: str

num_inputs

Number of input time steps for the task.

Type:: int

num_outputs

Number of output time steps for the task.

Type:: int

get_input_offsets() → list[datetime.timedelta]: Get the list of input time offsets.

get_output_offsets(**kwargs) → list[datetime.timedelta]: Get the list of output time offsets.

get_offsets(**kwargs) → list[datetime.timedelta]: Get the full list of input and output time offsets.

steps(mode: str = 'training') → Iterable[dict]: Get the steps for the task.

property num_input_timesteps: int: Number of input time steps.

property num_output_timesteps: int: Number of output time steps.

property num_steps: int: Number of training steps (rollout length).

get_metric_name(**_step_kwargs) → str: Get the metric name for the current step (if any).

get_input_offsets(**_kwargs) → list[timedelta]: Get the list of input time offsets.

get_output_offsets(**_kwargs) → list[timedelta]

Return the output offsets for a given step.

The default implementation returns self._output_offsets. Subclasses may override this to shift outputs per rollout step.

get_offsets(**_kwargs) → list[timedelta]

Get the list of offsets for a given mode (e.g. “training”, “validation”, “test”).

By default, this returns self._offsets, but can be overridden by subclasses to return different offsets per mode for example (e.g different rollout in training vs validation).

get_batch_input_indices(**kwargs) → list[int]: Positions of the input offsets within the full batch _offsets.

get_batch_output_indices(**kwargs) → list[int]

Positions of the output offsets within the full batch _offsets.

Parameters are forwarded to get_output_offset so that subclasses can parametrise the output selection (e.g. per rollout step).

get_inputs(batch: dict[str, Tensor], data_indices: dict[str, IndexCollection], **_kwargs) → dict[str, Tensor]

Extract model inputs from a batch.

Parameters:

batch (dict[str, torch.Tensor]) – Full batch keyed by dataset name.
data_indices (dict[str, IndexCollection]) – Data indices per dataset.

Returns:

Input tensors per dataset with shape (bs, num_inputs, grid, nvar).

Return type:

dict[str, torch.Tensor]

get_targets(batch: dict[str, Tensor], **kwargs) → dict[str, Tensor]

Extract model targets from a batch.

Parameters:

batch (dict[str, torch.Tensor]) – Full batch keyed by dataset name.
**kwargs – Forwarded to get_batch_output_indices (e.g. rollout_step).

Returns:

Target tensors per dataset with shape (bs, num_outputs, ensemble, grid, full_nvar) in DATA_FULL variable space (all variables including forcings).

Return type:

dict[str, torch.Tensor]

log_extra(*_args, **_kwargs) → None: Hook to log any task-specific information.

training_runtime_state_dict() → dict

Return training runtime state to be persisted in the training checkpoint.

Override in subclasses to include any mutable state that accumulates during training and must survive a job resume (e.g. curriculum counters). This state is stored at the training checkpoint level, not in the model state_dict, so that it is invisible to inference code.

load_training_runtime_state_dict(state: dict) → None: Restore training runtime state from a training checkpoint.

on_train_epoch_end(current_epoch: int) → None: Hook to update task state at the end of each training epoch (e.g. for curriculum learning).

fill_metadata(md_dict: dict) → None: Fill the metadata dictionary with task-specific information.

class anemoi.training.tasks.base.BaseSingleStepTask(input_offsets: list[timedelta], output_offsets: list[timedelta])

Bases: BaseTask

Base class for single-step tasks.

advance_input(*args, **_kwargs) → dict[str, Tensor]

Advance the input state for each dataset based on the task’s requirements.

This method can be overridden by specific tasks to implement custom logic for advancing the input state.

Returns:: dict[str, torch.Tensor]
Return type:: The advanced input state for each dataset.

BaseSingleStepTask

BaseSingleStepTask is a convenience subclass for tasks with a single training step (no rollout). Both TemporalDownscaler and BaseTimelessTask inherit from it.

Forecaster

Forecaster implements autoregressive rollout training. It is constructed with:

multistep_input — number of input time steps (e.g. 2 for [-6H, 0H]).
multistep_output — number of output time steps per rollout step (e.g. 1).
timestep — the model timestep as a frequency string (e.g. "6H").
rollout — optional dict configuring the rollout schedule (see RolloutConfig).
validation_rollout — number of rollout steps used during validation (default 1).

RolloutConfig

RolloutConfig encapsulates the progressive rollout schedule:

start — initial number of rollout steps at epoch 0.
epoch_increment — increase the rollout window by one every this many epochs (0 disables progression).
maximum — the rollout window is never increased beyond this value.

The current step count is stored in rollout.step and is increased by calling rollout.increase(), which is triggered by the trainer at the end of each epoch via on_train_epoch_end.

Multistep Input and Output

The forecaster task uses multistep_input and multistep_output to control how many time steps the model ingests as input and predicts in a single forward pass.

multistep_input: number of past timesteps provided as model input. When set to 1, only t_{0} is used.
multistep_output: number of future timesteps predicted per forward pass.

Set multistep_output greater than 1 to enable multi-output prediction. This reduces the number of forward passes needed to cover a rollout horizon.

Example:

task:
  _target_: anemoi.training.tasks.Forecaster
  multistep_input: 3
  multistep_output: 2
  timestep: "6H"
  rollout:
    start: 1
    epoch_increment: 1
    maximum: 6

Rollout behavior:

When time indices are inferred, the dataloader uses multistep_input + rollout * multistep_output to determine how many timesteps to load.
If multistep_output is greater than multistep_input, only the most recent multistep_input outputs are fed into the next rollout step.

TemporalDownscaler

TemporalDownscaler downscales to higher temporal resolution by generating intermediate time steps between two input times. It is constructed with:

input_timestep — coarse time resolution (e.g. "6H").
output_timestep — target fine resolution (e.g. "3H"). Must evenly divide input_timestep.
output_left_boundary — if True, include the t=0 frame in the output targets (default False).
output_right_boundary — if True, include the final t=input_timestep frame in the output targets (default False).

Example: input_timestep="6H", output_timestep="3H", output_left_boundary=True produces output offsets [0H, 3H] and input offsets [0H, 6H].

The default is to use the time aggregate loss when training any temporal downscaler.

Autoencoder

Autoencoder is a timeless task: both input and output are a single snapshot at t=0. It inherits from BaseTimelessTask which itself inherits from BaseSingleStepTask.

class anemoi.training.tasks.timeless.BaseTimelessTask(**_kwargs)

Bases: BaseSingleStepTask

Base class for timeless tasks.

Both input and output are a single snapshot at t=0.

class anemoi.training.tasks.timeless.Autoencoder(**_kwargs)

Bases: BaseTimelessTask

Autoencoding task implementation.