diff --git a/ci/requirements/py36-min-nep18.yml b/ci/requirements/py36-min-nep18.yml index cd2b1a18c77..dd543ce4ddf 100644 --- a/ci/requirements/py36-min-nep18.yml +++ b/ci/requirements/py36-min-nep18.yml @@ -11,6 +11,7 @@ dependencies: - msgpack-python=0.6 # remove once distributed is bumped. distributed GH3491 - numpy=1.17 - pandas=0.25 + - pint=0.13 - pip - pytest - pytest-cov @@ -18,5 +19,3 @@ dependencies: - scipy=1.2 - setuptools=41.2 - sparse=0.8 - - pip: - - pint==0.13 diff --git a/ci/requirements/py36.yml b/ci/requirements/py36.yml index aa2baf9dcce..a500173f277 100644 --- a/ci/requirements/py36.yml +++ b/ci/requirements/py36.yml @@ -28,6 +28,7 @@ dependencies: - numba - numpy - pandas + - pint - pip - pseudonetcdf - pydap @@ -44,4 +45,3 @@ dependencies: - zarr - pip: - numbagg - - pint diff --git a/ci/requirements/py37-windows.yml b/ci/requirements/py37-windows.yml index 8b12704d644..e9e5c7a900a 100644 --- a/ci/requirements/py37-windows.yml +++ b/ci/requirements/py37-windows.yml @@ -28,6 +28,7 @@ dependencies: - numba - numpy - pandas + - pint - pip - pseudonetcdf - pydap @@ -44,4 +45,3 @@ dependencies: - zarr - pip: - numbagg - - pint diff --git a/ci/requirements/py37.yml b/ci/requirements/py37.yml index 70c453e8776..dba3926596e 100644 --- a/ci/requirements/py37.yml +++ b/ci/requirements/py37.yml @@ -28,6 +28,7 @@ dependencies: - numba - numpy - pandas + - pint - pip - pseudonetcdf - pydap @@ -44,4 +45,3 @@ dependencies: - zarr - pip: - numbagg - - pint diff --git a/ci/requirements/py38-all-but-dask.yml b/ci/requirements/py38-all-but-dask.yml index 6d76eecbd6a..a375d9e1e5a 100644 --- a/ci/requirements/py38-all-but-dask.yml +++ b/ci/requirements/py38-all-but-dask.yml @@ -25,6 +25,7 @@ dependencies: - numba - numpy - pandas + - pint - pip - pseudonetcdf - pydap @@ -41,4 +42,3 @@ dependencies: - zarr - pip: - numbagg - - pint diff --git a/ci/requirements/py38.yml b/ci/requirements/py38.yml index 6f35138978c..7dff3a1bd97 100644 --- a/ci/requirements/py38.yml +++ b/ci/requirements/py38.yml @@ -28,6 +28,7 @@ dependencies: - numba - numpy - pandas + - pint - pip - pseudonetcdf - pydap @@ -44,4 +45,3 @@ dependencies: - zarr - pip: - numbagg - - pint diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 086cddee0a0..e4223f2b4e0 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -91,7 +91,7 @@ New Features - Support dask handling for :py:meth:`DataArray.idxmax`, :py:meth:`DataArray.idxmin`, :py:meth:`Dataset.idxmax`, :py:meth:`Dataset.idxmin`. (:pull:`3922`, :pull:`4135`) By `Kai Mühlbauer `_ and `Pascal Bourgault `_. -- More support for unit aware arrays with pint (:pull:`3643`, :pull:`3975`) +- More support for unit aware arrays with pint (:pull:`3643`, :pull:`3975`, :pull:`4163`) By `Justus Magin `_. - Support overriding existing variables in ``to_zarr()`` with ``mode='a'`` even without ``append_dim``, as long as dimension sizes do not change. diff --git a/xarray/core/common.py b/xarray/core/common.py index f759f4c32dd..67dc0fda461 100644 --- a/xarray/core/common.py +++ b/xarray/core/common.py @@ -1434,7 +1434,7 @@ def _full_like_variable(other, fill_value, dtype: DTypeLike = None): other.shape, fill_value, dtype=dtype, chunks=other.data.chunks ) else: - data = np.full_like(other, fill_value, dtype=dtype) + data = np.full_like(other.data, fill_value, dtype=dtype) return Variable(dims=other.dims, data=data, attrs=other.attrs) diff --git a/xarray/core/utils.py b/xarray/core/utils.py index 0542f850b02..668405ba574 100644 --- a/xarray/core/utils.py +++ b/xarray/core/utils.py @@ -247,6 +247,12 @@ def is_list_like(value: Any) -> bool: return isinstance(value, list) or isinstance(value, tuple) +def is_array_like(value: Any) -> bool: + return ( + hasattr(value, "ndim") and hasattr(value, "shape") and hasattr(value, "dtype") + ) + + def either_dict_or_kwargs( pos_kwargs: Optional[Mapping[Hashable, T]], kw_kwargs: Mapping[str, T], diff --git a/xarray/testing.py b/xarray/testing.py index 9681503414e..ec479ef09d4 100644 --- a/xarray/testing.py +++ b/xarray/testing.py @@ -11,7 +11,14 @@ from xarray.core.indexes import default_indexes from xarray.core.variable import IndexVariable, Variable -__all__ = ("assert_allclose", "assert_chunks_equal", "assert_equal", "assert_identical") +__all__ = ( + "assert_allclose", + "assert_chunks_equal", + "assert_duckarray_equal", + "assert_duckarray_allclose", + "assert_equal", + "assert_identical", +) def _decode_string_data(data): @@ -148,6 +155,62 @@ def compat_variable(a, b): raise TypeError("{} not supported by assertion comparison".format(type(a))) +def _format_message(x, y, err_msg, verbose): + diff = x - y + abs_diff = max(abs(diff)) + rel_diff = "not implemented" + + n_diff = int(np.count_nonzero(diff)) + n_total = diff.size + + fraction = f"{n_diff} / {n_total}" + percentage = float(n_diff / n_total * 100) + + parts = [ + "Arrays are not equal", + err_msg, + f"Mismatched elements: {fraction} ({percentage:.0f}%)", + f"Max absolute difference: {abs_diff}", + f"Max relative difference: {rel_diff}", + ] + if verbose: + parts += [ + f" x: {x!r}", + f" y: {y!r}", + ] + + return "\n".join(parts) + + +def assert_duckarray_allclose( + actual, desired, rtol=1e-07, atol=0, err_msg="", verbose=True +): + """ Like `np.testing.assert_allclose`, but for duckarrays. """ + __tracebackhide__ = True + + allclose = duck_array_ops.allclose_or_equiv(actual, desired, rtol=rtol, atol=atol) + assert allclose, _format_message(actual, desired, err_msg=err_msg, verbose=verbose) + + +def assert_duckarray_equal(x, y, err_msg="", verbose=True): + """ Like `np.testing.assert_array_equal`, but for duckarrays """ + __tracebackhide__ = True + + if not utils.is_array_like(x) and not utils.is_scalar(x): + x = np.asarray(x) + + if not utils.is_array_like(y) and not utils.is_scalar(y): + y = np.asarray(y) + + if (utils.is_array_like(x) and utils.is_scalar(y)) or ( + utils.is_scalar(x) and utils.is_array_like(y) + ): + equiv = (x == y).all() + else: + equiv = duck_array_ops.array_equiv(x, y) + assert equiv, _format_message(x, y, err_msg=err_msg, verbose=verbose) + + def assert_chunks_equal(a, b): """ Assert that chunksizes along chunked dimensions are equal. diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py index 40c5cfa267c..9021c4e7dbc 100644 --- a/xarray/tests/__init__.py +++ b/xarray/tests/__init__.py @@ -16,6 +16,10 @@ from xarray.core.duck_array_ops import allclose_or_equiv # noqa: F401 from xarray.core.indexing import ExplicitlyIndexed from xarray.core.options import set_options +from xarray.testing import ( # noqa: F401 + assert_duckarray_allclose, + assert_duckarray_equal, +) # import mpl and change the backend before other mpl imports try: diff --git a/xarray/tests/test_testing.py b/xarray/tests/test_testing.py index f4961af58e9..39ad250246b 100644 --- a/xarray/tests/test_testing.py +++ b/xarray/tests/test_testing.py @@ -1,7 +1,31 @@ +import numpy as np import pytest import xarray as xr +from . import has_dask + +try: + from dask.array import from_array as dask_from_array +except ImportError: + dask_from_array = lambda x: x + +try: + import pint + + unit_registry = pint.UnitRegistry(force_ndarray_like=True) + + def quantity(x): + return unit_registry.Quantity(x, "m") + + has_pint = True +except ImportError: + + def quantity(x): + return x + + has_pint = False + def test_allclose_regression(): x = xr.DataArray(1.01) @@ -30,3 +54,78 @@ def test_allclose_regression(): def test_assert_allclose(obj1, obj2): with pytest.raises(AssertionError): xr.testing.assert_allclose(obj1, obj2) + + +@pytest.mark.filterwarnings("error") +@pytest.mark.parametrize( + "duckarray", + ( + pytest.param(np.array, id="numpy"), + pytest.param( + dask_from_array, + id="dask", + marks=pytest.mark.skipif(not has_dask, reason="requires dask"), + ), + pytest.param( + quantity, + id="pint", + marks=[ + pytest.mark.skipif(not has_pint, reason="requires pint"), + pytest.mark.xfail( + reason="inconsistencies in the return value of pint's implementation of eq" + ), + ], + ), + ), +) +@pytest.mark.parametrize( + ["obj1", "obj2"], + ( + pytest.param([1e-10, 2], [0.0, 2.0], id="both arrays"), + pytest.param([1e-17, 2], 0.0, id="second scalar"), + pytest.param(0.0, [1e-17, 2], id="first scalar"), + ), +) +def test_assert_duckarray_equal_failing(duckarray, obj1, obj2): + # TODO: actually check the repr + a = duckarray(obj1) + b = duckarray(obj2) + with pytest.raises(AssertionError): + xr.testing.assert_duckarray_equal(a, b) + + +@pytest.mark.filterwarnings("error") +@pytest.mark.parametrize( + "duckarray", + ( + pytest.param(np.array, id="numpy"), + pytest.param( + dask_from_array, + id="dask", + marks=pytest.mark.skipif(not has_dask, reason="requires dask"), + ), + pytest.param( + quantity, + id="pint", + marks=[ + pytest.mark.skipif(not has_pint, reason="requires pint"), + pytest.mark.xfail( + reason="inconsistencies in the return value of pint's implementation of eq" + ), + ], + ), + ), +) +@pytest.mark.parametrize( + ["obj1", "obj2"], + ( + pytest.param([0, 2], [0.0, 2.0], id="both arrays"), + pytest.param([0, 0], 0.0, id="second scalar"), + pytest.param(0.0, [0, 0], id="first scalar"), + ), +) +def test_assert_duckarray_equal(duckarray, obj1, obj2): + a = duckarray(obj1) + b = duckarray(obj2) + + xr.testing.assert_duckarray_equal(a, b) diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py index 20a5f0e8613..619fa10116d 100644 --- a/xarray/tests/test_units.py +++ b/xarray/tests/test_units.py @@ -1,16 +1,16 @@ import functools import operator -from distutils.version import LooseVersion import numpy as np import pandas as pd import pytest import xarray as xr +from xarray.core import dtypes from xarray.core.npcompat import IS_NEP18_ACTIVE -from xarray.testing import assert_allclose, assert_equal, assert_identical -from .test_variable import _PAD_XR_NP_ARGS, VariableSubclassobjects +from . import assert_allclose, assert_duckarray_allclose, assert_equal, assert_identical +from .test_variable import _PAD_XR_NP_ARGS pint = pytest.importorskip("pint") DimensionalityError = pint.errors.DimensionalityError @@ -26,7 +26,7 @@ pytest.mark.skipif( not IS_NEP18_ACTIVE, reason="NUMPY_EXPERIMENTAL_ARRAY_FUNCTION is not enabled" ), - # pytest.mark.filterwarnings("ignore:::pint[.*]"), + pytest.mark.filterwarnings("error::pint.UnitStrippedWarning"), ] @@ -180,12 +180,7 @@ def attach_units(obj, units): new_obj = xr.Dataset(data_vars=data_vars, coords=coords, attrs=obj.attrs) elif isinstance(obj, xr.DataArray): # try the array name, "data" and None, then fall back to dimensionless - data_units = ( - units.get(obj.name, None) - or units.get("data", None) - or units.get(None, None) - or 1 - ) + data_units = units.get(obj.name, None) or units.get(None, None) or 1 data = array_attach_units(obj.data, data_units) @@ -264,7 +259,7 @@ def assert_units_equal(a, b): assert extract_units(a) == extract_units(b) -@pytest.fixture(params=[float, int]) +@pytest.fixture(params=[np.dtype(float), np.dtype(int)], ids=str) def dtype(request): return request.param @@ -364,14 +359,31 @@ def __repr__(self): return f"function_{self.name}" -def test_apply_ufunc_dataarray(dtype): +@pytest.mark.parametrize( + "variant", + ( + "data", + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", + ), +) +def test_apply_ufunc_dataarray(variant, dtype): + variants = { + "data": (unit_registry.m, 1, 1), + "dims": (1, unit_registry.m, 1), + "coords": (1, 1, unit_registry.m), + } + data_unit, dim_unit, coord_unit = variants.get(variant) func = functools.partial( xr.apply_ufunc, np.mean, input_core_dims=[["x"]], kwargs={"axis": -1} ) - array = np.linspace(0, 10, 20).astype(dtype) * unit_registry.m - x = np.arange(20) * unit_registry.s - data_array = xr.DataArray(data=array, dims="x", coords={"x": x}) + array = np.linspace(0, 10, 20).astype(dtype) * data_unit + x = np.arange(20) * dim_unit + u = np.linspace(-1, 1, 20) * coord_unit + data_array = xr.DataArray(data=array, dims="x", coords={"x": x, "u": ("x", u)}) expected = attach_units(func(strip_units(data_array)), extract_units(data_array)) actual = func(data_array) @@ -380,20 +392,39 @@ def test_apply_ufunc_dataarray(dtype): assert_identical(expected, actual) -def test_apply_ufunc_dataset(dtype): +@pytest.mark.parametrize( + "variant", + ( + "data", + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", + ), +) +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.get(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) * unit_registry.m - array2 = np.linspace(0, 10, 5).astype(dtype) * unit_registry.m + 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 - x = np.arange(5) * unit_registry.s - y = np.arange(10) * unit_registry.m + 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}, + coords={"x": x, "y": y, "u": ("y", u)}, ) expected = attach_units(func(strip_units(ds)), extract_units(ds)) @@ -403,10 +434,6 @@ def test_apply_ufunc_dataset(dtype): assert_identical(expected, actual) -# TODO: remove once pint==0.12 has been released -@pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose" -) @pytest.mark.parametrize( "unit,error", ( @@ -424,44 +451,61 @@ def test_apply_ufunc_dataset(dtype): "variant", ( "data", - pytest.param("dims", marks=pytest.mark.xfail(reason="indexes strip units")), + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), "coords", ), ) -@pytest.mark.parametrize("fill_value", (10, np.nan)) -def test_align_dataarray(fill_value, variant, unit, error, dtype): +@pytest.mark.parametrize("value", (10, dtypes.NA)) +def test_align_dataarray(value, variant, unit, error, dtype): + if variant == "coords" and ( + value != dtypes.NA or isinstance(unit, unit_registry.Unit) + ): + pytest.xfail( + reason=( + "fill_value is used for both data variables and coords. " + "See https://github.com/pydata/xarray/issues/4165" + ) + ) + + fill_value = dtypes.get_fill_value(dtype) if value == dtypes.NA else value + original_unit = unit_registry.m variants = { - "data": (unit, original_unit, original_unit), - "dims": (original_unit, unit, original_unit), - "coords": (original_unit, original_unit, unit), + "data": ((original_unit, unit), (1, 1), (1, 1)), + "dims": ((1, 1), (original_unit, unit), (1, 1)), + "coords": ((1, 1), (1, 1), (original_unit, unit)), } - data_unit, dim_unit, coord_unit = variants.get(variant) + ( + (data_unit1, data_unit2), + (dim_unit1, dim_unit2), + (coord_unit1, coord_unit2), + ) = variants.get(variant) - array1 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * original_unit - array2 = np.linspace(0, 8, 2 * 5).reshape(2, 5).astype(dtype) * data_unit - x = np.arange(2) * original_unit + array1 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * data_unit1 + array2 = np.linspace(0, 8, 2 * 5).reshape(2, 5).astype(dtype) * data_unit2 - y1 = np.arange(5) * original_unit - y2 = np.arange(2, 7) * dim_unit - y_a1 = np.array([3, 5, 7, 8, 9]) * original_unit - y_a2 = np.array([7, 8, 9, 11, 13]) * coord_unit + x = np.arange(2) * dim_unit1 + y1 = np.arange(5) * dim_unit1 + y2 = np.arange(2, 7) * dim_unit2 + + u1 = np.array([3, 5, 7, 8, 9]) * coord_unit1 + u2 = np.array([7, 8, 9, 11, 13]) * coord_unit2 coords1 = {"x": x, "y": y1} coords2 = {"x": x, "y": y2} if variant == "coords": - coords1["y_a"] = ("y", y_a1) - coords2["y_a"] = ("y", y_a2) + coords1["y_a"] = ("y", u1) + coords2["y_a"] = ("y", u2) data_array1 = xr.DataArray(data=array1, coords=coords1, dims=("x", "y")) data_array2 = xr.DataArray(data=array2, coords=coords2, dims=("x", "y")) - fill_value = fill_value * data_unit + fill_value = fill_value * data_unit2 func = function(xr.align, join="outer", fill_value=fill_value) - if error is not None and not ( - np.isnan(fill_value) and not isinstance(fill_value, Quantity) - ): + if error is not None and (value != dtypes.NA or isinstance(fill_value, Quantity)): with pytest.raises(error): func(data_array1, data_array2) @@ -469,7 +513,7 @@ def test_align_dataarray(fill_value, variant, unit, error, dtype): stripped_kwargs = { key: strip_units( - convert_units(value, {None: original_unit if data_unit != 1 else None}) + convert_units(value, {None: data_unit1 if data_unit2 != 1 else None}) ) for key, value in func.kwargs.items() } @@ -494,10 +538,6 @@ def test_align_dataarray(fill_value, variant, unit, error, dtype): assert_allclose(expected_b, actual_b) -# TODO: remove once pint==0.12 has been released -@pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose" -) @pytest.mark.parametrize( "unit,error", ( @@ -515,45 +555,61 @@ def test_align_dataarray(fill_value, variant, unit, error, dtype): "variant", ( "data", - pytest.param("dims", marks=pytest.mark.xfail(reason="indexes strip units")), + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), "coords", ), ) -@pytest.mark.parametrize("fill_value", (np.float64(10), np.float64(np.nan))) -def test_align_dataset(fill_value, unit, variant, error, dtype): +@pytest.mark.parametrize("value", (10, dtypes.NA)) +def test_align_dataset(value, unit, variant, error, dtype): + if variant == "coords" and ( + value != dtypes.NA or isinstance(unit, unit_registry.Unit) + ): + pytest.xfail( + reason=( + "fill_value is used for both data variables and coords. " + "See https://github.com/pydata/xarray/issues/4165" + ) + ) + + fill_value = dtypes.get_fill_value(dtype) if value == dtypes.NA else value + original_unit = unit_registry.m variants = { - "data": (unit, original_unit, original_unit), - "dims": (original_unit, unit, original_unit), - "coords": (original_unit, original_unit, unit), + "data": ((original_unit, unit), (1, 1), (1, 1)), + "dims": ((1, 1), (original_unit, unit), (1, 1)), + "coords": ((1, 1), (1, 1), (original_unit, unit)), } - data_unit, dim_unit, coord_unit = variants.get(variant) + ( + (data_unit1, data_unit2), + (dim_unit1, dim_unit2), + (coord_unit1, coord_unit2), + ) = variants.get(variant) - array1 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * original_unit - array2 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * data_unit + array1 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * data_unit1 + array2 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * data_unit2 - x = np.arange(2) * original_unit + x = np.arange(2) * dim_unit1 + y1 = np.arange(5) * dim_unit1 + y2 = np.arange(2, 7) * dim_unit2 - y1 = np.arange(5) * original_unit - y2 = np.arange(2, 7) * dim_unit - y_a1 = np.array([3, 5, 7, 8, 9]) * original_unit - y_a2 = np.array([7, 8, 9, 11, 13]) * coord_unit + u1 = np.array([3, 5, 7, 8, 9]) * coord_unit1 + u2 = np.array([7, 8, 9, 11, 13]) * coord_unit2 coords1 = {"x": x, "y": y1} coords2 = {"x": x, "y": y2} if variant == "coords": - coords1["y_a"] = ("y", y_a1) - coords2["y_a"] = ("y", y_a2) + coords1["u"] = ("y", u1) + coords2["u"] = ("y", u2) ds1 = xr.Dataset(data_vars={"a": (("x", "y"), array1)}, coords=coords1) ds2 = xr.Dataset(data_vars={"a": (("x", "y"), array2)}, coords=coords2) - fill_value = fill_value * data_unit + fill_value = fill_value * data_unit2 func = function(xr.align, join="outer", fill_value=fill_value) - if error is not None and not ( - np.isnan(fill_value) and not isinstance(fill_value, Quantity) - ): + if error is not None and (value != dtypes.NA or isinstance(fill_value, Quantity)): with pytest.raises(error): func(ds1, ds2) @@ -561,14 +617,14 @@ def test_align_dataset(fill_value, unit, variant, error, dtype): stripped_kwargs = { key: strip_units( - convert_units(value, {None: original_unit if data_unit != 1 else None}) + convert_units(value, {None: data_unit1 if data_unit2 != 1 else None}) ) for key, value in func.kwargs.items() } units_a = extract_units(ds1) units_b = extract_units(ds2) expected_a, expected_b = func( - strip_units(ds1), strip_units(convert_units(ds2, units_a)), **stripped_kwargs + strip_units(ds1), strip_units(convert_units(ds2, units_a)), **stripped_kwargs, ) expected_a = attach_units(expected_a, units_a) if isinstance(array2, Quantity): @@ -585,6 +641,7 @@ def test_align_dataset(fill_value, unit, variant, error, dtype): def test_broadcast_dataarray(dtype): + # uses align internally so more thorough tests are not needed array1 = np.linspace(0, 10, 2) * unit_registry.Pa array2 = np.linspace(0, 10, 3) * unit_registry.Pa @@ -606,6 +663,7 @@ def test_broadcast_dataarray(dtype): def test_broadcast_dataset(dtype): + # uses align internally so more thorough tests are not needed array1 = np.linspace(0, 10, 2) * unit_registry.Pa array2 = np.linspace(0, 10, 3) * unit_registry.Pa @@ -657,7 +715,9 @@ def test_broadcast_dataset(dtype): "variant", ( "data", - pytest.param("dims", marks=pytest.mark.xfail(reason="indexes strip units")), + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), "coords", ), ) @@ -665,31 +725,35 @@ def test_combine_by_coords(variant, unit, error, dtype): original_unit = unit_registry.m variants = { - "data": (unit, original_unit, original_unit), - "dims": (original_unit, unit, original_unit), - "coords": (original_unit, original_unit, unit), + "data": ((original_unit, unit), (1, 1), (1, 1)), + "dims": ((1, 1), (original_unit, unit), (1, 1)), + "coords": ((1, 1), (1, 1), (original_unit, unit)), } - data_unit, dim_unit, coord_unit = variants.get(variant) - - array1 = np.zeros(shape=(2, 3), dtype=dtype) * original_unit - array2 = np.zeros(shape=(2, 3), dtype=dtype) * original_unit - x = np.arange(1, 4) * 10 * original_unit - y = np.arange(2) * original_unit - z = np.arange(3) * original_unit - - other_array1 = np.ones_like(array1) * data_unit - other_array2 = np.ones_like(array2) * data_unit - other_x = np.arange(1, 4) * 10 * dim_unit - other_y = np.arange(2, 4) * dim_unit - other_z = np.arange(3, 6) * coord_unit + ( + (data_unit1, data_unit2), + (dim_unit1, dim_unit2), + (coord_unit1, coord_unit2), + ) = variants.get(variant) + + array1 = np.zeros(shape=(2, 3), dtype=dtype) * data_unit1 + array2 = np.zeros(shape=(2, 3), dtype=dtype) * data_unit1 + x = np.arange(1, 4) * 10 * dim_unit1 + y = np.arange(2) * dim_unit1 + u = np.arange(3) * coord_unit1 + + other_array1 = np.ones_like(array1) * data_unit2 + other_array2 = np.ones_like(array2) * data_unit2 + other_x = np.arange(1, 4) * 10 * dim_unit2 + other_y = np.arange(2, 4) * dim_unit2 + other_u = np.arange(3, 6) * coord_unit2 ds = xr.Dataset( data_vars={"a": (("y", "x"), array1), "b": (("y", "x"), array2)}, - coords={"x": x, "y": y, "z": ("x", z)}, + coords={"x": x, "y": y, "u": ("x", u)}, ) other = xr.Dataset( data_vars={"a": (("y", "x"), other_array1), "b": (("y", "x"), other_array2)}, - coords={"x": other_x, "y": other_y, "z": ("x", other_z)}, + coords={"x": other_x, "y": other_y, "u": ("x", other_u)}, ) if error is not None: @@ -728,7 +792,9 @@ def test_combine_by_coords(variant, unit, error, dtype): "variant", ( "data", - pytest.param("dims", marks=pytest.mark.xfail(reason="indexes strip units")), + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), "coords", ), ) @@ -736,18 +802,22 @@ def test_combine_nested(variant, unit, error, dtype): original_unit = unit_registry.m variants = { - "data": (unit, original_unit, original_unit), - "dims": (original_unit, unit, original_unit), - "coords": (original_unit, original_unit, unit), + "data": ((original_unit, unit), (1, 1), (1, 1)), + "dims": ((1, 1), (original_unit, unit), (1, 1)), + "coords": ((1, 1), (1, 1), (original_unit, unit)), } - data_unit, dim_unit, coord_unit = variants.get(variant) + ( + (data_unit1, data_unit2), + (dim_unit1, dim_unit2), + (coord_unit1, coord_unit2), + ) = variants.get(variant) - array1 = np.zeros(shape=(2, 3), dtype=dtype) * original_unit - array2 = np.zeros(shape=(2, 3), dtype=dtype) * original_unit + array1 = np.zeros(shape=(2, 3), dtype=dtype) * data_unit1 + array2 = np.zeros(shape=(2, 3), dtype=dtype) * data_unit1 - x = np.arange(1, 4) * 10 * original_unit - y = np.arange(2) * original_unit - z = np.arange(3) * original_unit + x = np.arange(1, 4) * 10 * dim_unit1 + y = np.arange(2) * dim_unit1 + z = np.arange(3) * coord_unit1 ds1 = xr.Dataset( data_vars={"a": (("y", "x"), array1), "b": (("y", "x"), array2)}, @@ -755,35 +825,35 @@ def test_combine_nested(variant, unit, error, dtype): ) ds2 = xr.Dataset( data_vars={ - "a": (("y", "x"), np.ones_like(array1) * data_unit), - "b": (("y", "x"), np.ones_like(array2) * data_unit), + "a": (("y", "x"), np.ones_like(array1) * data_unit2), + "b": (("y", "x"), np.ones_like(array2) * data_unit2), }, coords={ - "x": np.arange(3) * dim_unit, - "y": np.arange(2, 4) * dim_unit, - "z": ("x", np.arange(-3, 0) * coord_unit), + "x": np.arange(3) * dim_unit2, + "y": np.arange(2, 4) * dim_unit2, + "z": ("x", np.arange(-3, 0) * coord_unit2), }, ) ds3 = xr.Dataset( data_vars={ - "a": (("y", "x"), np.zeros_like(array1) * np.nan * data_unit), - "b": (("y", "x"), np.zeros_like(array2) * np.nan * data_unit), + "a": (("y", "x"), np.full_like(array1, fill_value=np.nan) * data_unit2), + "b": (("y", "x"), np.full_like(array2, fill_value=np.nan) * data_unit2), }, coords={ - "x": np.arange(3, 6) * dim_unit, - "y": np.arange(4, 6) * dim_unit, - "z": ("x", np.arange(3, 6) * coord_unit), + "x": np.arange(3, 6) * dim_unit2, + "y": np.arange(4, 6) * dim_unit2, + "z": ("x", np.arange(3, 6) * coord_unit2), }, ) ds4 = xr.Dataset( data_vars={ - "a": (("y", "x"), -1 * np.ones_like(array1) * data_unit), - "b": (("y", "x"), -1 * np.ones_like(array2) * data_unit), + "a": (("y", "x"), -1 * np.ones_like(array1) * data_unit2), + "b": (("y", "x"), -1 * np.ones_like(array2) * data_unit2), }, coords={ - "x": np.arange(6, 9) * dim_unit, - "y": np.arange(6, 8) * dim_unit, - "z": ("x", np.arange(6, 9) * coord_unit), + "x": np.arange(6, 9) * dim_unit2, + "y": np.arange(6, 8) * dim_unit2, + "z": ("x", np.arange(6, 9) * coord_unit2), }, ) @@ -828,22 +898,37 @@ def test_combine_nested(variant, unit, error, dtype): "variant", ( "data", - pytest.param("dims", marks=pytest.mark.xfail(reason="indexes strip units")), + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", ), ) def test_concat_dataarray(variant, unit, error, dtype): original_unit = unit_registry.m - variants = {"data": (unit, original_unit), "dims": (original_unit, unit)} - data_unit, dims_unit = variants.get(variant) + variants = { + "data": ((original_unit, unit), (1, 1), (1, 1)), + "dims": ((1, 1), (original_unit, unit), (1, 1)), + "coords": ((1, 1), (1, 1), (original_unit, unit)), + } + ( + (data_unit1, data_unit2), + (dim_unit1, dim_unit2), + (coord_unit1, coord_unit2), + ) = variants.get(variant) + + array1 = np.linspace(0, 5, 10).astype(dtype) * data_unit1 + array2 = np.linspace(-5, 0, 5).astype(dtype) * data_unit2 - array1 = np.linspace(0, 5, 10).astype(dtype) * unit_registry.m - array2 = np.linspace(-5, 0, 5).astype(dtype) * data_unit - x1 = np.arange(5, 15) * original_unit - x2 = np.arange(5) * dims_unit + x1 = np.arange(5, 15) * dim_unit1 + x2 = np.arange(5) * dim_unit2 + + u1 = np.linspace(1, 2, 10).astype(dtype) * coord_unit1 + u2 = np.linspace(0, 1, 5).astype(dtype) * coord_unit2 - arr1 = xr.DataArray(data=array1, coords={"x": x1}, dims="x") - arr2 = xr.DataArray(data=array2, coords={"x": x2}, dims="x") + arr1 = xr.DataArray(data=array1, coords={"x": x1, "u": ("x", u1)}, dims="x") + arr2 = xr.DataArray(data=array2, coords={"x": x2, "u": ("x", u2)}, dims="x") if error is not None: with pytest.raises(error): @@ -881,22 +966,37 @@ def test_concat_dataarray(variant, unit, error, dtype): "variant", ( "data", - pytest.param("dims", marks=pytest.mark.xfail(reason="indexes strip units")), + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", ), ) def test_concat_dataset(variant, unit, error, dtype): original_unit = unit_registry.m - variants = {"data": (unit, original_unit), "dims": (original_unit, unit)} - data_unit, dims_unit = variants.get(variant) + variants = { + "data": ((original_unit, unit), (1, 1), (1, 1)), + "dims": ((1, 1), (original_unit, unit), (1, 1)), + "coords": ((1, 1), (1, 1), (original_unit, unit)), + } + ( + (data_unit1, data_unit2), + (dim_unit1, dim_unit2), + (coord_unit1, coord_unit2), + ) = variants.get(variant) - array1 = np.linspace(0, 5, 10).astype(dtype) * unit_registry.m - array2 = np.linspace(-5, 0, 5).astype(dtype) * data_unit - x1 = np.arange(5, 15) * original_unit - x2 = np.arange(5) * dims_unit + array1 = np.linspace(0, 5, 10).astype(dtype) * data_unit1 + array2 = np.linspace(-5, 0, 5).astype(dtype) * data_unit2 - ds1 = xr.Dataset(data_vars={"a": ("x", array1)}, coords={"x": x1}) - ds2 = xr.Dataset(data_vars={"a": ("x", array2)}, coords={"x": x2}) + x1 = np.arange(5, 15) * dim_unit1 + x2 = np.arange(5) * dim_unit2 + + u1 = np.linspace(1, 2, 10).astype(dtype) * coord_unit1 + u2 = np.linspace(0, 1, 5).astype(dtype) * coord_unit2 + + ds1 = xr.Dataset(data_vars={"a": ("x", array1)}, coords={"x": x1, "u": ("x", u1)}) + ds2 = xr.Dataset(data_vars={"a": ("x", array2)}, coords={"x": x2, "u": ("x", u2)}) if error is not None: with pytest.raises(error): @@ -915,10 +1015,6 @@ def test_concat_dataset(variant, unit, error, dtype): assert_identical(expected, actual) -# TODO: remove once pint==0.12 has been released -@pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose" -) @pytest.mark.parametrize( "unit,error", ( @@ -936,7 +1032,9 @@ def test_concat_dataset(variant, unit, error, dtype): "variant", ( "data", - pytest.param("dims", marks=pytest.mark.xfail(reason="indexes strip units")), + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), "coords", ), ) @@ -944,29 +1042,33 @@ def test_merge_dataarray(variant, unit, error, dtype): original_unit = unit_registry.m variants = { - "data": (unit, original_unit, original_unit), - "dims": (original_unit, unit, original_unit), - "coords": (original_unit, original_unit, unit), + "data": ((original_unit, unit), (1, 1), (1, 1)), + "dims": ((1, 1), (original_unit, unit), (1, 1)), + "coords": ((1, 1), (1, 1), (original_unit, unit)), } - data_unit, dim_unit, coord_unit = variants.get(variant) - - array1 = np.linspace(0, 1, 2 * 3).reshape(2, 3).astype(dtype) * original_unit - x1 = np.arange(2) * original_unit - y1 = np.arange(3) * original_unit - u1 = np.linspace(10, 20, 2) * original_unit - v1 = np.linspace(10, 20, 3) * original_unit - - array2 = np.linspace(1, 2, 2 * 4).reshape(2, 4).astype(dtype) * data_unit - x2 = np.arange(2, 4) * dim_unit - z2 = np.arange(4) * original_unit - u2 = np.linspace(20, 30, 2) * coord_unit - w2 = np.linspace(10, 20, 4) * original_unit - - array3 = np.linspace(0, 2, 3 * 4).reshape(3, 4).astype(dtype) * data_unit - y3 = np.arange(3, 6) * dim_unit - z3 = np.arange(4, 8) * dim_unit - v3 = np.linspace(10, 20, 3) * coord_unit - w3 = np.linspace(10, 20, 4) * coord_unit + ( + (data_unit1, data_unit2), + (dim_unit1, dim_unit2), + (coord_unit1, coord_unit2), + ) = variants.get(variant) + + array1 = np.linspace(0, 1, 2 * 3).reshape(2, 3).astype(dtype) * data_unit1 + x1 = np.arange(2) * dim_unit1 + y1 = np.arange(3) * dim_unit1 + u1 = np.linspace(10, 20, 2) * coord_unit1 + v1 = np.linspace(10, 20, 3) * coord_unit1 + + array2 = np.linspace(1, 2, 2 * 4).reshape(2, 4).astype(dtype) * data_unit2 + x2 = np.arange(2, 4) * dim_unit2 + z2 = np.arange(4) * dim_unit1 + u2 = np.linspace(20, 30, 2) * coord_unit2 + w2 = np.linspace(10, 20, 4) * coord_unit1 + + array3 = np.linspace(0, 2, 3 * 4).reshape(3, 4).astype(dtype) * data_unit2 + y3 = np.arange(3, 6) * dim_unit2 + z3 = np.arange(4, 8) * dim_unit2 + v3 = np.linspace(10, 20, 3) * coord_unit2 + w3 = np.linspace(10, 20, 4) * coord_unit2 arr1 = xr.DataArray( name="a", @@ -993,31 +1095,22 @@ def test_merge_dataarray(variant, unit, error, dtype): return - units = {name: original_unit for name in list("axyzuvw")} - - convert_and_strip = lambda arr: strip_units(convert_units(arr, units)) - expected_units = { - "a": original_unit, - "u": original_unit, - "v": original_unit, - "w": original_unit, - "x": original_unit, - "y": original_unit, - "z": original_unit, + units = { + "a": data_unit1, + "u": coord_unit1, + "v": coord_unit1, + "w": coord_unit1, + "x": dim_unit1, + "y": dim_unit1, + "z": dim_unit1, } + convert_and_strip = lambda arr: strip_units(convert_units(arr, units)) - expected = convert_units( - attach_units( - xr.merge( - [ - convert_and_strip(arr1), - convert_and_strip(arr2), - convert_and_strip(arr3), - ] - ), - units, + expected = attach_units( + xr.merge( + [convert_and_strip(arr1), convert_and_strip(arr2), convert_and_strip(arr3)] ), - expected_units, + units, ) actual = xr.merge([arr1, arr2, arr3]) @@ -1026,10 +1119,6 @@ def test_merge_dataarray(variant, unit, error, dtype): assert_allclose(expected, actual) -# TODO: remove once pint==0.12 has been released -@pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose" -) @pytest.mark.parametrize( "unit,error", ( @@ -1047,7 +1136,9 @@ def test_merge_dataarray(variant, unit, error, dtype): "variant", ( "data", - pytest.param("dims", marks=pytest.mark.xfail(reason="indexes strip units")), + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), "coords", ), ) @@ -1055,43 +1146,47 @@ def test_merge_dataset(variant, unit, error, dtype): original_unit = unit_registry.m variants = { - "data": (unit, original_unit, original_unit), - "dims": (original_unit, unit, original_unit), - "coords": (original_unit, original_unit, unit), + "data": ((original_unit, unit), (1, 1), (1, 1)), + "dims": ((1, 1), (original_unit, unit), (1, 1)), + "coords": ((1, 1), (1, 1), (original_unit, unit)), } - data_unit, dim_unit, coord_unit = variants.get(variant) + ( + (data_unit1, data_unit2), + (dim_unit1, dim_unit2), + (coord_unit1, coord_unit2), + ) = variants.get(variant) - array1 = np.zeros(shape=(2, 3), dtype=dtype) * original_unit - array2 = np.zeros(shape=(2, 3), dtype=dtype) * original_unit + array1 = np.zeros(shape=(2, 3), dtype=dtype) * data_unit1 + array2 = np.zeros(shape=(2, 3), dtype=dtype) * data_unit1 - x = np.arange(11, 14) * original_unit - y = np.arange(2) * original_unit - z = np.arange(3) * original_unit + x = np.arange(11, 14) * dim_unit1 + y = np.arange(2) * dim_unit1 + u = np.arange(3) * coord_unit1 ds1 = xr.Dataset( data_vars={"a": (("y", "x"), array1), "b": (("y", "x"), array2)}, - coords={"x": x, "y": y, "u": ("x", z)}, + coords={"x": x, "y": y, "u": ("x", u)}, ) ds2 = xr.Dataset( data_vars={ - "a": (("y", "x"), np.ones_like(array1) * data_unit), - "b": (("y", "x"), np.ones_like(array2) * data_unit), + "a": (("y", "x"), np.ones_like(array1) * data_unit2), + "b": (("y", "x"), np.ones_like(array2) * data_unit2), }, coords={ - "x": np.arange(3) * dim_unit, - "y": np.arange(2, 4) * dim_unit, - "u": ("x", np.arange(-3, 0) * coord_unit), + "x": np.arange(3) * dim_unit2, + "y": np.arange(2, 4) * dim_unit2, + "u": ("x", np.arange(-3, 0) * coord_unit2), }, ) ds3 = xr.Dataset( data_vars={ - "a": (("y", "x"), np.full_like(array1, np.nan) * data_unit), - "b": (("y", "x"), np.full_like(array2, np.nan) * data_unit), + "a": (("y", "x"), np.full_like(array1, np.nan) * data_unit2), + "b": (("y", "x"), np.full_like(array2, np.nan) * data_unit2), }, coords={ - "x": np.arange(3, 6) * dim_unit, - "y": np.arange(4, 6) * dim_unit, - "u": ("x", np.arange(3, 6) * coord_unit), + "x": np.arange(3, 6) * dim_unit2, + "y": np.arange(4, 6) * dim_unit2, + "u": ("x", np.arange(3, 6) * coord_unit2), }, ) @@ -1104,15 +1199,9 @@ def test_merge_dataset(variant, unit, error, dtype): units = extract_units(ds1) convert_and_strip = lambda ds: strip_units(convert_units(ds, units)) - expected_units = {name: original_unit for name in list("abxyzu")} - expected = convert_units( - attach_units( - func( - [convert_and_strip(ds1), convert_and_strip(ds2), convert_and_strip(ds3)] - ), - units, - ), - expected_units, + expected = attach_units( + func([convert_and_strip(ds1), convert_and_strip(ds2), convert_and_strip(ds3)]), + units, ) actual = func([ds1, ds2, ds3]) @@ -1120,35 +1209,79 @@ def test_merge_dataset(variant, unit, error, dtype): assert_allclose(expected, actual) +@pytest.mark.parametrize( + "variant", + ( + "data", + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", + ), +) @pytest.mark.parametrize("func", (xr.zeros_like, xr.ones_like)) -def test_replication_dataarray(func, dtype): - array = np.linspace(0, 10, 20).astype(dtype) * unit_registry.s - data_array = xr.DataArray(data=array, dims="x") +def test_replication_dataarray(func, variant, dtype): + unit = unit_registry.m + + variants = { + "data": (unit, 1, 1), + "dims": (1, unit, 1), + "coords": (1, 1, unit), + } + data_unit, dim_unit, coord_unit = variants.get(variant) - numpy_func = getattr(np, func.__name__) - units = extract_units(numpy_func(data_array)) - expected = attach_units(func(data_array), units) + array = np.linspace(0, 10, 20).astype(dtype) * data_unit + x = np.arange(20) * dim_unit + u = np.linspace(0, 1, 20) * coord_unit + + data_array = xr.DataArray(data=array, dims="x", coords={"x": x, "u": ("x", u)}) + units = extract_units(data_array) + units.pop(data_array.name) + + expected = attach_units(func(strip_units(data_array)), units) actual = func(data_array) assert_units_equal(expected, actual) assert_identical(expected, actual) +@pytest.mark.parametrize( + "variant", + ( + "data", + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", + ), +) @pytest.mark.parametrize("func", (xr.zeros_like, xr.ones_like)) -def test_replication_dataset(func, dtype): - array1 = np.linspace(0, 10, 20).astype(dtype) * unit_registry.s - array2 = np.linspace(5, 10, 10).astype(dtype) * unit_registry.Pa - x = np.arange(20).astype(dtype) * unit_registry.m - y = np.arange(10).astype(dtype) * unit_registry.m - z = y.to(unit_registry.mm) +def test_replication_dataset(func, variant, dtype): + unit = unit_registry.m + + variants = { + "data": ((unit_registry.m, unit_registry.Pa), 1, 1), + "dims": ((1, 1), unit, 1), + "coords": ((1, 1), 1, unit), + } + (data_unit1, data_unit2), dim_unit, coord_unit = variants.get(variant) + + array1 = np.linspace(0, 10, 20).astype(dtype) * data_unit1 + array2 = np.linspace(5, 10, 10).astype(dtype) * data_unit2 + x = np.arange(20).astype(dtype) * dim_unit + y = np.arange(10).astype(dtype) * dim_unit + u = np.linspace(0, 1, 10) * coord_unit ds = xr.Dataset( data_vars={"a": ("x", array1), "b": ("y", array2)}, - coords={"x": x, "y": y, "z": ("y", z)}, + coords={"x": x, "y": y, "u": ("y", u)}, ) + units = { + name: unit + for name, unit in extract_units(ds).items() + if name not in ds.data_vars + } - numpy_func = getattr(np, func.__name__) - units = extract_units(ds.map(numpy_func)) expected = attach_units(func(strip_units(ds)), units) actual = func(ds) @@ -1157,37 +1290,40 @@ def test_replication_dataset(func, dtype): assert_identical(expected, actual) -@pytest.mark.xfail( - reason=( - "pint is undecided on how `full_like` should work, so incorrect errors " - "may be expected: hgrecco/pint#882" - ) -) @pytest.mark.parametrize( - "unit,error", + "variant", ( - pytest.param(1, DimensionalityError, id="no_unit"), + "data", pytest.param( - unit_registry.dimensionless, DimensionalityError, id="dimensionless" + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + pytest.param( + "coords", + marks=pytest.mark.xfail(reason="can't copy quantity into non-quantity"), ), - pytest.param(unit_registry.m, DimensionalityError, id="incompatible_unit"), - pytest.param(unit_registry.ms, None, id="compatible_unit"), - pytest.param(unit_registry.s, None, id="identical_unit"), ), - ids=repr, ) -def test_replication_full_like_dataarray(unit, error, dtype): - array = np.linspace(0, 5, 10) * unit_registry.s - data_array = xr.DataArray(data=array, dims="x") +def test_replication_full_like_dataarray(variant, dtype): + # since full_like will strip units and then use the units of the + # fill value, we don't need to try multiple units + unit = unit_registry.m - fill_value = -1 * unit - if error is not None: - with pytest.raises(error): - xr.full_like(data_array, fill_value=fill_value) + variants = { + "data": (unit, 1, 1), + "dims": (1, unit, 1), + "coords": (1, 1, unit), + } + data_unit, dim_unit, coord_unit = variants.get(variant) - return + array = np.linspace(0, 5, 10) * data_unit + x = np.arange(10) * dim_unit + u = np.linspace(0, 1, 10) * coord_unit + data_array = xr.DataArray(data=array, dims="x", coords={"x": x, "u": ("x", u)}) + + fill_value = -1 * unit_registry.degK - units = {**extract_units(data_array), **{None: unit if unit != 1 else None}} + units = extract_units(data_array) + units[data_array.name] = fill_value.units expected = attach_units( xr.full_like(strip_units(data_array), fill_value=strip_units(fill_value)), units ) @@ -1197,47 +1333,46 @@ def test_replication_full_like_dataarray(unit, error, dtype): assert_identical(expected, actual) -@pytest.mark.xfail( - reason=( - "pint is undecided on how `full_like` should work, so incorrect errors " - "may be expected: hgrecco/pint#882" - ) -) @pytest.mark.parametrize( - "unit,error", + "variant", ( - pytest.param(1, DimensionalityError, id="no_unit"), + "data", pytest.param( - unit_registry.dimensionless, DimensionalityError, id="dimensionless" + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + pytest.param( + "coords", + marks=pytest.mark.xfail(reason="can't copy quantity into non-quantity"), ), - pytest.param(unit_registry.m, DimensionalityError, id="incompatible_unit"), - pytest.param(unit_registry.ms, None, id="compatible_unit"), - pytest.param(unit_registry.s, None, id="identical_unit"), ), - ids=repr, ) -def test_replication_full_like_dataset(unit, error, dtype): - array1 = np.linspace(0, 10, 20).astype(dtype) * unit_registry.s - array2 = np.linspace(5, 10, 10).astype(dtype) * unit_registry.Pa - x = np.arange(20).astype(dtype) * unit_registry.m - y = np.arange(10).astype(dtype) * unit_registry.m - z = y.to(unit_registry.mm) +def test_replication_full_like_dataset(variant, dtype): + unit = unit_registry.m + + variants = { + "data": ((unit_registry.s, unit_registry.Pa), 1, 1), + "dims": ((1, 1), unit, 1), + "coords": ((1, 1), 1, unit), + } + (data_unit1, data_unit2), dim_unit, coord_unit = variants.get(variant) + + array1 = np.linspace(0, 10, 20).astype(dtype) * data_unit1 + array2 = np.linspace(5, 10, 10).astype(dtype) * data_unit2 + x = np.arange(20).astype(dtype) * dim_unit + y = np.arange(10).astype(dtype) * dim_unit + + u = np.linspace(0, 1, 10) * coord_unit ds = xr.Dataset( data_vars={"a": ("x", array1), "b": ("y", array2)}, - coords={"x": x, "y": y, "z": ("y", z)}, + coords={"x": x, "y": y, "u": ("y", u)}, ) - fill_value = -1 * unit - if error is not None: - with pytest.raises(error): - xr.full_like(ds, fill_value=fill_value) - - return + fill_value = -1 * unit_registry.degK units = { **extract_units(ds), - **{name: unit if unit != 1 else None for name in ds.data_vars}, + **{name: unit_registry.degK for name in ds.data_vars}, } expected = attach_units( xr.full_like(strip_units(ds), fill_value=strip_units(fill_value)), units @@ -1308,10 +1443,9 @@ def test_where_dataarray(fill_value, unit, error, dtype): def test_where_dataset(fill_value, unit, error, dtype): array1 = np.linspace(0, 5, 10).astype(dtype) * unit_registry.m array2 = np.linspace(-5, 0, 10).astype(dtype) * unit_registry.m - x = np.arange(10) * unit_registry.s - ds = xr.Dataset(data_vars={"a": ("x", array1), "b": ("x", array2)}, coords={"x": x}) - cond = x < 5 * unit_registry.s + ds = xr.Dataset(data_vars={"a": ("x", array1), "b": ("x", array2)}) + cond = array1 < 2 * unit_registry.m fill_value = fill_value * unit if error is not None and not ( @@ -1358,61 +1492,7 @@ def test_dot_dataarray(dtype): assert_identical(expected, actual) -def delete_attrs(*to_delete): - def wrapper(cls): - for item in to_delete: - setattr(cls, item, None) - - return cls - - return wrapper - - -@delete_attrs( - "test_getitem_with_mask", - "test_getitem_with_mask_nd_indexer", - "test_index_0d_string", - "test_index_0d_datetime", - "test_index_0d_timedelta64", - "test_0d_time_data", - "test_index_0d_not_a_time", - "test_datetime64_conversion", - "test_timedelta64_conversion", - "test_pandas_period_index", - "test_1d_reduce", - "test_array_interface", - "test___array__", - "test_copy_index", - "test_concat_number_strings", - "test_concat_fixed_len_str", - "test_concat_mixed_dtypes", - "test_pandas_datetime64_with_tz", - "test_pandas_data", - "test_multiindex", -) -class TestVariable(VariableSubclassobjects): - @staticmethod - def cls(dims, data, *args, **kwargs): - return xr.Variable( - dims, unit_registry.Quantity(data, unit_registry.m), *args, **kwargs - ) - - def example_1d_objects(self): - for data in [ - range(3), - 0.5 * np.arange(3), - 0.5 * np.arange(3, dtype=np.float32), - np.array(["a", "b", "c"], dtype=object), - ]: - yield (self.cls("x", data), data) - - # TODO: remove once pint==0.12 has been released - @pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose" - ) - def test_real_and_imag(self): - super().test_real_and_imag() - +class TestVariable: @pytest.mark.parametrize( "func", ( @@ -1454,22 +1534,14 @@ def test_aggregation(self, func, dtype): assert_units_equal(expected, actual) assert_allclose(expected, actual) - # TODO: remove once pint==0.12 has been released - @pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose" - ) def test_aggregate_complex(self): variable = xr.Variable("x", [1, 2j, np.nan] * unit_registry.m) expected = xr.Variable((), (0.5 + 1j) * unit_registry.m) actual = variable.mean() assert_units_equal(expected, actual) - xr.testing.assert_allclose(expected, actual) + assert_allclose(expected, actual) - # TODO: remove once pint==0.12 has been released - @pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose" - ) @pytest.mark.parametrize( "func", ( @@ -1526,7 +1598,7 @@ def test_numpy_methods(self, func, unit, error, dtype): actual = func(variable, *args, **kwargs) assert_units_equal(expected, actual) - xr.testing.assert_allclose(expected, actual) + assert_allclose(expected, actual) @pytest.mark.parametrize( "func", (method("item", 5), method("searchsorted", 5)), ids=repr @@ -1586,7 +1658,7 @@ def test_raw_numpy_methods(self, func, unit, error, dtype): actual = func(variable, *args, **kwargs) assert_units_equal(expected, actual) - np.testing.assert_allclose(expected, actual) + assert_duckarray_allclose(expected, actual) @pytest.mark.parametrize( "func", (method("isnull"), method("notnull"), method("count")), ids=repr @@ -1609,7 +1681,7 @@ def test_missing_value_detection(self, func): actual = func(variable) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "unit,error", @@ -1655,7 +1727,7 @@ def test_missing_value_fillna(self, unit, error): actual = variable.fillna(value=fill_value) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "unit", @@ -1766,12 +1838,8 @@ def test_isel(self, indices, dtype): actual = variable.isel(x=indices) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) - # TODO: remove once pint==0.12 has been released - @pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose" - ) @pytest.mark.parametrize( "unit,error", ( @@ -1828,7 +1896,7 @@ def test_1d_math(self, func, unit, error, dtype): actual = func(variable, y) assert_units_equal(expected, actual) - xr.testing.assert_allclose(expected, actual) + assert_allclose(expected, actual) @pytest.mark.parametrize( "unit,error", @@ -1877,43 +1945,30 @@ def test_masking(self, func, unit, error, dtype): actual = func(variable, cond, other) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) - def test_squeeze(self, dtype): + @pytest.mark.parametrize("dim", ("x", "y", "z", "t", "all")) + def test_squeeze(self, dim, dtype): shape = (2, 1, 3, 1, 1, 2) names = list("abcdef") + dim_lengths = dict(zip(names, shape)) array = np.ones(shape=shape) * unit_registry.m variable = xr.Variable(names, array) + kwargs = {"dim": dim} if dim != "all" and dim_lengths.get(dim, 0) == 1 else {} expected = attach_units( - strip_units(variable).squeeze(), extract_units(variable) + strip_units(variable).squeeze(**kwargs), extract_units(variable) ) - actual = variable.squeeze() + actual = variable.squeeze(**kwargs) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) - - names = tuple(name for name, size in zip(names, shape) if shape == 1) - for name in names: - expected = attach_units( - strip_units(variable).squeeze(dim=name), extract_units(variable) - ) - actual = variable.squeeze(dim=name) - - assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "func", ( method("coarsen", windows={"y": 2}, func=np.mean), - pytest.param( - method("quantile", q=[0.25, 0.75]), - marks=pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", - reason="quantile / nanquantile not implemented yet", - ), - ), + method("quantile", q=[0.25, 0.75]), pytest.param( method("rank", dim="x"), marks=pytest.mark.xfail(reason="rank not implemented for non-ndarray"), @@ -1940,7 +1995,7 @@ def test_computation(self, func, dtype): actual = func(variable) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "unit,error", @@ -1986,7 +2041,7 @@ def test_stack(self, dtype): actual = variable.stack(z=("x", "y")) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) def test_unstack(self, dtype): array = np.linspace(0, 5, 3 * 10).astype(dtype) * unit_registry.m @@ -1998,7 +2053,7 @@ def test_unstack(self, dtype): actual = variable.unstack(z={"x": 3, "y": 10}) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "unit,error", @@ -2038,7 +2093,7 @@ def test_concat(self, unit, error, dtype): actual = xr.Variable.concat([variable, other], dim="y") assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) def test_set_dims(self, dtype): array = np.linspace(0, 5, 3 * 10).reshape(3, 10).astype(dtype) * unit_registry.m @@ -2051,7 +2106,7 @@ def test_set_dims(self, dtype): actual = variable.set_dims(dims) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) def test_copy(self, dtype): array = np.linspace(0, 5, 10).astype(dtype) * unit_registry.m @@ -2064,7 +2119,7 @@ def test_copy(self, dtype): actual = variable.copy(data=other) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "unit", @@ -2105,45 +2160,43 @@ def test_no_conflicts(self, unit, dtype): assert expected == actual + @pytest.mark.parametrize( + "mode", + [ + "constant", + "mean", + "median", + "reflect", + "edge", + pytest.param( + "linear_ramp", + marks=pytest.mark.xfail( + reason="pint bug: https://github.com/hgrecco/pint/issues/1026" + ), + ), + "maximum", + "minimum", + "symmetric", + "wrap", + ], + ) @pytest.mark.parametrize("xr_arg, np_arg", _PAD_XR_NP_ARGS) - def test_pad_constant_values(self, dtype, xr_arg, np_arg): - data = np.arange(4 * 3 * 2).reshape(4, 3, 2).astype(dtype) * unit_registry.m + def test_pad(self, mode, xr_arg, np_arg): + data = np.arange(4 * 3 * 2).reshape(4, 3, 2) * unit_registry.m v = xr.Variable(["x", "y", "z"], data) - actual = v.pad(**xr_arg, mode="constant") - expected = xr.Variable( - v.dims, - np.pad( - v.data.astype(float), np_arg, mode="constant", constant_values=np.nan, - ), + expected = attach_units( + strip_units(v).pad(mode=mode, **xr_arg), extract_units(v), ) - xr.testing.assert_identical(expected, actual) - assert_units_equal(expected, actual) - assert isinstance(actual._data, type(v._data)) + actual = v.pad(mode=mode, **xr_arg) - # for the boolean array, we pad False - data = np.full_like(data, False, dtype=bool).reshape(4, 3, 2) - v = xr.Variable(["x", "y", "z"], data) - actual = v.pad(**xr_arg, mode="constant", constant_values=data.flat[0]) - expected = xr.Variable( - v.dims, - np.pad(v.data, np_arg, mode="constant", constant_values=v.data.flat[0]), - ) - xr.testing.assert_identical(actual, expected) assert_units_equal(expected, actual) + assert_equal(actual, expected) @pytest.mark.parametrize( "unit,error", ( - pytest.param( - 1, - DimensionalityError, - id="no_unit", - marks=pytest.mark.xfail( - LooseVersion(pint.__version__) < LooseVersion("0.10.2"), - reason="bug in pint's implementation of np.pad", - ), - ), + pytest.param(1, DimensionalityError, id="no_unit"), pytest.param( unit_registry.dimensionless, DimensionalityError, id="dimensionless" ), @@ -2176,17 +2229,16 @@ def test_pad_unit_constant_value(self, unit, error, dtype): actual = func(variable, constant_values=fill_value) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) class TestDataArray: - @pytest.mark.filterwarnings("error:::pint[.*]") @pytest.mark.parametrize( "variant", ( pytest.param( "with_dims", - marks=pytest.mark.xfail(reason="units in indexes are not supported"), + marks=pytest.mark.xfail(reason="indexes don't support units"), ), "with_coords", "without_coords", @@ -2215,7 +2267,6 @@ def test_init(self, variant, dtype): }.values() ) - @pytest.mark.filterwarnings("error:::pint[.*]") @pytest.mark.parametrize( "func", (pytest.param(str, id="str"), pytest.param(repr, id="repr")) ) @@ -2224,7 +2275,7 @@ def test_init(self, variant, dtype): ( pytest.param( "with_dims", - marks=pytest.mark.xfail(reason="units in indexes are not supported"), + marks=pytest.mark.xfail(reason="indexes don't support units"), ), pytest.param("with_coords"), pytest.param("without_coords"), @@ -2248,10 +2299,6 @@ def test_repr(self, func, variant, dtype): # warnings or errors, but does not check the result func(data_array) - # TODO: remove once pint==0.12 has been released - @pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose", - ) @pytest.mark.parametrize( "func", ( @@ -2345,7 +2392,7 @@ def test_unary_operations(self, func, dtype): actual = func(data_array) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "func", @@ -2365,14 +2412,21 @@ def test_binary_operations(self, func, dtype): actual = func(data_array) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "comparison", ( pytest.param(operator.lt, id="less_than"), pytest.param(operator.ge, id="greater_equal"), - pytest.param(operator.eq, id="equal"), + pytest.param( + operator.eq, + id="equal", + marks=pytest.mark.xfail( + # LooseVersion(pint.__version__) < "0.14", + reason="inconsistencies in the return values of pint's eq", + ), + ), ), ) @pytest.mark.parametrize( @@ -2416,7 +2470,7 @@ def test_comparison_operations(self, comparison, unit, error, dtype): ) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "units,error", @@ -2445,7 +2499,7 @@ def test_univariate_ufunc(self, units, error, dtype): actual = func(data_array) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.xfail(reason="needs the type register system for __array_ufunc__") @pytest.mark.parametrize( @@ -2487,11 +2541,11 @@ def test_bivariate_ufunc(self, unit, error, dtype): actual = np.maximum(data_array, 1 * unit) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) actual = np.maximum(1 * unit, data_array) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize("property", ("T", "imag", "real")) def test_numpy_properties(self, property, dtype): @@ -2508,7 +2562,7 @@ def test_numpy_properties(self, property, dtype): actual = getattr(data_array, property) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "func", @@ -2524,7 +2578,7 @@ def test_numpy_methods(self, func, dtype): actual = func(data_array) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) def test_item(self, dtype): array = np.arange(10).astype(dtype) * unit_registry.m @@ -2535,7 +2589,7 @@ def test_item(self, dtype): expected = func(strip_units(data_array)) * unit_registry.m actual = func(data_array) - np.testing.assert_allclose(expected, actual) + assert_duckarray_allclose(expected, actual) @pytest.mark.parametrize( "unit,error", @@ -2650,7 +2704,7 @@ def test_numpy_methods_with_args(self, func, unit, error, dtype): actual = func(data_array, *args, **kwargs) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "func", (method("isnull"), method("notnull"), method("count")), ids=repr @@ -2673,7 +2727,7 @@ def test_missing_value_detection(self, func, dtype): actual = func(data_array) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.xfail(reason="ffill and bfill lose units in data") @pytest.mark.parametrize("func", (method("ffill"), method("bfill")), ids=repr) @@ -2691,7 +2745,7 @@ def test_missing_value_filling(self, func, dtype): actual = func(data_array, dim="x") assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "unit,error", @@ -2740,7 +2794,7 @@ def test_fillna(self, fill_value, unit, error, dtype): actual = func(data_array, value=value) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) def test_dropna(self, dtype): array = ( @@ -2755,7 +2809,7 @@ def test_dropna(self, dtype): actual = data_array.dropna(dim="x") assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "unit", @@ -2784,7 +2838,7 @@ def test_isin(self, unit, dtype): actual = data_array.isin(values) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "variant", ("masking", "replacing_scalar", "replacing_array", "dropping") @@ -2838,7 +2892,7 @@ def test_where(self, variant, unit, error, dtype): actual = data_array.where(**kwargs) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.xfail(reason="uses numpy.vectorize") def test_interpolate_na(self): @@ -2854,7 +2908,7 @@ def test_interpolate_na(self): actual = data_array.interpolate_na(dim="x") assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "unit,error", @@ -2895,7 +2949,7 @@ def test_combine_first(self, unit, error, dtype): actual = data_array.combine_first(other) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "unit", @@ -2984,17 +3038,47 @@ def is_compatible(a, b): pytest.param(unit_registry.m, id="identical_unit"), ), ) - def test_broadcast_like(self, unit, dtype): - array1 = np.linspace(1, 2, 2 * 1).reshape(2, 1).astype(dtype) * unit_registry.Pa - array2 = np.linspace(0, 1, 2 * 3).reshape(2, 3).astype(dtype) * unit_registry.Pa + @pytest.mark.parametrize( + "variant", + ( + "data", + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", + ), + ) + def test_broadcast_like(self, variant, unit, dtype): + original_unit = unit_registry.m + + variants = { + "data": ((original_unit, unit), (1, 1), (1, 1)), + "dims": ((1, 1), (original_unit, unit), (1, 1)), + "coords": ((1, 1), (1, 1), (original_unit, unit)), + } + ( + (data_unit1, data_unit2), + (dim_unit1, dim_unit2), + (coord_unit1, coord_unit2), + ) = variants.get(variant) + + array1 = np.linspace(1, 2, 2 * 1).reshape(2, 1).astype(dtype) * data_unit1 + array2 = np.linspace(0, 1, 2 * 3).reshape(2, 3).astype(dtype) * data_unit2 + + x1 = np.arange(2) * dim_unit1 + x2 = np.arange(2) * dim_unit2 + y1 = np.array([0]) * dim_unit1 + y2 = np.arange(3) * dim_unit2 - x1 = np.arange(2) * unit_registry.m - x2 = np.arange(2) * unit - y1 = np.array([0]) * unit_registry.m - y2 = np.arange(3) * unit + u1 = np.linspace(0, 1, 2) * coord_unit1 + u2 = np.linspace(0, 1, 2) * coord_unit2 - arr1 = xr.DataArray(data=array1, coords={"x": x1, "y": y1}, dims=("x", "y")) - arr2 = xr.DataArray(data=array2, coords={"x": x2, "y": y2}, dims=("x", "y")) + arr1 = xr.DataArray( + data=array1, coords={"x": x1, "y": y1, "u": ("x", u1)}, dims=("x", "y") + ) + arr2 = xr.DataArray( + data=array2, coords={"x": x2, "y": y2, "u": ("x", u2)}, dims=("x", "y") + ) expected = attach_units( strip_units(arr1).broadcast_like(strip_units(arr2)), extract_units(arr1) @@ -3002,7 +3086,7 @@ def test_broadcast_like(self, unit, dtype): actual = arr1.broadcast_like(arr2) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "unit", @@ -3032,56 +3116,89 @@ def test_broadcast_equals(self, unit, dtype): assert expected == actual + def test_pad(self, dtype): + array = np.linspace(0, 5, 10).astype(dtype) * unit_registry.m + + data_array = xr.DataArray(data=array, dims="x") + units = extract_units(data_array) + + expected = attach_units(strip_units(data_array).pad(x=(2, 3)), units) + actual = data_array.pad(x=(2, 3)) + + assert_units_equal(expected, actual) + assert_equal(expected, actual) + + @pytest.mark.parametrize( + "variant", + ( + "data", + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", + ), + ) @pytest.mark.parametrize( "func", ( method("pipe", lambda da: da * 10), - method("assign_coords", y2=("y", np.arange(10) * unit_registry.mm)), + method("assign_coords", w=("y", np.arange(10) * unit_registry.mm)), method("assign_attrs", attr1="value"), - method("rename", x2="x_mm"), - method("swap_dims", {"x": "x2"}), - method( - "expand_dims", - dim={"z": np.linspace(10, 20, 12) * unit_registry.s}, - axis=1, + method("rename", u="v"), + pytest.param( + method("swap_dims", {"x": "u"}), + marks=pytest.mark.xfail(reason="indexes don't support units"), + ), + pytest.param( + method( + "expand_dims", + dim={"z": np.linspace(10, 20, 12) * unit_registry.s}, + axis=1, + ), + marks=pytest.mark.xfail(reason="indexes don't support units"), ), method("drop_vars", "x"), - method("reset_coords", names="x2"), + method("reset_coords", names="u"), method("copy"), method("astype", np.float32), ), ids=repr, ) - def test_content_manipulation(self, func, dtype): - quantity = ( - np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) - * unit_registry.pascal - ) - x = np.arange(quantity.shape[0]) * unit_registry.m - y = np.arange(quantity.shape[1]) * unit_registry.m - x2 = x.to(unit_registry.mm) + def test_content_manipulation(self, func, variant, dtype): + unit = unit_registry.m + + variants = { + "data": (unit, 1, 1), + "dims": (1, unit, 1), + "coords": (1, 1, unit), + } + data_unit, dim_unit, coord_unit = variants.get(variant) + + quantity = np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * data_unit + x = np.arange(quantity.shape[0]) * dim_unit + y = np.arange(quantity.shape[1]) * dim_unit + u = np.linspace(0, 1, quantity.shape[0]) * coord_unit data_array = xr.DataArray( - name="data", + name="a", data=quantity, - coords={"x": x, "x2": ("x", x2), "y": y}, + coords={"x": x, "u": ("x", u), "y": y}, dims=("x", "y"), ) stripped_kwargs = { key: array_strip_units(value) for key, value in func.kwargs.items() } - units = {**{"x_mm": x2.units, "x2": x2.units}, **extract_units(data_array)} + units = extract_units(data_array) + units["u"] = getattr(u, "units", None) + units["v"] = getattr(u, "units", None) expected = attach_units(func(strip_units(data_array), **stripped_kwargs), units) actual = func(data_array) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) - @pytest.mark.parametrize( - "func", (pytest.param(method("copy", data=np.arange(20))),), ids=repr - ) @pytest.mark.parametrize( "unit", ( @@ -3090,22 +3207,20 @@ def test_content_manipulation(self, func, dtype): pytest.param(unit_registry.degK, id="with_unit"), ), ) - def test_content_manipulation_with_units(self, func, unit, dtype): + def test_copy(self, unit, dtype): quantity = np.linspace(0, 10, 20, dtype=dtype) * unit_registry.pascal - x = np.arange(len(quantity)) * unit_registry.m - - data_array = xr.DataArray(data=quantity, coords={"x": x}, dims="x") + new_data = np.arange(20) - kwargs = {key: value * unit for key, value in func.kwargs.items()} + data_array = xr.DataArray(data=quantity, dims="x") expected = attach_units( - func(strip_units(data_array)), {None: unit, "x": x.units} + strip_units(data_array).copy(data=new_data), {None: unit} ) - actual = func(data_array, **kwargs) + actual = data_array.copy(data=new_data * unit) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "indices", @@ -3115,10 +3230,10 @@ def test_content_manipulation_with_units(self, func, unit, dtype): ), ) def test_isel(self, indices, dtype): + # TODO: maybe test for units in indexes? array = np.arange(10).astype(dtype) * unit_registry.s - x = np.arange(len(array)) * unit_registry.m - data_array = xr.DataArray(data=array, coords={"x": x}, dims="x") + data_array = xr.DataArray(data=array, dims="x") expected = attach_units( strip_units(data_array).isel(x=indices), extract_units(data_array) @@ -3126,7 +3241,7 @@ def test_isel(self, indices, dtype): actual = data_array.isel(x=indices) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( @@ -3171,7 +3286,7 @@ def test_sel(self, raw_values, unit, error, dtype): actual = data_array.sel(x=values) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( @@ -3216,7 +3331,7 @@ def test_loc(self, raw_values, unit, error, dtype): actual = data_array.loc[{"x": values}] assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( @@ -3261,8 +3376,9 @@ def test_drop_sel(self, raw_values, unit, error, dtype): actual = data_array.drop_sel(x=values) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) + @pytest.mark.parametrize("dim", ("x", "y", "z", "t", "all")) @pytest.mark.parametrize( "shape", ( @@ -3273,36 +3389,22 @@ def test_drop_sel(self, raw_values, unit, error, dtype): pytest.param((1, 10, 1, 20), id="first_and_last_dimension_squeezable"), ), ) - def test_squeeze(self, shape, dtype): + def test_squeeze(self, shape, dim, dtype): + names = "xyzt" + dim_lengths = dict(zip(names, shape)) names = "xyzt" - coords = { - name: np.arange(length).astype(dtype) - * (unit_registry.m if name != "t" else unit_registry.s) - for name, length in zip(names, shape) - } array = np.arange(10 * 20).astype(dtype).reshape(shape) * unit_registry.J - data_array = xr.DataArray( - data=array, coords=coords, dims=tuple(names[: len(shape)]) - ) + data_array = xr.DataArray(data=array, dims=tuple(names[: len(shape)])) + + kwargs = {"dim": dim} if dim != "all" and dim_lengths.get(dim, 0) == 1 else {} expected = attach_units( - strip_units(data_array).squeeze(), extract_units(data_array) + strip_units(data_array).squeeze(**kwargs), extract_units(data_array) ) - actual = data_array.squeeze() + actual = data_array.squeeze(**kwargs) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) - - # try squeezing the dimensions separately - names = tuple(dim for dim, coord in coords.items() if len(coord) == 1) - for index, name in enumerate(names): - expected = attach_units( - strip_units(data_array).squeeze(dim=name), extract_units(data_array) - ) - actual = data_array.squeeze(dim=name) - - assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "func", @@ -3310,14 +3412,10 @@ def test_squeeze(self, shape, dtype): ids=repr, ) def test_head_tail_thin(self, func, dtype): + # TODO: works like isel. Maybe also test units in indexes? array = np.linspace(1, 2, 10 * 5).reshape(10, 5) * unit_registry.degK - coords = { - "x": np.arange(10) * unit_registry.m, - "y": np.arange(5) * unit_registry.m, - } - - data_array = xr.DataArray(data=array, coords=coords, dims=("x", "y")) + data_array = xr.DataArray(data=array, dims=("x", "y")) expected = attach_units( func(strip_units(data_array)), extract_units(data_array) @@ -3325,12 +3423,8 @@ def test_head_tail_thin(self, func, dtype): actual = func(data_array) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) - # TODO: remove once pint==0.12 has been released - @pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose" - ) @pytest.mark.parametrize("variant", ("data", "coords")) @pytest.mark.parametrize( "func", @@ -3361,7 +3455,7 @@ def test_interp_reindex(self, variant, func, dtype): actual = func(data_array, x=new_x) assert_units_equal(expected, actual) - xr.testing.assert_allclose(expected, actual) + assert_allclose(expected, actual) @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( @@ -3402,12 +3496,8 @@ def test_interp_reindex_indexing(self, func, unit, error, dtype): actual = func(data_array, x=new_x) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) - # TODO: remove once pint==0.12 has been released - @pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose" - ) @pytest.mark.parametrize("variant", ("data", "coords")) @pytest.mark.parametrize( "func", @@ -3439,7 +3529,7 @@ def test_interp_reindex_like(self, variant, func, dtype): actual = func(data_array, other) assert_units_equal(expected, actual) - xr.testing.assert_allclose(expected, actual) + assert_allclose(expected, actual) @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( @@ -3482,7 +3572,7 @@ def test_interp_reindex_like_indexing(self, func, unit, error, dtype): actual = func(data_array, other) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "func", @@ -3505,7 +3595,7 @@ def test_stacking_stacked(self, func, dtype): actual = func(stacked) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.xfail(reason="indexes don't support units") def test_to_unstacked_dataset(self, dtype): @@ -3529,7 +3619,7 @@ def test_to_unstacked_dataset(self, dtype): actual = func(data_array) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) @pytest.mark.parametrize( "func", @@ -3564,55 +3654,71 @@ def test_stacking_reordering(self, func, dtype): actual = func(data_array) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) + @pytest.mark.parametrize( + "variant", + ( + "data", + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", + ), + ) @pytest.mark.parametrize( "func", ( method("diff", dim="x"), method("differentiate", coord="x"), method("integrate", dim="x"), - pytest.param( - method("quantile", q=[0.25, 0.75]), - marks=pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", - reason="quantile / nanquantile not implemented yet", - ), - ), + method("quantile", q=[0.25, 0.75]), method("reduce", func=np.sum, dim="x"), pytest.param(lambda x: x.dot(x), id="method_dot"), ), ids=repr, ) - def test_computation(self, func, dtype): - array = ( - np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * unit_registry.m - ) + def test_computation(self, func, variant, dtype): + unit = unit_registry.m - x = np.arange(array.shape[0]) * unit_registry.m - y = np.arange(array.shape[1]) * unit_registry.s + variants = { + "data": (unit, 1, 1), + "dims": (1, unit, 1), + "coords": (1, 1, unit), + } + data_unit, dim_unit, coord_unit = variants.get(variant) - data_array = xr.DataArray(data=array, coords={"x": x, "y": y}, dims=("x", "y")) + array = np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * data_unit + + x = np.arange(array.shape[0]) * dim_unit + y = np.arange(array.shape[1]) * dim_unit + + u = np.linspace(0, 1, array.shape[0]) * coord_unit + + data_array = xr.DataArray( + data=array, coords={"x": x, "y": y, "u": ("x", u)}, dims=("x", "y") + ) # we want to make sure the output unit is correct - units = { - **extract_units(data_array), - **( - {} - if isinstance(func, (function, method)) - else extract_units(func(array.reshape(-1))) - ), - } + units = extract_units(data_array) + if not isinstance(func, (function, method)): + units.update(extract_units(func(array.reshape(-1)))) expected = attach_units(func(strip_units(data_array)), units) actual = func(data_array) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) - # TODO: remove once pint==0.12 has been released - @pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", reason="pint bug in isclose" + @pytest.mark.parametrize( + "variant", + ( + "data", + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", + ), ) @pytest.mark.parametrize( "func", @@ -3632,25 +3738,37 @@ def test_computation(self, func, dtype): reason="numbagg functions are not supported by pint" ), ), + method("weighted", xr.DataArray(data=np.linspace(0, 1, 10), dims="y")), ), ids=repr, ) - def test_computation_objects(self, func, dtype): - array = ( - np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * unit_registry.m - ) + def test_computation_objects(self, func, variant, dtype): + unit = unit_registry.m - x = np.array([0, 0, 1, 2, 2]) * unit_registry.m - y = np.arange(array.shape[1]) * 3 * unit_registry.s + variants = { + "data": (unit, 1, 1), + "dims": (1, unit, 1), + "coords": (1, 1, unit), + } + data_unit, dim_unit, coord_unit = variants.get(variant) + + array = np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * data_unit + + x = np.array([0, 0, 1, 2, 2]) * dim_unit + y = np.arange(array.shape[1]) * 3 * dim_unit - data_array = xr.DataArray(data=array, coords={"x": x, "y": y}, dims=("x", "y")) + u = np.linspace(0, 1, 5) * coord_unit + + data_array = xr.DataArray( + data=array, coords={"x": x, "y": y, "u": ("x", u)}, dims=("x", "y") + ) units = extract_units(data_array) expected = attach_units(func(strip_units(data_array)).mean(), units) actual = func(data_array).mean() assert_units_equal(expected, actual) - xr.testing.assert_allclose(expected, actual) + assert_allclose(expected, actual) def test_resample(self, dtype): array = np.linspace(0, 5, 10).astype(dtype) * unit_registry.m @@ -3665,33 +3783,47 @@ def test_resample(self, dtype): actual = func(data_array).mean() assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) + @pytest.mark.parametrize( + "variant", + ( + "data", + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", + ), + ) @pytest.mark.parametrize( "func", ( - method("assign_coords", z=(["x"], np.arange(5) * unit_registry.s)), + method("assign_coords", z=("x", np.arange(5) * unit_registry.s)), method("first"), method("last"), - pytest.param( - method("quantile", q=[0.25, 0.5, 0.75], dim="x"), - marks=pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", - reason="quantile / nanquantile not implemented yet", - ), - ), + method("quantile", q=[0.25, 0.5, 0.75], dim="x"), ), ids=repr, ) - def test_grouped_operations(self, func, dtype): - array = ( - np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * unit_registry.m - ) + def test_grouped_operations(self, func, variant, dtype): + unit = unit_registry.m - x = np.arange(array.shape[0]) * unit_registry.m - y = np.arange(array.shape[1]) * 3 * unit_registry.s + variants = { + "data": (unit, 1, 1), + "dims": (1, unit, 1), + "coords": (1, 1, unit), + } + data_unit, dim_unit, coord_unit = variants.get(variant) + array = np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * data_unit + + x = np.arange(array.shape[0]) * dim_unit + y = np.arange(array.shape[1]) * 3 * dim_unit + + u = np.linspace(0, 1, array.shape[0]) * coord_unit - data_array = xr.DataArray(data=array, coords={"x": x, "y": y}, dims=("x", "y")) + data_array = xr.DataArray( + data=array, coords={"x": x, "y": y, "u": ("x", u)}, dims=("x", "y") + ) units = {**extract_units(data_array), **{"z": unit_registry.s, "q": None}} stripped_kwargs = { @@ -3708,10 +3840,9 @@ def test_grouped_operations(self, func, dtype): actual = func(data_array.groupby("y")) assert_units_equal(expected, actual) - xr.testing.assert_identical(expected, actual) + assert_identical(expected, actual) -@pytest.mark.filterwarnings("error::pint.UnitStrippedWarning") class TestDataset: @pytest.mark.parametrize( "unit,error", @@ -3796,8 +3927,7 @@ def test_init(self, shared, unit, error, dtype): ( "data", pytest.param( - "dims", - marks=pytest.mark.xfail(reason="units in indexes are not supported"), + "dims", marks=pytest.mark.xfail(reason="indexes don't support units"), ), "coords", ), @@ -4313,7 +4443,7 @@ def test_combine_first(self, variant, unit, error, dtype): ( "data", pytest.param( - "dims", marks=pytest.mark.xfail(reason="units in indexes not supported") + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") ), "coords", ), @@ -4473,6 +4603,19 @@ def test_broadcast_equals(self, unit, dtype): assert expected == actual + def test_pad(self, dtype): + a = np.linspace(0, 5, 10).astype(dtype) * unit_registry.Pa + b = np.linspace(-5, 0, 10).astype(dtype) * unit_registry.degK + + ds = xr.Dataset({"a": ("x", a), "b": ("x", b)}) + units = extract_units(ds) + + expected = attach_units(strip_units(ds).pad(x=(2, 3)), units) + actual = ds.pad(x=(2, 3)) + + assert_units_equal(expected, actual) + assert_equal(expected, actual) + @pytest.mark.parametrize( "func", (method("unstack"), method("reset_index", "v"), method("reorder_levels")), @@ -5011,13 +5154,7 @@ def test_interp_reindex_like_indexing(self, func, unit, error, dtype): method("diff", dim="x"), method("differentiate", coord="x"), method("integrate", coord="x"), - pytest.param( - method("quantile", q=[0.25, 0.75]), - marks=pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", - reason="nanquantile not implemented yet", - ), - ), + method("quantile", q=[0.25, 0.75]), method("reduce", func=np.sum, dim="x"), method("map", np.fabs), ), @@ -5067,13 +5204,7 @@ def test_computation(self, func, variant, dtype): "func", ( method("groupby", "x"), - pytest.param( - method("groupby_bins", "x", bins=2), - marks=pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", - reason="needs assert_allclose but that does not work with pint", - ), - ), + method("groupby_bins", "x", bins=2), method("coarsen", x=2), pytest.param( method("rolling", x=3), marks=pytest.mark.xfail(reason="strips units") @@ -5084,6 +5215,7 @@ def test_computation(self, func, variant, dtype): reason="numbagg functions are not supported by pint" ), ), + method("weighted", xr.DataArray(data=np.linspace(0, 1, 5), dims="y")), ), ids=repr, ) @@ -5122,11 +5254,7 @@ def test_computation_objects(self, func, variant, dtype): actual = func(ds).mean(*args) assert_units_equal(expected, actual) - # TODO: remove once pint 0.12 has been released - if LooseVersion(pint.__version__) <= "0.12": - assert_equal(expected, actual) - else: - assert_allclose(expected, actual) + assert_allclose(expected, actual) @pytest.mark.parametrize( "variant", @@ -5177,13 +5305,7 @@ def test_resample(self, variant, dtype): method("assign_coords", v=("x", np.arange(5) * unit_registry.s)), method("first"), method("last"), - pytest.param( - method("quantile", q=[0.25, 0.5, 0.75], dim="x"), - marks=pytest.mark.xfail( - LooseVersion(pint.__version__) <= "0.12", - reason="nanquantile not implemented", - ), - ), + method("quantile", q=[0.25, 0.5, 0.75], dim="x"), ), ids=repr, )