Function test_apply_ufunc_dataset

xarray/tests/test_units.py:428–456

(variant, dtype)

    ),
)
def test_apply_ufunc_dataset(variant, dtype):
    variants = {
        "data": (unit_registry.m, 1, 1),
        "dims": (1, unit_registry.m, 1),
        "coords": (1, 1, unit_registry.s),
    }
    data_unit, dim_unit, coord_unit = variants[variant]
    func = functools.partial(
        xr.apply_ufunc, np.mean, input_core_dims=[["x"]], kwargs={"axis": -1}
    )

    array1 = np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * data_unit
    array2 = np.linspace(0, 10, 5).astype(dtype) * data_unit

    x = np.arange(5) * dim_unit
    y = np.arange(10) * dim_unit

    u = np.linspace(-1, 1, 10) * coord_unit

    ds = xr.Dataset(
        data_vars={"a": (("x", "y"), array1), "b": ("x", array2)},
        coords={"x": x, "y": y, "u": ("y", u)},
    )

    expected = attach_units(func(strip_units(ds)), extract_units(ds))
    actual = func(ds)

    assert_units_equal(expected, actual)
    assert_identical(expected, actual)


@pytest.mark.parametrize(