How to pass coordinates to map, apply_ufunc, or whatever

[openSourcerer9000]

It looks like someone else just asked this but was shot down. I need to apply a function over a gigantic raster that is aware of the coords at each pixel it is iterating over. How do I do this?

Example trying apply_ufunc, may not be the right direction:

da = xr.DataArray(np.random.randn(4, 5,1), dims=("x", "y","band"), coords={"x": range(4),"y":range(10,15),"band":[1],"spatial_ref":'EPSG:CODE'})
da.name = 'z'
ds = da.to_dataset()
ds

def innerfunc(z):
    return z + z['x'] + z['y'] #here is where I want to query the coordinates
def outerfunc(a):
    return xr.apply_ufunc(innerfunc,a,input_core_dims=[['x','y']])
outerfunc(ds)

...

----> 2     return z + z['x'] + z['y']
      3 def outerfunc(a):
      4     return xr.apply_ufunc(innerfunc,a,input_core_dims=[['x','y']])
      5 outerfunc(ds)

IndexError: only integers, slices (`:`), ellipsis (`...`), numpy.newaxis (`None`) and integer or boolean arrays are valid indices

[TomNicholas]

Hi there,

I think what's confusing here is the way that passing a Dataset to apply_ufunc causes the func to be mapped over the different variables in the dataset. I did not really appreciate that it did that, and it's certainly not clear from the current docstring, so thanks for flagging that so we can improve it!

As for your question I think you should simply explicitly pass the coordinate variables you need into xr.apply_ufunc as dataarray args, then you can operate on their raw data how you like, i.e.

def innerfunc(z, x, y):
    return z + x + y

def outerfunc(a):
    return xr.apply_ufunc(innerfunc, a, a['x'], a['y'], input_core_dims=[['x','y'], ['x'], ['y']])

In this case this will give an error though, because innerfunc would be trying to add arrays of different shape:

<ipython-input-50-538a01d37e0c> in innerfunc(z, x, y)
      1 def innerfunc(z, x, y):
----> 2     return z + x + y
      3 
      4 def outerfunc(a):
      5     return xr.apply_ufunc(innerfunc, a, a['x'], a['y'], input_core_dims=[['x','y'], ['x'], ['y']])

ValueError: operands could not be broadcast together with shapes (1,4,5) (4,)

Perhaps something closer to what you actually want involves broadcasting over the core dimensions first, like this

def innerfunc(z, x, y):
     return z + x + y

def outerfunc(a):
     xx, yy = xr.broadcast(a['x'], a['y'])
     return xr.apply_ufunc(innerfunc, a, xx, yy, 
                           input_core_dims=[['x','y'], ['x','y'], ['x', 'y']],
                           output_core_dims=[['x', 'y']])

which should perform the addition in innerfunc on the broadcast raw arrays to give you

In [49]: outerfunc(ds)
Out[49]: 
<xarray.Dataset>
Dimensions:      (band: 1, x: 4, y: 5)
Coordinates:
  * x            (x) int64 0 1 2 3
  * y            (y) int64 10 11 12 13 14
  * band         (band) int64 1
    spatial_ref  <U9 'EPSG:CODE'
Data variables:
    z            (band, x, y) float64 10.73 9.821 10.46 ... 14.16 14.24 17.05

Does that help?