Concatenating DataArray objects lazily?

[krikru]

As part of training a neural network with climate data as input, stored as netCDF files, each one containing a time series covering the same time period, I need to be able to load an arbitrary time point from each netCDF file (the same time point for each file) and concatenate those into one NumPy array.

If I use iris, I can just load the netCDF files into cubes, extract a Dask array from each of those cubes, and concatenate those Dask arrays into a single, large Dask array, and since this doesn't load any of the actual data (only the metadata) into memory (RAM), it goes almost instantly. I can then extract a slice corresponding to the time point I want from the concatenated Dask array, and call compute to get the time slice as a NumPy array.

Is there any similar way to do this in xarray? That is, can I load the same netCDF files and concatenate those into one xarray object of some kind, without loading anything else than metadata into memory, and then just extract a single time slice from that xarray object and first then read it to memory by requesting to get it as a NumPy array? Or do I have to load the netCDF files into separate xarray objects, extract time slices from those as separate NumPy arrays, and thus read them to memory, and then concatenate those?

I have tried doing xr.concat(variables, dim="level"), where variables is a list of xarray.DataArray objects that I load from the netCDF files in question, and "level" is an already existing dimension I want the variables to be concatenated along. However, this operation seems to take forever (in contrast with the Dask array concatenation which happens instantaneously), so it seems to me like xarray tries to load all data into memory before doing the concatenation instead of delaying loading the data to memory until I decide to access the values attribute (of the concatenated array), and then just loading the data corresponding to that time slice into memory (by loading from each netCDF file into memory and concatenating behind the scenes).

[TomNicholas]

Hey @krikru , you'll want to read #4628. Basically if you ensure your xarray variables are backed by dask arrays then concat should be lazy, but right now if your variables are backed by numpy arrays (i.e. they have been .load()-ed into memory) then xarray will do a lot of things lazily, but not concat (yet).

[krikru]

@TomNicholas Thanks for your answer. What does it mean that my xarray variables are backed by dask arrays? How do I ensure that?

[TomNicholas]

Hi @krikru

What does it mean that my xarray variables are backed by dask arrays?

Xarray objects wrap arrays. Each variable in an xarray.Dataset for example can can wrap either a numpy in-memory array, or a dask out-of-memory array. Xarray treats these basically the same way, so all operations you perform with Dataset work the same way. You can see which type of array is underneath in the xarray objects' printable repr. For example here the temperature variable is backed by a dask array (which itself wraps many small numpy arrays)

<xarray.Dataset>
Dimensions:      (time: 30, latitude: 180, longitude: 180)
Coordinates:
  * time         (time) datetime64[ns] 2015-01-01 2015-01-02 ... 2015-01-30
  * longitude    (longitude) int64 0 1 2 3 4 5 6 ... 173 174 175 176 177 178 179
  * latitude     (latitude) float64 89.5 88.5 87.5 86.5 ... -87.5 -88.5 -89.5
Data variables:
    temperature  (time, latitude, longitude) float64 dask.array<chunksize=(10, 180, 180), meta=np.ndarray>

You should read our docs page on using dask with xarray.

How do I ensure that?

Using the chunks argument when opening files, see https://docs.xarray.dev/en/stable/user-guide/dask.html#reading-and-writing-data