Support distribution of `xarray`

[ClaudiaComito]

See https://docs.xarray.dev/en/stable/

If I understand correctly, an xarray object is made up of the actual data array (np.ndarray), and 1-D coordinates arrays (dictionaries?) that map data dimensions and indices to meaningful physical quantities.

For example, if xarray is a matrix of coordinates (date, temperature), users will be able to perform

mean_temp = xarray['2010_01_01', '2010_12_31'].mean()

Feature functionality

Enable distribution of xarray object, allow named dimensions, keep track of coordinates arrays, one of which will be distributed.

Example, :

ht_xarray = ht.array(xarray, split="date")
ht_mean_temp = ht_xarray['2010_01_01', '2010_12_31'].mean() # memory-distr. operation

Check out Pytorch's named tensors functionality.

Additional context
Initiating collaboration with N. Koldunov @koldunovn at Alfred Wegener Institute (Helmholtz centre for polar and marine research).
Also interesting for @kleinert-f, @ben-bou
Tagging @bhagemeier for help with implementation.

[koldunovn]

Coordinates are not necessarily 1D arrays. For example for curvilinear grids covering the Earth's surface one would have to describe positions of (e.g. centers) of grid points as 2D arrays (one for lon and one for lat), and coordinate arrays then will be 2D.

[ClaudiaComito]

Got it, thanks @koldunovn, edited original post.

[ClaudiaComito]

Implementation discussed during devs meeting Sept 28. Preliminary scheme:

• introduce dndarray.coordinates attribute (default=None)
• ht.array() to accept xarray input, interpret string split, populate dndarray.coordinates, chunk data AND split coordinate
• every operation needs to check if not coordinates: or similar, quantify and minimize performance hit

@Markus-Goetz would be good if you chime in

[ClaudiaComito]

tagging @Mystic-Slice as they have shown interest in working on this 👋

[ClaudiaComito]

Interesting discussion over at array-api about single-node parallelism, but xarray also mentioned in distributed execution context.

@TomNicholas you might be interested in this effort as well.

[Mystic-Slice]

Hi everyone!
I have been thinking about this project for some time now and here are my initial thoughts.

DXarray:

Xarray uses np.ndarray for its data storage. And since numpy does not support computation on GPU, the xarray objects cannot be directly chunked and stored within the new DXarray object.

There are two solutions:

1. Using CuPy arrays within xarray. see (cupy-xarray)
   • As far as I know, this must be very simple and straightforward to use.
   • This sounds like the best way to get most if not all the functionalities of xarray without having to re-implement.
   • My only problem with CuPy is that, it is strictly GPU-based. It will limit flexibility if we dont switch between numpy and cupy throughout the code. (But maybe cupy-xarray handles that internally and we dont have to worry about that?)
2. Implement our own version of xarray using torch.Tensors.
   • Except for a few functionalities (like array broadcasting, grouping and data alignment), everything else must be easy enough to implement by just translating the labels to corresponding indices and redirecting to the existing DNDarray methods.
   • Example:

   >>> xarray = xr.DataArray(
           data=np.arange(48).reshape(4, 2, 6),
           dims=("u", "v", "time"),
           coords={
               "u": [-3.2, 2.1, 5.3, 6.5],
               "v": [-1, 2.6],
               "time": pd.date_range("2009-01-05", periods=6, freq="M"),
           },
       )
   
   # Single select
   >>> arr.sel(u=5.3, time="2009-04-30") # array([27, 33])
       # Translated to 
   >>> arr[2, :, 3] # array([27, 33])
   
   # Multi select
   >>> arr.sel(u=[-3.2, 6.5], time=slice("2009-02-28", "2009-03-31")) # array([[[ 1,  2], [ 7,  8]], [[37, 38], [43, 44]]])
       # Translated to
   >>> arr[[0, 3], :, slice(1, 2 + 1)] # array([[[ 1,  2], [ 7,  8]], [[37, 38], [43, 44]]])

   • But this means we will have to implement any new feature/functionality that is released by Xarray in the future.

I would love to know what you all think. Any suggestion is highly appreciated!

[TomNicholas]

Hi everyone, thanks for tagging me here!

I'm a bit unclear what you would like to achieve here - are you talking about:

(a) making DNDArray wrap xarray objects?
(b) wrapping DNDArray inside xarray objects? (related Q: does DNDArray obey the python array API standard?)
(c) wrapping DNDArray inside xarray objects but also with an understanding of chunks? (similar to dask - this is what the issue you linked to is discussing @ClaudiaComito)
(d) just creating DNDArray objects from xarray objects?

Some miscellaneous comments:

It will limit flexibility if we dont switch between numpy and cupy throughout the code.

This is the point of the get_array_module pattern - you should then not have to have separate code paths for working with GPU vs CPU data.

Implement our own version of xarray using torch.Tensors

There have been some discussions of wrapping pytorch tensors in xarray. We also have an ongoing project to publicly expose xarray.Variable, a lower-level abstraction better-suited for use as a building block of efforts to make named tensors like torch.Tensor. In a better-coordinated world this would have already happened, and pytorch would now have an optional dependency on xarray. 🤷‍♂️

[Mystic-Slice]

Hi everyone, thanks for tagging me here!

I'm a bit unclear what you would like to achieve here - are you talking about:

(a) making DNDArray wrap xarray objects? (b) wrapping DNDArray inside xarray objects? (related Q: does DNDArray obey the python array API standard?) (c) wrapping DNDArray inside xarray objects but also with an understanding of chunks? (similar to dask - this is what the issue you linked to is discussing @ClaudiaComito) (d) just creating DNDArray objects from xarray objects?

Some miscellaneous comments:

It will limit flexibility if we dont switch between numpy and cupy throughout the code.

This is the point of the get_array_module pattern - you should then not have to have separate code paths for working with GPU vs CPU data.

Implement our own version of xarray using torch.Tensors

There have been some discussions of wrapping pytorch tensors in xarray. We also have an ongoing project to publicly expose xarray.Variable, a lower-level abstraction better-suited for use as a building block of efforts to make named tensors like torch.Tensor. In a better-coordinated world this would have already happened, and pytorch would now have an optional dependency on xarray. 🤷‍♂️

Hi @TomNicholas!
We want to create a Distributed-Xarray. This should be able to be chunked and distributed to different processes. Most users would also prefer to use GPU-based computation.
So, I guess we are trying to achieve option (a) or create a new distributed data structure of our own that mimics Xarray (sounds like a lot of work😅).

This should allow users to manipulate huge amounts of data while also being able to work with the more-intuitive Xarray API.

@ClaudiaComito should be able to shed more light on this.
Thanks for taking the time to speak with us!

[koldunovn]

Would the experience of xarray with Dask make creation of the data structure you want? Also there are implementations with GPU support https://xarray.dev/blog/xarray-kvikio