Write back a slice of a dask array

[Shaunakde]

Hey,
I have a dataset with 3 dimensions ('time', 'x' and 'y'). I want to apply this function foo along the time dimension:

def foo(arr):
    lo, hi = np.percentile(arr, (1,99))
    arr = np.clip(arr, lo, hi)
    arr = (arr - lo) / (hi - lo)
    return arr

Basically, I want to scale each "image" in the time dimension with a function like foo

In numpy, I could just do something like:

for i in range(data.shape[0]):
   data[i] = foo(data[i])

but since the data is stored in dask arrays, I am unable to write back the modified data. How would one go about doing this in xarray/dask?

Shaunak

[dcherian]

You can vectorize that using the axis argument of np.percentile.

def foo(arr):
    lo, hi = np.percentile(arr, (1,99), axis=0)
    arr = np.clip(arr, lo, hi)
    arr = (arr - lo) / (hi - lo)
    return arr

data = foo(data) # should work for dask arrays.

With an xarray thing, see .quantile and .clip

[Shaunakde]

Thanks for this reply. Will try this solution. For now, I was able to use apply_ufuncs

dcube = xr.apply_ufunc(foo, dcube, input_core_dims=[["time"]], output_core_dims=[["time"]], dask = 'allowed', vectorize=True)

The key was passing the dask=allowed and vectorize=True params.