Mann-Whitney U rank test for xarray with Benjamini and Hochberg correction

stats.py

import scipy as sp
import xarray as xr
import statsmodels as sm
import numpy as np
import statsmodels.stats.multitest


def mannwhitney(da1, da2, dim="time", global_alpha=0.05):
    """xarray wrapper for Mann-Whitney U test

    Paramters
    ---------
    da1, da2 : xr.DataArray
        arrays of samples
    dim : str, default: "time"
        Dimension along which to compute the test
    global_alpha, float, default: 0.05
        Global alpha of Benjamini and Hochberg correction
    
    
    """
    dim = [dim] if isinstance(dim, str) else dim

    def _mannwhitneyu(x, y):
        return sp.stats.mannwhitneyu(x, y, axis=-1, nan_policy="omit").pvalue

    # use xr.apply_ufunc to handle vectorization
    result = xr.apply_ufunc(
        _mannwhitneyu,
        da1,
        da2,
        input_core_dims=[dim, dim],
        output_core_dims=[[]],
        exclude_dims=set(dim),
        dask="parallelized",
        output_dtypes=[float],
    ).compute()

    # apply Benjamini and Hochberg correction
    shape = result.shape
    
    values = result.values.ravel()
    notnull = np.isfinite(values)

    p_adjust = sm.stats.multitest.multipletests(
        values[notnull], alpha=global_alpha, method="fdr_bh"
    )[0]
    values[notnull] = p_adjust
    result.values = values.reshape(shape)

    return result

Needs a newer scipy version (I have not explicitly checked but 0.14 works)

To test:

import xarray as xr
air = xr.tutorial.open_dataset("air_temperature")
da1 = air.air.isel(time=slice(None, 20))
da2 = air.air.isel(time=slice(20, 30))
mannwhitney(da1, da2)