Anemoi dataset support for mllam training

[annnnnetens]

ECMWF is responsible for a Machine Learning for Weather Prediction project called Anemoi. It is more versatile than Neural LAM, but also more convoluted as a result, so each project has its pros and cons. Anemoi-datasets are also based on Zarr files, similar to mllam-datasets, and are batched similarly. Furthermore, there are tens of terabytes of available data to members at https://anemoi.ecmwf.int/, including gridded LAM data such as the CERRA dataset. Because of this, it would be extremely useful to members of this project to support reading Anemoi datasets.

Our goal at LEGMC would be to try training a model with more variables than radar reflectance. The datasets at mllam-data-prep mostly contain just radar data, and we are considering using more NWP data.

An example of an Anemoi-dataset can be found at https://anemoi.readthedocs.io/projects/training/en/latest/user-guide/download-era5-o96.html . The data itself is pretty much ready to be used in Neural LAM, but the metadata has to be changed first.

Another note is that Anemoi-datasets allow more flexible datasets, such as various grids, global data and now also point observations. This is not something that needs to be addressed, it is completely reasonable that such data would just be rejected in mllam.

For starters, when directly trying to use inputs.dataset.variables: - 10u  in MLLAM with an anemoi-dataset, I get an error
KeyError: "Could not find the variable 10u in the dataset. The available variables are ['stdev', 'count', 'dates', 'mean', 'sums', 'has_nans', 'minimum', 'latitudes', 'squares', 'longitudes', 'maximum', 'data']"

[sadamov]

Thanks @annnnnetens! Anemoi support is something we've discussed on and off for a couple of years now, and CERRA plus the wider anemoi data pool would be genuinely valuable for LAM training.

First, on the error: inputs.dataset.variables: - 10u is mllam-data-prep config, and the KeyError is mllam-data-prep failing to find 10u. The reason is that an anemoi-dataset doesn't store fields as separate named arrays. Everything is packed into a single data array (roughly (time, variable, ensemble, cell)), and the variable names plus per-variable statistics (mean, stdev, ...) live in the metadata rather than as named variables. So mllam-data-prep only sees the top-level arrays the error lists (data, mean, latitudes, ...) and can't resolve 10u.

The clean way to support anemoi is a dedicated AnemoiDatastore. neural-lam doesn't depend on mllam-data-prep, it depends on the BaseDatastore interface (base.py, introduced in #66), so a new backend slots in next to MDPDatastore and NpyFilesDatastoreMEPS (init.py) and bypasses mllam-data-prep entirely. anemoi already ships mean/stdev and variable metadata, so get_standardization_dataarray / get_vars_names come almost for free. It could also avoid the duplication MDPDatastore incurs: that one materialises a derived zarr on disk (mdp.py), whereas an anemoi store is already prepared, so an AnemoiDatastore could read it in place and do the category split + standardisation lazily. We'd have to double check this holds against the full contract, but if it does it's great!

The crux: how do we actually read an anemoi store? My understanding is that anemoi-datasets reimplemented their own xarray-like access layer on top of zarr rather than exposing a plain xarray dataset. If anemoi-datasets exposed a proper xarray view (lazy, dask-backed), an AnemoiDatastore becomes almost trivial, and honestly fixing it upstream in anemoi-datasets might be the lower-effort path for everyone. Do you know whether they already expose something like that, or whether there's appetite on the anemoi side for it?

On the flexible-grid point: rejecting global and point-obs data is fine for a first version, but I wouldn't frame non-regular grids as permanently out of scope. We do have plans for irregular grids. Daniel Holmberg et al. have been using such features (https://arxiv.org/abs/2605.15470), and there's relevant ongoing work (e.g. #184, #211, #323, #473,...). So global/obs is coming...

For a first PR I'd keep scope tight: state + forcing on a single regular anemoi grid, standardisation from the built-in stats, read in place where possible. Would LEGMC be up for prototyping this?

[leifdenby]

Thanks for bringing this up @annnnnetens! Yes we have been discussing this off an on, it would be great to have!

@sadamov : The clean way to support anemoi is a dedicated AnemoiDatastore.

I completely agree. One starting point you could consider here is looking at the code we have written in mlwp-data-loaders for loading datasets from anemoi-datasets into a well-defined xr.Dataset-based datastructure: https://github.com/mlwp-tools/mlwp-data-loaders/blob/main/src/mlwp_data_loaders/loaders/anemoi/anemoi_datasets.py. From this you could extract the pieces to satisfy the interface (implement the methods and properties) of the BaseDatastore abstract base class. Or you could just use that code for inspiration.

There is the caveat here though (that I also mentioned on slack):

• there is (still) no specification for zarr datasets written with anemoi-datasets, there has been attempts to create a spec but those efforts have stalled
• the "loader" above in mlwp-data-loader therefore will stop working when what anemoi-datasets writes changes, the "format" has changed multiple times in the past, and I think is likely to change again. It is not even clear to me that all the datasets in the anemoi "registry" all follow the same structure.

You are welcome to use mlwp-data-loaders directly, I am about to release 0.1.0 of that package, and over time I hope to discover what revisions of anemoi-datasets formats exist and adapt the code to handle them.

As I wrote on slack you might also find useful inspiration in https://gist.github.com/MeraX/7c80b20689c999d5088e17327f9eb113#file-grib_to_anemoi_zarr-py.