Proposal: Move xarray loaders to earthkit-data

[ninaburg]

Hi,

I’m Nina (Pilot project on Emerging Technologies). Part of our work is to help bring earthkit-data into wider use, and I would like to contribute here.

In anemoi-datasets, several sources still use the custom xarray_support layer. Since earthkit-data now handles NetCDF, Zarr, OPeNDAP, kerchunk, etc., I’d like to start replacing code that can be cleanly migrated.

Likely candidates:

• create/sources/xarray.py
• create/sources/xarray_zarr.py
• create/sources/xarray_kerchunk.py
• create/sources/opendap.py
• create/sources/netcdf.py
• The whole xarray_support/ directory (fieldlist, variable, coordinates, flavour, patch)

I already tested a small prototype of a NetCDF loader using earthkit-data and would be happy to open a PR if this direction makes sense.

Before proceeding:

• Is this migration aligned with your plans, and are there parts of xarray_support that should remain?

Thanks,
Nina

@frazane adding you here so you're aware of this development.

[khintz]

Hi Nina,

I’d like to add my perspective to this. This is a very relevant discussion, and you deserve real credit for raising the question early rather than later.

One should be very careful about moving more functionality away from xarray towards earthkit-data, on the contrary, I would like to see the opposite happen. xarray is the de facto community standard, and most of the meteorological/climate ecosystem is built around xarray data structures (MetPy, Pangeo, Kerchunk, Intake, Zarr, etc).
Reducing xarray usage risks isolating Anemoi, making it harder for external users to adopt and engage in its development.

At least as of today, Earthkit-data is not a community-driven library, at least not at the scale of xarray. xarray solves much of what is needed already, such as field lists, coordinates, metadata propagation, lazy loading, Dask integration.

If one of the goals of Anemoi-datasets is to provide datasets/formats and data loaders that are easy to re-use in different ML pipelines (including non-ECMWF ones), which I don't know if that's the case, keeping xarray as the core common type ensures that the data is directly consumable by the majority of ML workflows in Europe.

I think it would be preferable to avoid removing the existing xarray_support layer unless there is a clear, community-driven reason. Earthkit-data can, of course, be a backend option, but, in my opinion, xarray should be the primary data model.

[observingClouds]

I would like to see ECMWF adapt projects like xarray more widely and actually increase their dependence on this global-community driven and established code-bases rather than developing something that already exists. It would be great if anemoi could completely drop the dependency on earthkit-data. To load data from various stores (file-based, memory, bucket, compressed, huggingface,...) there is a widely adopted framework to do so and those are fsspec-drivers. Drivers that are missing are those for the special ECMWF stores, like FDB, but even for those the community starts to create drivers, e.g. for ECFS.

Any development time into earthkit-data seems therefore better invested into fsspec-drivers that can be much more generalized and are widely accepted by the community.

[cosunae]

I think the context of this discussion was to have a better interplay between earthkit data and anemoi-datasets and discuss where we should use earthkit instead of implementing duplications in anemoi-datasets. I think the proposal to replace all supported array backends by xarray changes functionality supported, so I believe would be best discussing in a separate issue.
In short, I think the earthkit framework already supports fully the xarray ecosystem. Enforcing it as suggested would break many operational workflows and requirements.
I am happy to reflect them (as it was already discussed in the past) in a separate issue if you would like to open one.

[leifdenby]

I think it was good to discuss this at the ATS meeting today (and I look forward to discussing this more at the MLPP meeting next week), and if I understand it correctly what we are discussing here is what implementation to use for loading data into anemoi-datasets. Switching to using earthkit-data to doing this makes me nervous, both because

• earthkit-data actually just proxies calls to xarray or directly to pandas/cfgrib, rather than encouraging users to simply use xarray directly. This adds another layer of indirection where debugging now involves working out if something is wrong in anemoi-datasets, earthkit-data or xarray for example.
• but more fundamentally, I am concerned because of implementation and design of the data structure earthkit-data loads multi-dimensional arrays into. If you compare the definition of the Field class in earthkit-data https://github.com/ecmwf/earthkit-data/blob/develop/src/earthkit/data/core/fieldlist.py#L91 (which I assume is for representing a single multi-dimensional array) with DataArray in xarray https://github.com/pydata/xarray/blob/main/xarray/core/dataarray.py#L257, the former has no docstring defined, inherits from an opaque class called Base and for example has a number of properties tied to MARS infrastructure (e.g. mars_grid, mars_area). It is very hard to reason about what a Field is, what operations can be done on them (can two Fieldss be added together for example?) and it does not feel like a general-purpose data structure for geospatial data. To me the key to working with geospatial data is the datastructure that is implemented to represent it, and I am concerned that using a Field (and FieldList for that matter, https://github.com/ecmwf/earthkit-data/blob/develop/src/earthkit/data/core/fieldlist.py#L995) could lead to a lot of confusing debugging.

But, this all speaks to a more general issue (that I appreciate may not have been the purpose of the proposition here @cosunae, but it is a related issue): both anemoi-datasets and to my mind earthkit-data implement their own multidimensional array data-structures with their own APIs for how to operate on these. I do think that we could save a lot of time and make much better tools if we all build on the existing implementation in xarray instread. I raised over a year ago that anemoi-datasets could have at minimum have used a similar API to xarray (I'll attach a side-by-side comparison that I made here, anemoi_datasets_vs_xarray.py.txt), and I also built an alternative tool that tries to do the same operations as anemoi-datasets, but leveraging the existing xarray API (and uses xr.Dataset and xr.DataArray objects internally for multi-dimensional arrays) rather than implementing things from scratch: https://github.com/mllam/mllam-data-prep. I am not suggesting this should be the tool going forward, but something along those lines would be easier to develop and maintain I think. Using xarray's xr.Dataset/xr.DataArray data structures will also give us things like GPU-native computation (using the same xarray API) for free https://xarray.dev/blog/xarray-kvikio, https://github.com/pangeo-data/ncar-hackathon-xarray-on-gpus

[cosunae]

The questions raised deserve discussion and it might be related, but is not the same. The discussion we want trigger is not whether we switch to earthkit data for loading the data. anemoi-datasets is already doing that. And the functionality that was implemented in anemoi-dataset (replicated code now), was done in a similar fashion.
The intention here is to remove duplicated code to better maintain the ecosystem.