How to apply a custom function to xarray.DataArray.coarsen.reduce()?

I have a (2x2) NumPy array:

ar = np.array([[2, 0],[3, 0]])

and the same one in the form of xarray.DataArray:

da = xr.DataArray(ar, dims=['x', 'y'], coords=[[0, 1], [0, 1]])

I am trying to downsample the 2d array spatially using a custom function to find the mode (i.e., the most frequently occurring value):

def find_mode(window):
    # find the mode over all axes 
    uniq = np.unique(window, return_counts=True)
    return uniq[0][np.argmax(uniq[1])]

The find_mode() works well for ar as find_mode(ar) gives 0.

However, it doesn't work for da (i.e., da.coarsen(x=2, y=2).reduce(find_mode)), with an error:

TypeError: find_mode() got an unexpected keyword argument 'axis'

Thank you so much for your attention and participation.

The signature for functions passed to DatasetCoarsen.reduce must include axis and kwargs. A good example is np.sum. So your function would need to look something like:

def find_mode(window, axis=None, **kwargs):
    # find the mode over all axes 
    uniq = np.unique(window, return_counts=True)
    
    ret = uniq[0][np.argmax(uniq[1])]
    ret = np.atleast_2d(ret)

    return ret

Depending on your application, you may want to use the axis argument (tuple of integers) in place of your [0] and [1] indexing steps.

Note: I've added np.atleast_2d here to make sure the return array is 2D. This is a bit ugly so I recommend thinking a bit more about this part in your application. The key thing to understand is that the return array needs to have the same number of dimensions as the input array.