More pint compatibility: silence UnitStrippedWarnings (#4163)

* globally promote UnitStrippedWarning to errors

* separately test apply_ufunc with units in dims, coords and data

* split the DataArray align test into data, dims and coords tests

* use dtypes instead of python types and use a dtype specific fill value

* rewrite the dataset align tests

* compare with dtypes.NA instead of using np.isnan

* mention the issue in the xfail reason

* make sure the combine_* variants are properly separated from each other

* improve the test case names

* note that broadcast uses align

* properly separate the test cases for concat

* always use the same reason when xfailing units in indexes tests

* also check that the replication functions work with dims and units

* apply full_like to the data instead of the variable

* check full_like with units in dims, data and coords separately

* clearly separate the test variants of the merge tests

* don't use indexes for the dataset where tests

* replace numpy.testing.assert_allclose with assert numpy.allclose

* remove a conditional xfail that depends on a very old pint version

* use assert_identical from the local namespace

* properly separate between the broadcast_like test variants

* don't accept "data" as an alias of the DataArray's data

* properly separate between the variants of the content manipulation tests

* use assert np.allclose(...) instead of np.testing.assert_allclose(...)

* don't test units in indexes in the isel tests

* don't use units in indexes for the head / tail / thin tests

* properly separate the variants of more tests

* rewrite the squeeze tests

* use assert_allclose from the module's namespace

* rewrite the copy tests

* xfail the equal comparison for a pint version lower than 0.14

* try to implement a duckarray friendly assert_array_equal

* add tests for not raising an assertion error

* skip only the dask test if it isn't installed

* also check using pint if available

* add a duckarray version of np.testing.assert_allclose

* add both to __all__

* make both available in xarray.tests

* don't inherit from VariableSubtests since that was not written to test duck arrays.

* test the constant pad mode along with all other modes

* remove most pint version checks, now that pint 0.13 has been released

* use conda to install pint

* xfail the DataArray comparison test until pint's dev version fixed it

* add tests for the pad method of DataArray and Dataset

* add tests for weighted

* update whats-new.rst

* replace assert np.allclose(...) with assert_duckarray_allclose(...)

* fix the dask fallback

* xfail the pint tests for now since there's a bug in pint

* use utils.is_array_like and utils.is_scalar

Author: keewis

ci/requirements/py36-min-nep18.yml

@@ -11,12 +11,11 @@ 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
   - pytest-env
   - scipy=1.2
   - setuptools=41.2
   - sparse=0.8
-  - pip:
-      - pint==0.13

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

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

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

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

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

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 <https://github.com/kmuehlbauer>`_ and `Pascal Bourgault <https://github.com/aulemahal>`_.
-- 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 <https://github.com/keewis>`_.
 - Support overriding existing variables in ``to_zarr()`` with ``mode='a'`` even
   without ``append_dim``, as long as dimension sizes do not change.

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)
 

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],

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.

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:

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)