Utility Functions

The utilities module provides helper functions and classes for working with anemoi workflows, particularly for managing multi-dimensional expansions.

earthkit.workflows.plugins.anemoi.utils.expansion_qube_from_metadata(metadata: dict[str, Metadata], lead_time: LEAD_TIME) → dict[str, Qube]

Create a Qube object from model metadata and lead time.

This function constructs a Qube object by analysing the model’s metadata to identify surface, pressure level, and model level variables. It creates a hierarchical qube structure organising variables by their vertical coordinate type and expands them across time steps.

Parameters:

metadata (dict[str, Metadata]) – Model metadata containing variable definitions, including their vertical coordinate information (surface, pressure levels, model levels) and the model’s time step. The dict keys are dataset names and the values must be anemoi.inference.metadata.Metadata instances. The forecast lead time as an integer or string (e.g., “7D” for 7 days). This determines the number of time steps in the expansion. If an integer is provided, it is interpreted as hours.
lead_time (LEAD_TIME) – The forecast lead time as an integer or string (e.g., “7D” for 7 days). This determines the number of time steps in the expansion. If an integer is provided, it is interpreted as hours.

Returns:

A dictionary of qube objects with a hierarchical structure containing up to three branches: surface, pressure, and model level variables. Each branch contains the appropriate parameters and coordinates.

Return type:

dict[str, Qube]

Notes

The function creates a qube with the following structure:

Surface variables: expanded over (step, param)
Pressure level variables: expanded over (step, param, level)
Model level variables: expanded over (step, param, level)

Only variables marked as “diagnostic” or “prognostic” that don’t have MARS requests are included. The time steps are calculated from the model’s time step size up to the specified lead time.

Examples

Create expansion coordinates for a 5-day forecast:

>>> from anemoi.inference.checkpoint import Checkpoint
>>> ckpt = Checkpoint("path/to/checkpoint.ckpt")
>>> metadata = ckpt.metadata
>>> qube = expansion_qube(metadata, "5D")
>>> qube.axes()
{'step': {6, 12, 18, ..., 120}, 'param': {...}, 'level': {...}}

Use with an action to expand across all dimensions:

>>> expanded_action = action.expand_as_qube(qube)

See also

Qube: The Qube class for manual qube construction

earthkit.workflows.plugins.anemoi.utils.expansion_qube_from_variables(variables: dict[str, list[str]], variables_metadata: dict[str, dict[str, dict]], model_step: int, lead_time: LEAD_TIME) → dict[str, Qube]

Create a Qube object from a list of variable names, their metadata, model step, and lead time.

This function constructs a Qube object by analysing the provided list of variable names and their corresponding metadata to identify surface, pressure level, and model level variables. It creates a hierarchical qube structure organising variables by their vertical coordinate type and expands them across time steps.

Parameters:

variables (dict[str, list[str]]) – A dictionary mapping dataset names to lists of variable names to include in the qube. These should correspond to keys in the variables_metadata dictionary.
variables_metadata (dict[str, dict[str, dict]]) – A dictionary mapping variable names to their metadata objects. Each metadata object should contain information about whether the variable is a surface level, pressure level, or model level variable, as well as its parameter and level information. Structure: {dataset_name: {variable_name: variable_metadata}}
model_step (int) – The model’s time step in seconds. This is used to calculate the time steps for expansion.
lead_time (LEAD_TIME) – The forecast lead time as an integer or string (e.g., “7D” for 7 days). This determines the number of time steps in the expansion. If an integer is provided, it is interpreted as hours.

Returns:

A dictionary mapping dataset names to qube objects with a hierarchical structure containing up to three branches: surface, pressure, and model level variables. Each branch contains the appropriate parameters and coordinates.

Return type:

dict[str, Qube]

Notes

The function creates a qube with the following structure:

Surface variables: expanded over (step, param)
Pressure level variables: expanded over (step, param, level)
Model level variables: expanded over (step, param, level)

Only variables that are present in the provided list and have corresponding metadata are included. The time steps are calculated from the model’s time step size up to the specified lead time.

Examples

Create expansion coordinates for specific variables:

>>> variables = ["2t", "msl", "10u"]
>>> variables_metadata = {
        "2t": ...,
        "msl": ...,
        "10u": ...,
    }
>>> model_step = 3600  # 1 hour in seconds
>>> lead_time = "5D"
>>> qube = expansion_qube_from_variables(variables, variables_metadata, model_step, lead_time)
>>> qube.axes()
{'step': {6, 12, 18, ..., 120}, 'param': {130, 134, 165}, 'level': {500, 1000}}