Utility Functions

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

class earthkit.workflows.plugins.anemoi.utils.Expansion(qube: Qube)

Bases: object

Utilise Qubed to expand actions along the qube’s axes.

The Expansion class provides a mechanism for expanding workflow actions across multiple dimensions using the Qube data structure.

The expansion process:

Processes each dimension in the qube structure
Expands actions along each axis sequentially
Splits actions into separate branches when multiple child dimensions exist
Creates a hierarchical structure organised by metadata names

Parameters:: qube (Qube) – A Qube object representing the multi-dimensional structure to expand over. The qube defines the axes (dimensions) and their values that will be used to expand the action.

Examples

Create an expansion for surface and pressure level variables:

>>> from qubed import Qube
>>> from earthkit.workflows.plugins.anemoi.utils import Expansion
>>>
>>> # Create surface variables qube
>>> surface = Qube.from_datacube({
...     "step": [6, 12, 18],
...     "param": ["2t", "10u", "10v"]
... })
>>> surface.add_metadata({"name": "surface"})
>>>
>>> # Create pressure level variables qube
>>> pressure = Qube.from_datacube({
...     "step": [6, 12, 18],
...     "param": ["t", "u", "v"],
...     "level": [500, 850, 1000]
... })
>>> pressure.add_metadata({"name": "pressure"})
>>>
>>> # Combine into hierarchical structure
>>> combined = surface | pressure
>>> expansion = Expansion(combined)
>>>
>>> # Expand an action
>>> expanded_action = expansion.expand(action)

See also

expansion_coordinates: Create an Expansion from model metadata

axes() → dict[str, set[str]]

Return a dictionary of all the axes and their values in the qube.

Returns:: Dictionary mapping axis names to sets of values for that axis.
Return type:: dict[str, set[str]]

Examples

>>> qube = Qube.from_datacube({
...     "step": [6, 12],
...     "param": ["t", "q"]
... })
>>> expansion = Expansion(qube)
>>> expansion.axes()
{'step': {6, 12}, 'param': {'t', 'q'}}

drop_axis(axis: str | list[str]) → Self

Drop one or more axes from the qube.

This creates a new Expansion instance with the specified axis/axes removed from the qube structure. The original Expansion is not modified.

Parameters:: axis (str | list[str]) – The name(s) of the axis/axes to drop from the qube.
Returns:: A new Expansion instance with the specified axis/axes removed.
Return type:: Self

Examples

Drop a single axis:

>>> qube = Qube.from_datacube({
...     "step": [6, 12],
...     "param": ["t", "q"],
...     "level": [500, 850]
... })
>>> expansion = Expansion(qube)
>>> new_expansion = expansion.drop_axis("level")
>>> "level" in new_expansion.axes()
False

Drop multiple axes:

>>> new_expansion = expansion.drop_axis(["step", "level"])
>>> list(new_expansion.axes().keys())
['param']

expand(action: Action) → Action

Expand the action according to the qube structure.

This method recursively expands an action across all dimensions defined in the qube. For qubes with a single child, it expands sequentially through the dimensions. For qubes with multiple children, it splits the action into separate branches, each named according to the child’s metadata (if present) or using alphabetical naming as a fallback.

Parameters:: action (Action) – The workflow action to expand across the qube’s dimensions.
Returns:: The expanded action with all dimensions applied. The action will have a hierarchical structure if the qube has multiple children.
Return type:: Action

Notes

The expansion algorithm works as follows:

If the qube has no children, the action is returned unchanged
The action is expanded along the first child’s dimensions
If multiple children exist, the action is split into branches
Each branch is recursively expanded with its child’s dimensions
Child branches are named using metadata or alphabetical fallback

Examples

Simple single-dimension expansion:

>>> qube = Qube.from_datacube({"step": [6, 12, 18]})
>>> expansion = Expansion(qube)
>>> expanded = expansion.expand(action)
# Action is now expanded over step=[6, 12, 18]

Hierarchical expansion with surface and pressure levels:

>>> # Qube structure:
>>> # root, step=6/12
>>> # ├── param=2t/10u/10v (surface)
>>> # └── param=t/u/v, level=500/850/1000 (pressure)
>>>
>>> expansion = Expansion(qube)
>>> expanded = expansion.expand(action)
# Action is expanded over step, then split into /surface and /pressure
# branches, each with their respective param and level dimensions

Drop an axis before expansion:

>>> expansion = Expansion(qube)
>>> no_step = expansion.drop_axis("step")
>>> expanded = no_step.expand(action)
# Action expanded over remaining dimensions only

See also

drop_axis: Remove axes before expansion
axes: View available axes in the qube

earthkit.workflows.plugins.anemoi.utils.expansion_coordinates(metadata: Metadata, lead_time: int | str | timedelta) → Expansion

Create an Expansion object from model metadata and lead time.

This function constructs an Expansion 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 (Metadata) – Model metadata containing variable definitions, including their vertical coordinate information (surface, pressure levels, model levels) and the model’s time step.
lead_time (int | str | timedelta) – 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:

An Expansion object with a hierarchical qube structure containing three branches: surface, pressure, and model level variables. Each branch contains the appropriate parameters and coordinates.

Return type:

Expansion

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
>>> expansion = expansion_coordinates(metadata, "5D")
>>> expansion.axes()
{'step': {6, 12, 18, ..., 120}, 'param': {...}, 'level': {...}}

Use with an action to expand across all dimensions:

>>> expanded_action = expansion.expand(action)