Closed ClaudiaComito closed 3 months ago
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.
Got it, thanks @koldunovn, edited original post.
Implementation discussed during devs meeting Sept 28. Preliminary scheme:
dndarray.coordinates
attribute (default=None
)ht.array()
to accept xarray
input, interpret string split
, populate dndarray.coordinates
, chunk data AND split coordinate if not coordinates:
or similar, quantify and minimize performance hit@Markus-Goetz would be good if you chime in
tagging @Mystic-Slice as they have shown interest in working on this :wave:
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.
Hi everyone! I have been thinking about this project for some time now and here are my initial thoughts.
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:
Implement our own version of xarray
using torch.Tensor
s.
DNDarray
methods.
>>> 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"),
},
)
arr.sel(u=5.3, time="2009-04-30") # array([27, 33])
Translated to
arr[2, :, 3] # array([27, 33])
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!
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. :man_shrugging:
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) wrappingDNDArray
inside xarray objects? (related Q: doesDNDArray
obey the python array API standard?) (c) wrappingDNDArray
inside xarray objects but also with an understanding ofchunks
? (similar to dask - this is what the issue you linked to is discussing @ClaudiaComito) (d) just creatingDNDArray
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 liketorch.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!
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
Meanwhile I have implemented a basic idea of DXarray
in the corresponding branch. In the near future I plan to go on with:
DXarray.to_xarray()
: convert distributed DXarray
to a xarray.DataArray
(similar to DNDarray.to_numpy()
)resplit
for DXarray
from_xarray': convert a non-distributed
xarray.DataArrayto distributed
DXarray`DXarray
(depending on how this can be done for an xarray
-like structure)value
-array only (such as mean, etc.) Things that might get a bit more complicated:
@ClaudiaComito in #1154 it seems you started wrapping heat objects inside xarray, which is awesome! I recently improved the documentation on wrapping numpy-like xarrays with xarray objects (https://github.com/pydata/xarray/pull/7911 and https://github.com/pydata/xarray/pull/7951).
Those extra pages in the docs aren't released yet, but for now you can view them here (on wrapping numpy-like arrays) and here (on wrapping distributed numpy-like arrays).
In the branch 1031-support-distribution-of-xarray there is now available:
DXarray
implementing sth similar as xarray
but using DNDarrays
as underlying objectsDXarray
xarray.DataArrays
's
(still missing: unit tests...) I will stop here, until we have discussed in the team the "wrapping-approach" proposed by TomNicholas, because such an approach would be much easier to implement (if applicable to Heat).
This issue is stale because it has been open for 60 days with no activity.
This issue was closed because it has been inactive for 60 days since being marked as stale.
See https://docs.xarray.dev/en/stable/
If I understand correctly, an
xarray
object is made up of the actualdata
array (np.ndarray), and1-Dcoordinates
arrays (dictionaries?) that mapdata
dimensions and indices to meaningful physical quantities.For example, if
xarray
is a matrix of coordinates(date, temperature)
, users will be able to performFeature functionality
Enable distribution of xarray object, allow named dimensions, keep track of
coordinates
arrays, one of which will be distributed.Example, :
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.