Simple numpy operations with xarray dataarray

[gewitterblitz]

Hi, I wanted to do the following simple numpy difference operations using xarray dataarray:

dvdy = np.empty_like(ds2)
dvdy[:, 1:-2] = (ds2[:, 2:-1] - ds2[:, 1:-2]) / 1000

However, I get the following error:

I can resolve the issue if I use ds2.values but that would mean loading the data into memory. I don't want to do that because the actual dataarray is much bigger (out-of-memory). The idea is to apply a chain of functions to my bigger dataarray using dask's client.submit function.

Is there a better way of doing this?

[TomNicholas]

Hi @gewitterblitz, would you mind posting a minimal reproducible example please? I ask because when I try to recreate your differencing operation it works as expected:

In [1]: import xarray as xr

In [2]: import numpy as np

In [3]: da = xr.DataArray(np.arange(10))

In [4]: da[2:-1] - da[1:-2]
Out[4]: 
<xarray.DataArray (dim_0: 7)>
array([1, 1, 1, 1, 1, 1, 1])
Dimensions without coordinates: dim_0

and I don't have enough information from your post to recreate your dvdy variable.

[gewitterblitz]

Hi @TomNicholas, thanks for your reply. Here you go!

da = xr.DataArray(np.random.rand(10,10,10),[("x", np.arange(10)),("y", np.arange(10)),("z", np.arange(10))])

dvdy = np.empty_like(da)

dvdy[:,1:-2,:] = (da[:, 2:-1,:] - da[:, 1:-2,:])

[TomNicholas]

Thanks @gewitterblitz .

So firstly you realise that dvdy in your example is a np.ndarray, not an xr.DataArray? I would not expect this assignment to work between different types of objects. However if we promote dvdy to a coordinate-less dataarray (for example like this)

dvdy = xr.DataArray(np.empty_like(da), dims=['x', 'y', 'z'])

then we can start to see what the problem is.

In [18]: da.shape
Out[18]: (10, 10, 10)

In [19]: da[:, 2:-1,:].shape
Out[19]: (10, 7, 10)

In [20]: da[:, 1:-2,:].shape
Out[20]: (10, 7, 10)

In [21]: (da[:, 2:-1,:] - da[:, 1:-2,:]).shape
Out[21]: (10, 6, 10)

In [22]: dvdy[:,1:-2,:].shape
Out[22]: (10, 7, 10)

The differencing operation is executing, but produces a differently-shaped array than the one you are trying to assign to. The reason for this is that by default xarray is doing automatic alignment using your coordinates, so the result only has the intersection of the coordinates of the input.

We can see this more clearly if we try giving dvdy some coordinates by creating it like this instead:

dvdy = xr.zeros_like(da)

Now if we do your operation we get a different, clearer, error:

In [6]: dvdy[:,1:-2,:] = (da[:, 2:-1,:] - da[:, 1:-2,:])
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-6-2492895e89d3> in <module>
----> 1 dvdy[:,1:-2,:] = (da[:, 2:-1,:] - da[:, 1:-2,:])

~/Documents/Work/Code/xarray/xarray/core/dataarray.py in __setitem__(self, key, value)
    749             obj = self[key]
    750             if isinstance(value, DataArray):
--> 751                 assert_coordinate_consistent(value, obj.coords.variables)
    752             # DataArray key -> Variable key
    753             key = {

~/Documents/Work/Code/xarray/xarray/core/coordinates.py in assert_coordinate_consistent(obj, coords)
    387         # make sure there are no conflict in dimension coordinates
    388         if k in coords and k in obj.coords and not coords[k].equals(obj[k].variable):
--> 389             raise IndexError(
    390                 f"dimension coordinate {k!r} conflicts between "
    391                 f"indexed and indexing objects:\n{obj[k]}\nvs.\n{coords[k]}"

IndexError: dimension coordinate 'y' conflicts between indexed and indexing objects:
<xarray.DataArray 'y' (y: 6)>
array([2, 3, 4, 5, 6, 7])
Coordinates:
  * y        (y) int64 2 3 4 5 6 7
vs.
<xarray.IndexVariable 'y' (y: 7)>
array([1, 2, 3, 4, 5, 6, 7])

You should have a read of the documentation on alignment, but actually I think your specific case might be easier solved by using a method like DataArray.roll or DataArray.shift (paying attention to the roll_coords keyword argument).

Finally if the reason you're doing this is to take differences of arrays living on geospatial grids, you might also be interested in the xGCM package.

(EDIT: Edited significantly to add example of dvdy with and without coordinates.)

[gewitterblitz]

Thanks @TomNicholas, that's good to know.

[TomNicholas]

Okay great, does that answer your question? Feel free to comment again if not!

[gewitterblitz]

Yes, it does. However, I was trying to use some legacy code (multiple functions; non-vectorized) for vorticity calculation and inclusion this in the pipeline was not the most efficient way.

I chose to write a vectorized version which uses dask gufunc and does the job.

I really appreciate your quick response and follow-up though :)