0

I am trying to apply a function to an Xarray dataset, using dataset.where mask to decide where to apply it. I am not sure how to do it.

Some context: the dataset has two variables (A and B), which are two overlapping raster images (same size and coordinates). I want to run a function (defined by me), that runs a calculation on each gridcell of the raster image, using the two variables. This is not a u_func, it is just a function that takes the values of the two overlapping gridcells, does some calculations and returns a third value I want to save to a third raster image (C).

How do I do that?

I have gotten this far, but it does not work because ds.A passes the entire array to the function, not just the value of ds.A in that particular gridcell:

def my_func(x, y):
    ..do things
    return result

ds = xr.Dataset({
    "A": xarray.open_rasterio("A.tif"), 
    "B": xarray.open_rasterio("B.tif"),
    "C": xarray.open_rasterio("C.tif"),
})

ds.C.where(ds.A > 0, my_func(ds.A, ds.B))
ffgg
  • 134
  • 1
  • 8
  • 1
    `Dataset.where` expects a scalar or array of some kind in order to replace the values in the original array with the values in the other array. Your issue is therefore only with `my_func(ds.A, ds.B)` because the function gets whole arrays and expects single grids. You need to create an auxiliar/wrapper function that loops over gridcells and passes single grids to `my_func`, this could be done with `for` loops or with [`numba.guvectorize`](https://numba.pydata.org/numba-doc/0.35.0/reference/jit-compilation.html#numba.guvectorize) – OriolAbril May 26 '20 at 14:40
  • The answer will therefore depend on the type of function of which we now nothing yet. I'd recommend searching for answers on vectorizing functions or ufunc creation to see if some meets your needs. Otherwise modify the question or post a new one. – OriolAbril May 26 '20 at 14:43
  • 1
    ok yes that's what I wanted to know, I think I was looking for an equivalent of numpy [apply_along_axis](https://numpy.org/doc/stable/reference/generated/numpy.apply_along_axis.html), which doesn't require me to create a ufunc. – ffgg May 27 '20 at 15:25
  • Great, glad I could help! – OriolAbril May 27 '20 at 15:57

0 Answers0