xray nd-groupby

xray_nd_groupby.py

import functools

def nd_groupby_apply(ds, dims, func):
    if isinstance(dims, str):
        dims = [dims]
    if len(dims) > 1:
        func = functools.partial(nd_groupby_apply, dims=dims[1:], func=func) 
    return ds.groupby(dims[0]).apply(func)

def nd_group_over_apply(ds, dims, func):
    if isinstance(dims, str):
        dims = [dims]
    by_dims = [d for d in ds.dims if d not in dims]
    return nd_groupby_apply(ds, by_dims, func)