XArray apply_ufunc for function with multiple outputs?

Viewed 395

Simple question, but can't wrap my head around the documentation:

Given two DataArrays, how can one use apply_ufunc such that the outputs of the functions are collected in a new dimension?

For example:

test1 = xr.DataArray(np.linspace(1, 6, 6).reshape(3, 2))
test2 = xr.DataArray(np.linspace(6, 1, 6).reshape(3, 2))

def foo(a, b):
    return a+b, a-b

xr.apply_ufunc(foo, test1, test2)

This returns an error: ValueError: applied function returned data with unexpected number of dimensions: 3 vs 2, for dimensions ('dim_0', 'dim_1')

Any ideas of how to do this?

0 Answers
Related