How it works

The operations described in this section of the documentation (subsetting, selecting, combining, cutout, etc.) are not independent features bolted onto a reader. They are composable building blocks: every operation takes one or more datasets as input and returns something that is, again, a dataset. This page explains that idea.

Note

These building blocks do not yet have an official name in the documentation. Throughout this page we refer to them as dataset views (by analogy with NumPy views: lazy objects that behave like the data they wrap). Other candidates are operators, adapters, nodes or lenses — suggestions are welcome.

Everything composes

The single most important property of anemoi-datasets is that the result of an operation behaves exactly like a dataset: it has the same methods and attributes (shape, variables, dates, statistics, __getitem__, …) and it loads its data lazily. This is a closure property: operations take datasets and return datasets, so they can be chained without limit.

Because of this, you can build a complex dataset step by step, feeding the result of one open_dataset call into the next:

from anemoi.datasets import open_dataset

# Each call to open_dataset returns an object that behaves like a dataset.
# Because the result is itself a dataset, it can be passed to open_dataset
# again: the operations *compose*.

# 1. Open a dataset stored on disk (a leaf).
ds = open_dataset("aifs-ea-an-oper-0001-mars-o96-1979-2022-6h-v8")

# 2. Subset it in time (wraps the leaf).
ds = open_dataset(ds, start=2000, end=2020)

# 3. Select a few variables (wraps the subset).
ds = open_dataset(ds, select=["2t", "msl", "10u", "10v"])

# 4. Change the frequency (wraps the selection).
ds = open_dataset(ds, frequency="12h")

# The result behaves exactly like the original dataset.
print(ds.shape)
print(ds.variables)
print(ds[0])

Each step wraps the previous one. You do not have to materialise the intermediate results: nothing is read from disk until you index the final object.

The same network can be expressed as a single nested call:

from anemoi.datasets import open_dataset

# The same network of operations as in `compose_steps.py`, but described
# in a single, nested call. The operations are applied from the inside
# out: open -> subset -> select -> change frequency.
ds = open_dataset(
    {
        "dataset": "aifs-ea-an-oper-0001-mars-o96-1979-2022-6h-v8",
        "start": 2000,
        "end": 2020,
        "select": ["2t", "msl", "10u", "10v"],
        "frequency": "12h",
    }
)

print(ds.shape)

The operations can equally be passed as keyword arguments to a single call, without wrapping them in a dictionary:

from anemoi.datasets import open_dataset

# The same network again, but with the operations passed as keyword
# arguments to a single call rather than collected in a dictionary.
ds = open_dataset(
    "aifs-ea-an-oper-0001-mars-o96-1979-2022-6h-v8",
    start=2000,
    end=2020,
    select=["2t", "msl", "10u", "10v"],
    frequency="12h",
)

print(ds.shape)

Subsetting, selecting, joining, concatenating, cutting out, rescaling, filling missing dates, … all follow the same rule, which is why they can be combined freely.

A network of interacting views

When you call open_dataset, the keyword arguments are interpreted one by one. Each recognised operation wraps its input in a new dataset view and returns it, so the arguments build up a network (a tree, or more generally a directed acyclic graph) of objects:

the leaves are the actual zarr stores on disk or in the cloud;
the internal nodes are the views created by start / end, select, frequency, join, cutout, and so on;
the root is the object returned to you.

No data is copied when the network is built. Instead, when you access ds[i], the request travels from the root down to the leaves: each view transforms the index and/or the data it receives from its children and passes the result back up. A time subset shifts indices, a variable selection picks columns, a join stacks the arrays of its children, a cutout merges grid points, and so on. The data is read from the zarr stores only where and when it is actually needed.

You can inspect the network with the tree() method:

from anemoi.datasets import open_dataset

ds = open_dataset(
    {
        "dataset": "aifs-ea-an-oper-0001-mars-o96-1979-2022-6h-v8",
        "start": 2000,
        "end": 2020,
        "select": ["2t", "msl"],
        "frequency": "12h",
    }
)

# `tree()` returns the network of objects that was built. Printing it shows
# how the operations are nested, from the outermost operation down to the
# leaf (the zarr store).
print(ds.tree())

For the example above, the views form a simple chain from the root (the object returned to you) down to the leaf (the store):

Select (select=2t, msl)
└── Subset (frequency=12h, start=2000, end=2020)
    └── GriddedZarr aifs-ea-an-oper-0001-mars-o96-1979-2022-6h-v8

The chain is linear because every operation here wraps a single dataset. As soon as datasets are combined, the tree branches. For instance, the join configuration used later on this page (two datasets combined, the second one reduced to two variables, the whole thing subset in time and frequency) builds this network:

Subset (frequency=12h, start=2000, end=2020)
└── Join
    ├── GriddedZarr aifs-ea-an-oper-0001-mars-o96-1979-2022-6h-v8
    └── Select (select=tp, cp)
        └── GriddedZarr extra-variables-o96-1979-2022-6h-v1

The leaves are always the actual stores; everything above them is a lazy view. Reading ds[i] sends the request down through this tree and assembles the answer on the way back up.

Combining two datasets always makes the tree branch, whatever the operation. Two datasets covering consecutive periods can be concatenated along time:

from anemoi.datasets import open_dataset

# Two datasets covering consecutive periods, concatenated along the time
# axis into a single, longer dataset.
ds = open_dataset(
    concat=[
        "aifs-ea-an-oper-0001-mars-o96-1979-2000-6h-v8",
        "aifs-ea-an-oper-0001-mars-o96-2001-2022-6h-v8",
    ]
)

Concat
├── GriddedZarr aifs-ea-an-oper-0001-mars-o96-1979-2000-6h-v8
└── GriddedZarr aifs-ea-an-oper-0001-mars-o96-2001-2022-6h-v8

and a high-resolution regional dataset can be cut into a global one, so that its grid points replace the global ones where the two overlap:

from anemoi.datasets import open_dataset

# A high-resolution regional dataset cut into a global one: the regional
# grid points replace the global ones wherever the two overlap.
ds = open_dataset(
    cutout=[
        "metno-regional-2p5km-2020-2022-6h-v1",
        "aifs-ea-an-oper-0001-mars-o96-2020-2022-6h-v8",
    ]
)

Cutout
├── GriddedZarr metno-regional-2p5km-2020-2022-6h-v1
└── GriddedZarr aifs-ea-an-oper-0001-mars-o96-2020-2022-6h-v8

Because the whole processing chain is captured by the network of views, the datasets are metadata rich: the metadata() method records every operation that was applied, the sources that were combined, the provenance of the run and the supporting arrays (latitudes, longitudes, etc.).

from anemoi.datasets import open_dataset

ds = open_dataset(
    {
        "dataset": "aifs-ea-an-oper-0001-mars-o96-1979-2022-6h-v8",
        "start": 2000,
        "end": 2020,
        "select": ["2t", "msl"],
        "frequency": "12h",
    }
)

# The metadata records every operation applied to the data, the sources
# that were combined, the provenance of the run, and the supporting arrays
# (latitudes, longitudes, etc.). Anything that reads the dataset can later
# reconstruct exactly how it was produced.
metadata = ds.metadata()

print(sorted(metadata.keys()))
print(metadata["start_date"], metadata["end_date"])
print(metadata["variables"])

This metadata is stored in the model checkpoints at training time, so that the exact data pipeline used to train a model is always recoverable.

Why a single `open_dataset` with a dictionary

The network of views is entirely described by the arguments passed to open_dataset. Because those arguments can be given as a single (possibly deeply nested) dictionary, the whole pipeline can be described as data rather than code:

# A dataset network described entirely as data. Because the description is
# a plain dictionary, it can live in a YAML file and be loaded by a training
# framework such as Hydra, without any Python code changes.
dataset:
  join:
    - aifs-ea-an-oper-0001-mars-o96-1979-2022-6h-v8
    - dataset: extra-variables-o96-1979-2022-6h-v1
      select: [tp, cp]
  start: 2000
  end: 2020
  frequency: 12h

import yaml

from anemoi.datasets import open_dataset

# The whole network of operations can be stored as data (here YAML) and
# passed to open_dataset as a single dictionary. This is what allows a
# training configuration file to fully describe the training data.
with open("compose_config.yaml") as file:
    config = yaml.safe_load(file)

ds = open_dataset(config["dataset"])

print(ds.shape)

This is the reason the package exposes a single open_dataset function driven by a dictionary, rather than a collection of separate functions. A training framework — typically driven by YAML files through Hydra — can let users describe an arbitrarily complex combination of datasets in their configuration file, and hand that configuration straight to open_dataset. The same configuration fully and reproducibly defines the training data.

How it works

Everything composes

A network of interacting views

Why a single open_dataset with a dictionary

Why a single `open_dataset` with a dictionary