使用 xarray 滚动的有限差分

Question

My goal is to compute a derivative of a moving window of a multidimensional dataset along a given dimension, where the dataset is stored as Xarray DataArray or DataSet .

In the simplest case, given a 2D array I would like to compute a moving difference across multiple entries in one dimension, eg:

data = np.kron(np.linspace(0,1,10), np.linspace(1,4,6) ).reshape(10,6)
T=3
reducedArray = np.zeros_like(data)
for i in range(data.shape[1]):
    if i < T:
        reducedArray[:,i] = data[:,i] - data[:,0]
    else:
        reducedArray[:,i] = data[:,i] - data[:,i-T]

where the if i <T condition ensures that input and output contain proper values (ie, no nan s) and are of identical shape.

Xarray's diff aims to perform a finite-difference approximation of a given derivative order using nearest-neighbours, so it is not suitable here, hence the question:

Is it possible to perform this operation using Xarray functions only?

The rolling weighted average example appears to be something similar, but still too distinct due to the usage of NumPy routines. I've been thinking that something along the lines of the following should work:

xr2DDataArray = xr.DataArray(
    data,
    dims=('x','y'),
    coords={'x':np.linspace(0,1,10), 'y':np.linspace(1,4,6)}
)
r = xr2DDataArray.rolling(x=T,min_periods=2)
r.reduce( redFn )

I am struggling with the definition of redFn here,though.

Caveat The actual dataset to which the operation is to be applied will have a size of ~10GiB, so a solution that does not blow up the memory requirements will be highly appreciated!

Answer 1

I have never used xarray , so maybe I am mistaken, but I think you can get the result you want avoiding using loops and conditionals. This is at least twice faster than your example for numpy arrays:

data = np.kron(np.linspace(0,1,10), np.linspace(1,4,6)).reshape(10,6)
reducedArray = np.empty_like(data)
reducedArray[:, T:] = data[:, T:] - data[:, :-T]
reducedArray[:, :T] = data[:, :T] - data[:, 0, np.newaxis]

I imagine the improvement will be higher when using DataArray s. It does not use xarray functions but neither depends on numpy functions. I am confident that translating this to xarray will be straightforward, I know that it works if there are no coords , but once you include them, you get an error because of the coords mismatch ( coords of data[:, T:] and of data[:, :-T] are different). Sadly, I can't do better now.