Function original_and_stack_negative

xarray/tests/test_computation.py:525–532

(obj)

    assert_identical(stacked_dataset, stack_negative(dataset.groupby("x")))

    def original_and_stack_negative(obj):
        def func(x):
            return (x, np.stack([x, -x], axis=-1))

        result = apply_ufunc(func, obj, output_core_dims=[[], ["sign"]])
        if isinstance(result[1], xr.Dataset | xr.DataArray):
            result[1].coords["sign"] = [1, -1]
        return result

    out0, out1 = original_and_stack_negative(array)
    assert_identical(array, out0)