How to get apply_ufunc work with scipy's Mann-Whitney U test?

[nicrie]

I want to perform heaps of Mann-Whitney U tests given two DataArrays. Although scipy.stats.mannwhitneyu works for multi-dimensional input, I cannot make it work with xarrary's apply_ufunc.
I can make it work when using vectorize=True, however since speed is important for me and scipy's implementation allows for multidimensional input, I want to avoid the vectorized solution.

Here's a MWP:

import numpy as np
import scipy.stats
import xarray as xr

# Set a seed for reproducibility
np.random.seed(0)

# Define the dimensions and shape
samples_dim = 20
season_dim = 4
category_dim = 3

# Coordinates
samples_coords = np.arange(samples_dim)
season_coords = ['winter', 'spring', 'summer', 'autumn']
category_coords = ['fish', 'mammals', 'birds']

# Create the data for the arrays
data1 = np.random.rand(samples_dim, season_dim, category_dim)
data2 = np.random.rand(samples_dim, season_dim, category_dim)

# Create the xarray DataArrays
da1 = xr.DataArray(data1, dims=("samples", "season", "category"), coords=[samples_coords, season_coords, category_coords])
da2 = xr.DataArray(data2, dims=("samples", "season", "category"), coords=[samples_coords, season_coords, category_coords])


def MWU(x, y, dim):
    ''' xarray wrapper for scipy.stats.mannwhitneyu '''
    return xr.apply_ufunc(
        scipy.stats.mannwhitneyu,
        x,
        y,
        input_core_dims=[[dim], [dim]],
        output_core_dims=[[], []],
        vectorize=False  # I don't want to vectorize because speed is important
    )

MWU(da1, da2, dim='samples')

With vectorize=False I get the following error message:

ValueError: size of dimension 'season' on inputs was unexpectedly changed by applied function from 4 to 3. Only dimensions specified in exclude_dims with xarray.apply_ufunc are allowed to change size.

Note that the following works:

scipy.stats.mannwhitneyu(da1.values, da2.values, axis=0)

This would probably be my fallback solution ;)

Thanks for your help! :)

[kmuehlbauer]

@nicrie The core dimensions would need to be specified this way:

input_core_dims=[["samples", "season", "category"], ["samples", "season", "category"]],
output_core_dims=[["season", "category"], ["season", "category"]],

As the target function works on all dimensions, they are actually all input core dimensions. For the output you just have to specify the remaining dimensions.

[keewis]

The issue is this (from the documentation of the input_core_dims parameter):

Core dimensions are automatically moved to the last axes of input variables before applying func, which facilitates using NumPy style generalized ufuncs [2].

which means that we can't just use the default axis=0 but have to set axis=-1. This should work:

def MWU(x, y, dim):
    """ xarray wrapper for scipy.stats.mannwhitneyu """
    return xr.apply_ufunc(
        scipy.stats.mannwhitneyu,
        x,
        y,
        input_core_dims=[[dim], [dim]],
        output_core_dims=[[], []],
        vectorize=False  # I don't want to vectorize because speed is important
        kwargs={"axis": -1},
    )

[kmuehlbauer]

@keewis That's actually very nice, as it's only relying on the one dimension name. I'd favour this as a more general answer.

[nicrie]

I changed the accepted answer accordingly to point potential readers towards this solution.

[nicrie]

Thanks to both of you @kmuehlbauer & @keewis !
Both solutions work, speed is the same. I will accept @kmuehlbauer answer since he was the first ;)