array")
def test_short_data_repr_html_non_str_keys(dataset):
@@ -108,8 +108,8 @@ def test_summarize_attrs_with_unsafe_attr_name_and_value():
def test_repr_of_dataarray(dataarray):
formatted = fh.array_repr(dataarray)
assert "dim_0" in formatted
- # has an expandable data section
- assert formatted.count("class='xr-array-in' type='checkbox' >") == 1
+ # has an expanded data section
+ assert formatted.count("class='xr-array-in' type='checkbox' checked>") == 1
# coords and attrs don't have an items so they'll be be disabled and collapsed
assert (
formatted.count("class='xr-section-summary-in' type='checkbox' disabled >") == 2
@@ -137,3 +137,22 @@ def test_repr_of_dataset(dataset):
)
assert "<U4" in formatted or ">U4" in formatted
assert "<IA>" in formatted
+
+
+def test_repr_text_fallback(dataset):
+ formatted = fh.dataset_repr(dataset)
+
+ # Just test that the "pre" block used for fallback to plain text is present.
+ assert "" in formatted
+
+
+def test_variable_repr_html():
+ v = xr.Variable(["time", "x"], [[1, 2, 3], [4, 5, 6]], {"foo": "bar"})
+ assert hasattr(v, "_repr_html_")
+ with xr.set_options(display_style="html"):
+ html = v._repr_html_().strip()
+ # We don't do a complete string identity since
+ # html output is probably subject to change, is long and... reasons.
+ # Just test that something reasonable was produced.
+ assert html.startswith("")
+ assert "xarray.Variable" in html
diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py
index 866d5fb0899..aa54c8f36f1 100644
--- a/xarray/tests/test_groupby.py
+++ b/xarray/tests/test_groupby.py
@@ -538,4 +538,16 @@ def test_groupby_bins_timeseries():
assert_identical(actual, expected)
+def test_groupby_none_group_name():
+ # GH158
+ # xarray should not fail if a DataArray's name attribute is None
+
+ data = np.arange(10) + 10
+ da = xr.DataArray(data) # da.name = None
+ key = xr.DataArray(np.floor_divide(data, 2))
+
+ mean = da.groupby(key).mean()
+ assert "group" in mean.dims
+
+
# TODO: move other groupby tests from test_dataset and test_dataarray over here
diff --git a/xarray/tests/test_interp.py b/xarray/tests/test_interp.py
index 0502348160e..7a0dda216e2 100644
--- a/xarray/tests/test_interp.py
+++ b/xarray/tests/test_interp.py
@@ -699,3 +699,21 @@ def test_3641():
times = xr.cftime_range("0001", periods=3, freq="500Y")
da = xr.DataArray(range(3), dims=["time"], coords=[times])
da.interp(time=["0002-05-01"])
+
+
+@requires_scipy
+@pytest.mark.parametrize("method", ["nearest", "linear"])
+def test_decompose(method):
+ da = xr.DataArray(
+ np.arange(6).reshape(3, 2),
+ dims=["x", "y"],
+ coords={"x": [0, 1, 2], "y": [-0.1, -0.3]},
+ )
+ x_new = xr.DataArray([0.5, 1.5, 2.5], dims=["x1"])
+ y_new = xr.DataArray([-0.15, -0.25], dims=["y1"])
+ x_broadcast, y_broadcast = xr.broadcast(x_new, y_new)
+ assert x_broadcast.ndim == 2
+
+ actual = da.interp(x=x_new, y=y_new, method=method).drop(("x", "y"))
+ expected = da.interp(x=x_broadcast, y=y_broadcast, method=method).drop(("x", "y"))
+ assert_allclose(actual, expected)
diff --git a/xarray/tests/test_missing.py b/xarray/tests/test_missing.py
index 731cd165244..bc186c8bd15 100644
--- a/xarray/tests/test_missing.py
+++ b/xarray/tests/test_missing.py
@@ -534,6 +534,18 @@ def test_get_clean_interp_index_potential_overflow():
get_clean_interp_index(da, "time")
+@pytest.mark.parametrize("index", ([0, 2, 1], [0, 1, 1]))
+def test_get_clean_interp_index_strict(index):
+ da = xr.DataArray([0, 1, 2], dims=("x",), coords={"x": index})
+
+ with pytest.raises(ValueError):
+ get_clean_interp_index(da, "x")
+
+ clean = get_clean_interp_index(da, "x", strict=False)
+ np.testing.assert_array_equal(index, clean)
+ assert clean.dtype == np.float64
+
+
@pytest.fixture
def da_time():
return xr.DataArray(
diff --git a/xarray/tests/test_plot.py b/xarray/tests/test_plot.py
index bf1f9ed60bb..938f403e01b 100644
--- a/xarray/tests/test_plot.py
+++ b/xarray/tests/test_plot.py
@@ -111,6 +111,12 @@ class TestPlot(PlotTestCase):
def setup_array(self):
self.darray = DataArray(easy_array((2, 3, 4)))
+ def test_accessor(self):
+ from ..plot.plot import _PlotMethods
+
+ assert DataArray.plot is _PlotMethods
+ assert isinstance(self.darray.plot, _PlotMethods)
+
def test_label_from_attrs(self):
da = self.darray.copy()
assert "" == label_from_attrs(da)
@@ -136,14 +142,14 @@ def test_label_from_attrs(self):
def test1d(self):
self.darray[:, 0, 0].plot()
- with raises_regex(ValueError, "None"):
+ with raises_regex(ValueError, "x must be one of None, 'dim_0'"):
self.darray[:, 0, 0].plot(x="dim_1")
with raises_regex(TypeError, "complex128"):
(self.darray[:, 0, 0] + 1j).plot()
def test_1d_bool(self):
- xr.ones_like(self.darray[:, 0, 0], dtype=np.bool).plot()
+ xr.ones_like(self.darray[:, 0, 0], dtype=bool).plot()
def test_1d_x_y_kw(self):
z = np.arange(10)
@@ -155,14 +161,31 @@ def test_1d_x_y_kw(self):
for aa, (x, y) in enumerate(xy):
da.plot(x=x, y=y, ax=ax.flat[aa])
- with raises_regex(ValueError, "cannot"):
+ with raises_regex(ValueError, "Cannot specify both"):
da.plot(x="z", y="z")
- with raises_regex(ValueError, "None"):
- da.plot(x="f", y="z")
+ error_msg = "must be one of None, 'z'"
+ with raises_regex(ValueError, f"x {error_msg}"):
+ da.plot(x="f")
+
+ with raises_regex(ValueError, f"y {error_msg}"):
+ da.plot(y="f")
+
+ def test_multiindex_level_as_coord(self):
+ da = xr.DataArray(
+ np.arange(5),
+ dims="x",
+ coords=dict(a=("x", np.arange(5)), b=("x", np.arange(5, 10))),
+ )
+ da = da.set_index(x=["a", "b"])
+
+ for x in ["a", "b"]:
+ h = da.plot(x=x)[0]
+ assert_array_equal(h.get_xdata(), da[x].values)
- with raises_regex(ValueError, "None"):
- da.plot(x="z", y="f")
+ for y in ["a", "b"]:
+ h = da.plot(y=y)[0]
+ assert_array_equal(h.get_ydata(), da[y].values)
# Test for bug in GH issue #2725
def test_infer_line_data(self):
@@ -211,7 +234,7 @@ def test_2d_line(self):
self.darray[:, :, 0].plot.line(x="dim_0", hue="dim_1")
self.darray[:, :, 0].plot.line(y="dim_0", hue="dim_1")
- with raises_regex(ValueError, "cannot"):
+ with raises_regex(ValueError, "Cannot"):
self.darray[:, :, 0].plot.line(x="dim_1", y="dim_0", hue="dim_1")
def test_2d_line_accepts_legend_kw(self):
@@ -854,21 +877,22 @@ def test_norm_sets_vmin_vmax(self):
vmin = self.data.min()
vmax = self.data.max()
- for norm, extend in zip(
+ for norm, extend, levels in zip(
[
+ mpl.colors.Normalize(),
mpl.colors.Normalize(),
mpl.colors.Normalize(vmin + 0.1, vmax - 0.1),
mpl.colors.Normalize(None, vmax - 0.1),
mpl.colors.Normalize(vmin + 0.1, None),
],
- ["neither", "both", "max", "min"],
+ ["neither", "neither", "both", "max", "min"],
+ [7, None, None, None, None],
):
test_min = vmin if norm.vmin is None else norm.vmin
test_max = vmax if norm.vmax is None else norm.vmax
- cmap_params = _determine_cmap_params(self.data, norm=norm)
-
+ cmap_params = _determine_cmap_params(self.data, norm=norm, levels=levels)
assert cmap_params["vmin"] == test_min
assert cmap_params["vmax"] == test_max
assert cmap_params["extend"] == extend
@@ -1013,7 +1037,7 @@ def test_1d_raises_valueerror(self):
self.plotfunc(self.darray[0, :])
def test_bool(self):
- xr.ones_like(self.darray, dtype=np.bool).plot()
+ xr.ones_like(self.darray, dtype=bool).plot()
def test_complex_raises_typeerror(self):
with raises_regex(TypeError, "complex128"):
@@ -1031,6 +1055,16 @@ def test_nonnumeric_index_raises_typeerror(self):
with raises_regex(TypeError, r"[Pp]lot"):
self.plotfunc(a)
+ def test_multiindex_raises_typeerror(self):
+ a = DataArray(
+ easy_array((3, 2)),
+ dims=("x", "y"),
+ coords=dict(x=("x", [0, 1, 2]), a=("y", [0, 1]), b=("y", [2, 3])),
+ )
+ a = a.set_index(y=("a", "b"))
+ with raises_regex(TypeError, r"[Pp]lot"):
+ self.plotfunc(a)
+
def test_can_pass_in_axis(self):
self.pass_in_axis(self.plotmethod)
@@ -1139,15 +1173,16 @@ def test_positional_coord_string(self):
assert "y_long_name [y_units]" == ax.get_ylabel()
def test_bad_x_string_exception(self):
- with raises_regex(ValueError, "x and y must be coordinate variables"):
+
+ with raises_regex(ValueError, "x and y cannot be equal."):
+ self.plotmethod(x="y", y="y")
+
+ error_msg = "must be one of None, 'x', 'x2d', 'y', 'y2d'"
+ with raises_regex(ValueError, f"x {error_msg}"):
self.plotmethod("not_a_real_dim", "y")
- with raises_regex(
- ValueError, "x must be a dimension name if y is not supplied"
- ):
+ with raises_regex(ValueError, f"x {error_msg}"):
self.plotmethod(x="not_a_real_dim")
- with raises_regex(
- ValueError, "y must be a dimension name if x is not supplied"
- ):
+ with raises_regex(ValueError, f"y {error_msg}"):
self.plotmethod(y="not_a_real_dim")
self.darray.coords["z"] = 100
@@ -1182,6 +1217,27 @@ def test_non_linked_coords_transpose(self):
# simply ensure that these high coords were passed over
assert np.min(ax.get_xlim()) > 100.0
+ def test_multiindex_level_as_coord(self):
+ da = DataArray(
+ easy_array((3, 2)),
+ dims=("x", "y"),
+ coords=dict(x=("x", [0, 1, 2]), a=("y", [0, 1]), b=("y", [2, 3])),
+ )
+ da = da.set_index(y=["a", "b"])
+
+ for x, y in (("a", "x"), ("b", "x"), ("x", "a"), ("x", "b")):
+ self.plotfunc(da, x=x, y=y)
+
+ ax = plt.gca()
+ assert x == ax.get_xlabel()
+ assert y == ax.get_ylabel()
+
+ with raises_regex(ValueError, "levels of the same MultiIndex"):
+ self.plotfunc(da, x="a", y="b")
+
+ with raises_regex(ValueError, "y must be one of None, 'a', 'b', 'x'"):
+ self.plotfunc(da, x="a", y="y")
+
def test_default_title(self):
a = DataArray(easy_array((4, 3, 2)), dims=["a", "b", "c"])
a.coords["c"] = [0, 1]
@@ -2048,6 +2104,12 @@ def setUp(self):
ds.B.attrs["units"] = "Bunits"
self.ds = ds
+ def test_accessor(self):
+ from ..plot.dataset_plot import _Dataset_PlotMethods
+
+ assert Dataset.plot is _Dataset_PlotMethods
+ assert isinstance(self.ds.plot, _Dataset_PlotMethods)
+
@pytest.mark.parametrize(
"add_guide, hue_style, legend, colorbar",
[
diff --git a/xarray/tests/test_testing.py b/xarray/tests/test_testing.py
index 041b7341ade..f4961af58e9 100644
--- a/xarray/tests/test_testing.py
+++ b/xarray/tests/test_testing.py
@@ -1,3 +1,5 @@
+import pytest
+
import xarray as xr
@@ -5,3 +7,26 @@ def test_allclose_regression():
x = xr.DataArray(1.01)
y = xr.DataArray(1.02)
xr.testing.assert_allclose(x, y, atol=0.01)
+
+
+@pytest.mark.parametrize(
+ "obj1,obj2",
+ (
+ pytest.param(
+ xr.Variable("x", [1e-17, 2]), xr.Variable("x", [0, 3]), id="Variable",
+ ),
+ pytest.param(
+ xr.DataArray([1e-17, 2], dims="x"),
+ xr.DataArray([0, 3], dims="x"),
+ id="DataArray",
+ ),
+ pytest.param(
+ xr.Dataset({"a": ("x", [1e-17, 2]), "b": ("y", [-2e-18, 2])}),
+ xr.Dataset({"a": ("x", [0, 2]), "b": ("y", [0, 1])}),
+ id="Dataset",
+ ),
+ ),
+)
+def test_assert_allclose(obj1, obj2):
+ with pytest.raises(AssertionError):
+ xr.testing.assert_allclose(obj1, obj2)
diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py
index 2826dc2479c..20a5f0e8613 100644
--- a/xarray/tests/test_units.py
+++ b/xarray/tests/test_units.py
@@ -7,9 +7,8 @@
import pytest
import xarray as xr
-from xarray.core import formatting
from xarray.core.npcompat import IS_NEP18_ACTIVE
-from xarray.testing import assert_allclose, assert_identical
+from xarray.testing import assert_allclose, assert_equal, assert_identical
from .test_variable import _PAD_XR_NP_ARGS, VariableSubclassobjects
@@ -27,11 +26,6 @@
pytest.mark.skipif(
not IS_NEP18_ACTIVE, reason="NUMPY_EXPERIMENTAL_ARRAY_FUNCTION is not enabled"
),
- # TODO: remove this once pint has a released version with __array_function__
- pytest.mark.skipif(
- not hasattr(unit_registry.Quantity, "__array_function__"),
- reason="pint does not implement __array_function__ yet",
- ),
# pytest.mark.filterwarnings("ignore:::pint[.*]"),
]
@@ -51,10 +45,23 @@ def dimensionality(obj):
def compatible_mappings(first, second):
return {
key: is_compatible(unit1, unit2)
- for key, (unit1, unit2) in merge_mappings(first, second)
+ for key, (unit1, unit2) in zip_mappings(first, second)
}
+def merge_mappings(base, *mappings):
+ result = base.copy()
+ for m in mappings:
+ result.update(m)
+
+ return result
+
+
+def zip_mappings(*mappings):
+ for key in set(mappings[0]).intersection(*mappings[1:]):
+ yield key, tuple(m[key] for m in mappings)
+
+
def array_extract_units(obj):
if isinstance(obj, (xr.Variable, xr.DataArray, xr.Dataset)):
obj = obj.data
@@ -257,50 +264,11 @@ def assert_units_equal(a, b):
assert extract_units(a) == extract_units(b)
-def assert_equal_with_units(a, b):
- # works like xr.testing.assert_equal, but also explicitly checks units
- # so, it is more like assert_identical
- __tracebackhide__ = True
-
- if isinstance(a, xr.Dataset) or isinstance(b, xr.Dataset):
- a_units = extract_units(a)
- b_units = extract_units(b)
-
- a_without_units = strip_units(a)
- b_without_units = strip_units(b)
-
- assert a_without_units.equals(b_without_units), formatting.diff_dataset_repr(
- a, b, "equals"
- )
- assert a_units == b_units
- else:
- a = a if not isinstance(a, (xr.DataArray, xr.Variable)) else a.data
- b = b if not isinstance(b, (xr.DataArray, xr.Variable)) else b.data
-
- assert type(a) == type(b) or (
- isinstance(a, Quantity) and isinstance(b, Quantity)
- )
-
- # workaround until pint implements allclose in __array_function__
- if isinstance(a, Quantity) or isinstance(b, Quantity):
- assert (
- hasattr(a, "magnitude") and hasattr(b, "magnitude")
- ) and np.allclose(a.magnitude, b.magnitude, equal_nan=True)
- assert (hasattr(a, "units") and hasattr(b, "units")) and a.units == b.units
- else:
- assert np.allclose(a, b, equal_nan=True)
-
-
@pytest.fixture(params=[float, int])
def dtype(request):
return request.param
-def merge_mappings(*mappings):
- for key in set(mappings[0]).intersection(*mappings[1:]):
- yield key, tuple(m[key] for m in mappings)
-
-
def merge_args(default_args, new_args):
from itertools import zip_longest
@@ -329,19 +297,29 @@ def __call__(self, obj, *args, **kwargs):
all_args = merge_args(self.args, args)
all_kwargs = {**self.kwargs, **kwargs}
+ xarray_classes = (
+ xr.Variable,
+ xr.DataArray,
+ xr.Dataset,
+ xr.core.groupby.GroupBy,
+ )
+
+ if not isinstance(obj, xarray_classes):
+ # remove typical xarray args like "dim"
+ exclude_kwargs = ("dim", "dims")
+ all_kwargs = {
+ key: value
+ for key, value in all_kwargs.items()
+ if key not in exclude_kwargs
+ }
+
func = getattr(obj, self.name, None)
+
if func is None or not isinstance(func, Callable):
# fall back to module level numpy functions if not a xarray object
if not isinstance(obj, (xr.Variable, xr.DataArray, xr.Dataset)):
numpy_func = getattr(np, self.name)
func = partial(numpy_func, obj)
- # remove typical xarray args like "dim"
- exclude_kwargs = ("dim", "dims")
- all_kwargs = {
- key: value
- for key, value in all_kwargs.items()
- if key not in exclude_kwargs
- }
else:
raise AttributeError(f"{obj} has no method named '{self.name}'")
@@ -425,6 +403,10 @@ 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",
(
@@ -512,6 +494,10 @@ 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",
(
@@ -929,6 +915,10 @@ 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",
(
@@ -1036,6 +1026,10 @@ 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",
(
@@ -1385,7 +1379,6 @@ def wrapper(cls):
"test_datetime64_conversion",
"test_timedelta64_conversion",
"test_pandas_period_index",
- "test_1d_math",
"test_1d_reduce",
"test_array_interface",
"test___array__",
@@ -1413,13 +1406,20 @@ def example_1d_objects(self):
]:
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()
+
@pytest.mark.parametrize(
"func",
(
method("all"),
method("any"),
- method("argmax"),
- method("argmin"),
+ method("argmax", dim="x"),
+ method("argmin", dim="x"),
method("argsort"),
method("cumprod"),
method("cumsum"),
@@ -1443,13 +1443,33 @@ def test_aggregation(self, func, dtype):
)
variable = xr.Variable("x", array)
- units = extract_units(func(array))
+ numpy_kwargs = func.kwargs.copy()
+ if "dim" in func.kwargs:
+ numpy_kwargs["axis"] = variable.get_axis_num(numpy_kwargs.pop("dim"))
+
+ units = extract_units(func(array, **numpy_kwargs))
expected = attach_units(func(strip_units(variable)), units)
actual = func(variable)
assert_units_equal(expected, actual)
- xr.testing.assert_identical(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)
+ # 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",
(
@@ -1660,7 +1680,7 @@ def test_missing_value_fillna(self, unit, error):
method("equals"),
pytest.param(
method("identical"),
- marks=pytest.mark.skip(reason="behaviour of identical is unclear"),
+ marks=pytest.mark.skip(reason="behavior of identical is undecided"),
),
),
ids=repr,
@@ -1748,6 +1768,10 @@ def test_isel(self, indices, dtype):
assert_units_equal(expected, actual)
xr.testing.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",
(
@@ -1885,7 +1909,10 @@ def test_squeeze(self, dtype):
method("coarsen", windows={"y": 2}, func=np.mean),
pytest.param(
method("quantile", q=[0.25, 0.75]),
- marks=pytest.mark.xfail(reason="nanquantile not implemented"),
+ marks=pytest.mark.xfail(
+ LooseVersion(pint.__version__) <= "0.12",
+ reason="quantile / nanquantile not implemented yet",
+ ),
),
pytest.param(
method("rank", dim="x"),
@@ -2161,8 +2188,8 @@ class TestDataArray:
"with_dims",
marks=pytest.mark.xfail(reason="units in indexes are not supported"),
),
- pytest.param("with_coords"),
- pytest.param("without_coords"),
+ "with_coords",
+ "without_coords",
),
)
def test_init(self, variant, dtype):
@@ -2221,24 +2248,36 @@ 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",
(
+ function("all"),
+ function("any"),
pytest.param(
- function("all"),
- marks=pytest.mark.xfail(reason="not implemented by pint yet"),
+ function("argmax"),
+ marks=pytest.mark.skip(
+ reason="calling np.argmax as a function on xarray objects is not "
+ "supported"
+ ),
),
pytest.param(
- function("any"),
- marks=pytest.mark.xfail(reason="not implemented by pint yet"),
+ function("argmin"),
+ marks=pytest.mark.skip(
+ reason="calling np.argmin as a function on xarray objects is not "
+ "supported"
+ ),
),
- function("argmax"),
- function("argmin"),
function("max"),
function("mean"),
pytest.param(
function("median"),
- marks=pytest.mark.xfail(reason="not implemented by xarray"),
+ marks=pytest.mark.skip(
+ reason="median does not work with dataarrays yet"
+ ),
),
function("min"),
pytest.param(
@@ -2249,38 +2288,24 @@ def test_repr(self, func, variant, dtype):
function("std"),
function("var"),
function("cumsum"),
- pytest.param(
- function("cumprod"),
- marks=pytest.mark.xfail(reason="not implemented by pint yet"),
- ),
- pytest.param(
- method("all"),
- marks=pytest.mark.xfail(reason="not implemented by pint yet"),
- ),
- pytest.param(
- method("any"),
- marks=pytest.mark.xfail(reason="not implemented by pint yet"),
- ),
- method("argmax"),
- method("argmin"),
+ function("cumprod"),
+ method("all"),
+ method("any"),
+ method("argmax", dim="x"),
+ method("argmin", dim="x"),
method("max"),
method("mean"),
method("median"),
method("min"),
pytest.param(
method("prod"),
- marks=pytest.mark.xfail(
- reason="comparison of quantity with ndarrays in nanops not implemented"
- ),
+ marks=pytest.mark.xfail(reason="not implemented by pint yet"),
),
method("sum"),
method("std"),
method("var"),
method("cumsum"),
- pytest.param(
- method("cumprod"),
- marks=pytest.mark.xfail(reason="pint does not implement cumprod yet"),
- ),
+ method("cumprod"),
),
ids=repr,
)
@@ -2290,13 +2315,18 @@ def test_aggregation(self, func, dtype):
)
data_array = xr.DataArray(data=array, dims="x")
+ numpy_kwargs = func.kwargs.copy()
+ if "dim" in numpy_kwargs:
+ numpy_kwargs["axis"] = data_array.get_axis_num(numpy_kwargs.pop("dim"))
+
# units differ based on the applied function, so we need to
# first compute the units
units = extract_units(func(array))
expected = attach_units(func(strip_units(data_array)), units)
actual = func(data_array)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_allclose(expected, actual)
@pytest.mark.parametrize(
"func",
@@ -2314,7 +2344,8 @@ def test_unary_operations(self, func, dtype):
expected = attach_units(func(strip_units(data_array)), units)
actual = func(data_array)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"func",
@@ -2333,7 +2364,8 @@ def test_binary_operations(self, func, dtype):
expected = attach_units(func(strip_units(data_array)), units)
actual = func(data_array)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"comparison",
@@ -2383,7 +2415,8 @@ def test_comparison_operations(self, comparison, unit, error, dtype):
strip_units(convert_units(to_compare_with, expected_units)),
)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"units,error",
@@ -2411,9 +2444,10 @@ def test_univariate_ufunc(self, units, error, dtype):
)
actual = func(data_array)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
- @pytest.mark.xfail(reason="xarray's `np.maximum` strips units")
+ @pytest.mark.xfail(reason="needs the type register system for __array_ufunc__")
@pytest.mark.parametrize(
"unit,error",
(
@@ -2422,7 +2456,12 @@ def test_univariate_ufunc(self, units, error, dtype):
unit_registry.dimensionless, DimensionalityError, id="dimensionless"
),
pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
- pytest.param(unit_registry.mm, None, id="compatible_unit"),
+ pytest.param(
+ unit_registry.mm,
+ None,
+ id="compatible_unit",
+ marks=pytest.mark.xfail(reason="pint converts to the wrong units"),
+ ),
pytest.param(unit_registry.m, None, id="identical_unit"),
),
)
@@ -2433,7 +2472,7 @@ def test_bivariate_ufunc(self, unit, error, dtype):
if error is not None:
with pytest.raises(error):
- np.maximum(data_array, 0 * unit)
+ np.maximum(data_array, 1 * unit)
return
@@ -2441,16 +2480,18 @@ def test_bivariate_ufunc(self, unit, error, dtype):
expected = attach_units(
np.maximum(
strip_units(data_array),
- strip_units(convert_units(0 * unit, expected_units)),
+ strip_units(convert_units(1 * unit, expected_units)),
),
expected_units,
)
- actual = np.maximum(data_array, 0 * unit)
- assert_equal_with_units(expected, actual)
+ actual = np.maximum(data_array, 1 * unit)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
- actual = np.maximum(0 * unit, data_array)
- assert_equal_with_units(expected, actual)
+ actual = np.maximum(1 * unit, data_array)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize("property", ("T", "imag", "real"))
def test_numpy_properties(self, property, dtype):
@@ -2466,7 +2507,8 @@ def test_numpy_properties(self, property, dtype):
)
actual = getattr(data_array, property)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"func",
@@ -2481,16 +2523,86 @@ def test_numpy_methods(self, func, dtype):
expected = attach_units(strip_units(data_array), units)
actual = func(data_array)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
+
+ def test_item(self, dtype):
+ array = np.arange(10).astype(dtype) * unit_registry.m
+ data_array = xr.DataArray(data=array)
+
+ func = method("item", 2)
+
+ expected = func(strip_units(data_array)) * unit_registry.m
+ actual = func(data_array)
+
+ np.testing.assert_allclose(expected, actual)
+
+ @pytest.mark.parametrize(
+ "unit,error",
+ (
+ pytest.param(1, DimensionalityError, id="no_unit"),
+ pytest.param(
+ unit_registry.dimensionless, DimensionalityError, id="dimensionless"
+ ),
+ pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
+ pytest.param(unit_registry.cm, None, id="compatible_unit"),
+ pytest.param(unit_registry.m, None, id="identical_unit"),
+ ),
+ )
+ @pytest.mark.parametrize(
+ "func",
+ (
+ method("searchsorted", 5),
+ pytest.param(
+ function("searchsorted", 5),
+ marks=pytest.mark.xfail(
+ reason="xarray does not implement __array_function__"
+ ),
+ ),
+ ),
+ ids=repr,
+ )
+ def test_searchsorted(self, func, unit, error, dtype):
+ array = np.arange(10).astype(dtype) * unit_registry.m
+ data_array = xr.DataArray(data=array)
+
+ scalar_types = (int, float)
+ args = list(value * unit for value in func.args)
+ kwargs = {
+ key: (value * unit if isinstance(value, scalar_types) else value)
+ for key, value in func.kwargs.items()
+ }
+
+ if error is not None:
+ with pytest.raises(error):
+ func(data_array, *args, **kwargs)
+
+ return
+
+ units = extract_units(data_array)
+ expected_units = extract_units(func(array, *args, **kwargs))
+ stripped_args = [strip_units(convert_units(value, units)) for value in args]
+ stripped_kwargs = {
+ key: strip_units(convert_units(value, units))
+ for key, value in kwargs.items()
+ }
+ expected = attach_units(
+ func(strip_units(data_array), *stripped_args, **stripped_kwargs),
+ expected_units,
+ )
+ actual = func(data_array, *args, **kwargs)
+
+ assert_units_equal(expected, actual)
+ np.testing.assert_allclose(expected, actual)
@pytest.mark.parametrize(
"func",
(
method("clip", min=3, max=8),
pytest.param(
- method("searchsorted", v=5),
+ function("clip", a_min=3, a_max=8),
marks=pytest.mark.xfail(
- reason="searchsorted somehow requires a undocumented `keys` argument"
+ reason="xarray does not implement __array_function__"
),
),
),
@@ -2513,28 +2625,32 @@ def test_numpy_methods_with_args(self, func, unit, error, dtype):
data_array = xr.DataArray(data=array)
scalar_types = (int, float)
+ args = list(value * unit for value in func.args)
kwargs = {
key: (value * unit if isinstance(value, scalar_types) else value)
for key, value in func.kwargs.items()
}
if error is not None:
with pytest.raises(error):
- func(data_array, **kwargs)
+ func(data_array, *args, **kwargs)
return
units = extract_units(data_array)
- expected_units = extract_units(func(array, **kwargs))
+ expected_units = extract_units(func(array, *args, **kwargs))
+ stripped_args = [strip_units(convert_units(value, units)) for value in args]
stripped_kwargs = {
key: strip_units(convert_units(value, units))
for key, value in kwargs.items()
}
expected = attach_units(
- func(strip_units(data_array), **stripped_kwargs), expected_units
+ func(strip_units(data_array), *stripped_args, **stripped_kwargs),
+ expected_units,
)
- actual = func(data_array, **kwargs)
+ actual = func(data_array, *args, **kwargs)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"func", (method("isnull"), method("notnull"), method("count")), ids=repr
@@ -2551,15 +2667,13 @@ def test_missing_value_detection(self, func, dtype):
)
* unit_registry.degK
)
- x = np.arange(array.shape[0]) * unit_registry.m
- y = np.arange(array.shape[1]) * unit_registry.m
-
- data_array = xr.DataArray(data=array, coords={"x": x, "y": y}, dims=("x", "y"))
+ data_array = xr.DataArray(data=array)
expected = func(strip_units(data_array))
actual = func(data_array)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.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)
@@ -2576,7 +2690,8 @@ def test_missing_value_filling(self, func, dtype):
)
actual = func(data_array, dim="x")
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"unit,error",
@@ -2586,12 +2701,7 @@ def test_missing_value_filling(self, func, dtype):
unit_registry.dimensionless, DimensionalityError, id="dimensionless"
),
pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
- pytest.param(
- unit_registry.cm,
- None,
- id="compatible_unit",
- marks=pytest.mark.xfail(reason="fillna converts to value's unit"),
- ),
+ pytest.param(unit_registry.cm, None, id="compatible_unit"),
pytest.param(unit_registry.m, None, id="identical_unit"),
),
)
@@ -2629,7 +2739,8 @@ def test_fillna(self, fill_value, unit, error, dtype):
)
actual = func(data_array, value=value)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
def test_dropna(self, dtype):
array = (
@@ -2643,18 +2754,13 @@ def test_dropna(self, dtype):
expected = attach_units(strip_units(data_array).dropna(dim="x"), units)
actual = data_array.dropna(dim="x")
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"unit",
(
- pytest.param(
- 1,
- id="no_unit",
- marks=pytest.mark.xfail(
- reason="pint's isin implementation does not work well with mixed args"
- ),
- ),
+ pytest.param(1, id="no_unit"),
pytest.param(unit_registry.dimensionless, id="dimensionless"),
pytest.param(unit_registry.s, id="incompatible_unit"),
pytest.param(unit_registry.cm, id="compatible_unit"),
@@ -2677,22 +2783,11 @@ def test_isin(self, unit, dtype):
) & array.check(unit)
actual = data_array.isin(values)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
- "variant",
- (
- pytest.param(
- "masking",
- marks=pytest.mark.xfail(reason="array(nan) is not a quantity"),
- ),
- "replacing_scalar",
- "replacing_array",
- pytest.param(
- "dropping",
- marks=pytest.mark.xfail(reason="array(nan) is not a quantity"),
- ),
- ),
+ "variant", ("masking", "replacing_scalar", "replacing_array", "dropping")
)
@pytest.mark.parametrize(
"unit,error",
@@ -2742,22 +2837,24 @@ def test_where(self, variant, unit, error, dtype):
)
actual = data_array.where(**kwargs)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
- @pytest.mark.xfail(reason="interpolate strips units")
- def test_interpolate_na(self, dtype):
+ @pytest.mark.xfail(reason="uses numpy.vectorize")
+ def test_interpolate_na(self):
array = (
np.array([-1.03, 0.1, 1.4, np.nan, 2.3, np.nan, np.nan, 9.1])
* unit_registry.m
)
x = np.arange(len(array))
- data_array = xr.DataArray(data=array, coords={"x": x}, dims="x").astype(dtype)
+ data_array = xr.DataArray(data=array, coords={"x": x}, dims="x")
units = extract_units(data_array)
expected = attach_units(strip_units(data_array).interpolate_na(dim="x"), units)
actual = data_array.interpolate_na(dim="x")
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"unit,error",
@@ -2767,18 +2864,8 @@ def test_interpolate_na(self, dtype):
unit_registry.dimensionless, DimensionalityError, id="dimensionless"
),
pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
- pytest.param(
- unit_registry.cm,
- None,
- id="compatible_unit",
- marks=pytest.mark.xfail(reason="depends on reindex"),
- ),
- pytest.param(
- unit_registry.m,
- None,
- id="identical_unit",
- marks=pytest.mark.xfail(reason="depends on reindex"),
- ),
+ pytest.param(unit_registry.cm, None, id="compatible_unit",),
+ pytest.param(unit_registry.m, None, id="identical_unit",),
),
)
def test_combine_first(self, unit, error, dtype):
@@ -2807,7 +2894,8 @@ def test_combine_first(self, unit, error, dtype):
)
actual = data_array.combine_first(other)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"unit",
@@ -2829,7 +2917,17 @@ def test_combine_first(self, unit, error, dtype):
"coords",
),
)
- @pytest.mark.parametrize("func", (method("equals"), method("identical")), ids=repr)
+ @pytest.mark.parametrize(
+ "func",
+ (
+ method("equals"),
+ pytest.param(
+ method("identical"),
+ marks=pytest.mark.skip(reason="the behavior of identical is undecided"),
+ ),
+ ),
+ ids=repr,
+ )
def test_comparisons(self, func, variation, unit, dtype):
def is_compatible(a, b):
a = a if a is not None else 1
@@ -2903,7 +3001,8 @@ def test_broadcast_like(self, unit, dtype):
)
actual = arr1.broadcast_like(arr2)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"unit",
@@ -2950,7 +3049,6 @@ def test_broadcast_equals(self, unit, dtype):
method("reset_coords", names="x2"),
method("copy"),
method("astype", np.float32),
- method("item", 1),
),
ids=repr,
)
@@ -2978,7 +3076,8 @@ def test_content_manipulation(self, func, dtype):
expected = attach_units(func(strip_units(data_array), **stripped_kwargs), units)
actual = func(data_array)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"func", (pytest.param(method("copy", data=np.arange(20))),), ids=repr
@@ -3004,7 +3103,9 @@ def test_content_manipulation_with_units(self, func, unit, dtype):
)
actual = func(data_array, **kwargs)
- assert_equal_with_units(expected, actual)
+
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"indices",
@@ -3024,7 +3125,8 @@ def test_isel(self, indices, dtype):
)
actual = data_array.isel(x=indices)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.xfail(reason="indexes don't support units")
@pytest.mark.parametrize(
@@ -3067,7 +3169,9 @@ def test_sel(self, raw_values, unit, error, dtype):
extract_units(data_array),
)
actual = data_array.sel(x=values)
- assert_equal_with_units(expected, actual)
+
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.xfail(reason="indexes don't support units")
@pytest.mark.parametrize(
@@ -3110,7 +3214,9 @@ def test_loc(self, raw_values, unit, error, dtype):
extract_units(data_array),
)
actual = data_array.loc[{"x": values}]
- assert_equal_with_units(expected, actual)
+
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.xfail(reason="indexes don't support units")
@pytest.mark.parametrize(
@@ -3153,7 +3259,9 @@ def test_drop_sel(self, raw_values, unit, error, dtype):
extract_units(data_array),
)
actual = data_array.drop_sel(x=values)
- assert_equal_with_units(expected, actual)
+
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"shape",
@@ -3181,7 +3289,9 @@ def test_squeeze(self, shape, dtype):
strip_units(data_array).squeeze(), extract_units(data_array)
)
actual = data_array.squeeze()
- assert_equal_with_units(expected, actual)
+
+ 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)
@@ -3190,7 +3300,9 @@ def test_squeeze(self, shape, dtype):
strip_units(data_array).squeeze(dim=name), extract_units(data_array)
)
actual = data_array.squeeze(dim=name)
- assert_equal_with_units(expected, actual)
+
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"func",
@@ -3212,49 +3324,46 @@ def test_head_tail_thin(self, func, dtype):
)
actual = func(data_array)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
- @pytest.mark.xfail(reason="indexes don't support units")
+ # 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(
- "unit,error",
+ "func",
(
- pytest.param(1, DimensionalityError, id="no_unit"),
pytest.param(
- unit_registry.dimensionless, DimensionalityError, id="dimensionless"
+ method("interp"), marks=pytest.mark.xfail(reason="uses scipy")
),
- pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
- pytest.param(unit_registry.cm, None, id="compatible_unit"),
- pytest.param(unit_registry.m, None, id="identical_unit"),
+ method("reindex"),
),
+ ids=repr,
)
- def test_interp(self, unit, error):
- array = np.linspace(1, 2, 10 * 5).reshape(10, 5) * unit_registry.degK
- new_coords = (np.arange(10) + 0.5) * unit
- coords = {
- "x": np.arange(10) * unit_registry.m,
- "y": np.arange(5) * unit_registry.m,
+ def test_interp_reindex(self, variant, func, dtype):
+ variants = {
+ "data": (unit_registry.m, 1),
+ "coords": (1, unit_registry.m),
}
+ data_unit, coord_unit = variants.get(variant)
- data_array = xr.DataArray(array, coords=coords, dims=("x", "y"))
+ array = np.linspace(1, 2, 10).astype(dtype) * data_unit
+ y = np.arange(10) * coord_unit
- if error is not None:
- with pytest.raises(error):
- data_array.interp(x=new_coords)
-
- return
+ x = np.arange(10)
+ new_x = np.arange(10) + 0.5
+ data_array = xr.DataArray(array, coords={"x": x, "y": ("x", y)}, dims="x")
units = extract_units(data_array)
- expected = attach_units(
- strip_units(data_array).interp(
- x=strip_units(convert_units(new_coords, {None: unit_registry.m}))
- ),
- units,
- )
- actual = data_array.interp(x=new_coords)
+ expected = attach_units(func(strip_units(data_array), x=new_x), units)
+ actual = func(data_array, x=new_x)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_allclose(expected, actual)
- @pytest.mark.xfail(reason="indexes strip units")
+ @pytest.mark.xfail(reason="indexes don't support units")
@pytest.mark.parametrize(
"unit,error",
(
@@ -3267,79 +3376,70 @@ def test_interp(self, unit, error):
pytest.param(unit_registry.m, None, id="identical_unit"),
),
)
- def test_interp_like(self, unit, error):
- array = np.linspace(1, 2, 10 * 5).reshape(10, 5) * unit_registry.degK
- coords = {
- "x": (np.arange(10) + 0.3) * unit_registry.m,
- "y": (np.arange(5) + 0.3) * unit_registry.m,
- }
-
- data_array = xr.DataArray(array, coords=coords, dims=("x", "y"))
- other = xr.DataArray(
- data=np.empty((20, 10)) * unit_registry.degK,
- coords={"x": np.arange(20) * unit, "y": np.arange(10) * unit},
- dims=("x", "y"),
- )
+ @pytest.mark.parametrize(
+ "func", (method("interp"), method("reindex")), ids=repr,
+ )
+ def test_interp_reindex_indexing(self, func, unit, error, dtype):
+ array = np.linspace(1, 2, 10).astype(dtype)
+ x = np.arange(10) * unit_registry.m
+ new_x = (np.arange(10) + 0.5) * unit
+ data_array = xr.DataArray(array, coords={"x": x}, dims="x")
if error is not None:
with pytest.raises(error):
- data_array.interp_like(other)
+ func(data_array, x=new_x)
return
units = extract_units(data_array)
expected = attach_units(
- strip_units(data_array).interp_like(
- strip_units(convert_units(other, units))
+ func(
+ strip_units(data_array),
+ x=strip_units(convert_units(new_x, {None: unit_registry.m})),
),
units,
)
- actual = data_array.interp_like(other)
+ actual = func(data_array, x=new_x)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
- @pytest.mark.xfail(reason="indexes don't support units")
+ # 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(
- "unit,error",
+ "func",
(
- pytest.param(1, DimensionalityError, id="no_unit"),
pytest.param(
- unit_registry.dimensionless, DimensionalityError, id="dimensionless"
+ method("interp_like"), marks=pytest.mark.xfail(reason="uses scipy")
),
- pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
- pytest.param(unit_registry.cm, None, id="compatible_unit"),
- pytest.param(unit_registry.m, None, id="identical_unit"),
+ method("reindex_like"),
),
+ ids=repr,
)
- def test_reindex(self, unit, error, dtype):
- array = (
- np.linspace(1, 2, 10 * 5).reshape(10, 5).astype(dtype) * unit_registry.degK
- )
- new_coords = (np.arange(10) + 0.5) * unit
- coords = {
- "x": np.arange(10) * unit_registry.m,
- "y": np.arange(5) * unit_registry.m,
+ def test_interp_reindex_like(self, variant, func, dtype):
+ variants = {
+ "data": (unit_registry.m, 1),
+ "coords": (1, unit_registry.m),
}
+ data_unit, coord_unit = variants.get(variant)
- data_array = xr.DataArray(array, coords=coords, dims=("x", "y"))
- func = method("reindex")
-
- if error is not None:
- with pytest.raises(error):
- func(data_array, x=new_coords)
+ array = np.linspace(1, 2, 10).astype(dtype) * data_unit
+ coord = np.arange(10) * coord_unit
- return
+ x = np.arange(10)
+ new_x = np.arange(-2, 2) + 0.5
+ data_array = xr.DataArray(array, coords={"x": x, "y": ("x", coord)}, dims="x")
+ other = xr.DataArray(np.empty_like(new_x), coords={"x": new_x}, dims="x")
- expected = attach_units(
- func(
- strip_units(data_array),
- x=strip_units(convert_units(new_coords, {None: unit_registry.m})),
- ),
- {None: unit_registry.degK},
- )
- actual = func(data_array, x=new_coords)
+ units = extract_units(data_array)
+ expected = attach_units(func(strip_units(data_array), other), units)
+ actual = func(data_array, other)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_allclose(expected, actual)
@pytest.mark.xfail(reason="indexes don't support units")
@pytest.mark.parametrize(
@@ -3354,38 +3454,35 @@ def test_reindex(self, unit, error, dtype):
pytest.param(unit_registry.m, None, id="identical_unit"),
),
)
- def test_reindex_like(self, unit, error, dtype):
- array = (
- np.linspace(1, 2, 10 * 5).reshape(10, 5).astype(dtype) * unit_registry.degK
- )
- coords = {
- "x": (np.arange(10) + 0.3) * unit_registry.m,
- "y": (np.arange(5) + 0.3) * unit_registry.m,
- }
+ @pytest.mark.parametrize(
+ "func", (method("interp_like"), method("reindex_like")), ids=repr,
+ )
+ def test_interp_reindex_like_indexing(self, func, unit, error, dtype):
+ array = np.linspace(1, 2, 10).astype(dtype)
+ x = np.arange(10) * unit_registry.m
+ new_x = (np.arange(-2, 2) + 0.5) * unit
- data_array = xr.DataArray(array, coords=coords, dims=("x", "y"))
- other = xr.DataArray(
- data=np.empty((20, 10)) * unit_registry.degK,
- coords={"x": np.arange(20) * unit, "y": np.arange(10) * unit},
- dims=("x", "y"),
- )
+ data_array = xr.DataArray(array, coords={"x": x}, dims="x")
+ other = xr.DataArray(np.empty_like(new_x), {"x": new_x}, dims="x")
if error is not None:
with pytest.raises(error):
- data_array.reindex_like(other)
+ func(data_array, other)
return
units = extract_units(data_array)
expected = attach_units(
- strip_units(data_array).reindex_like(
- strip_units(convert_units(other, units))
+ func(
+ strip_units(data_array),
+ strip_units(convert_units(other, {None: unit_registry.m})),
),
units,
)
- actual = data_array.reindex_like(other)
+ actual = func(data_array, other)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"func",
@@ -3407,7 +3504,8 @@ def test_stacking_stacked(self, func, dtype):
expected = attach_units(func(strip_units(stacked)), {"data": unit_registry.m})
actual = func(stacked)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.xfail(reason="indexes don't support units")
def test_to_unstacked_dataset(self, dtype):
@@ -3430,7 +3528,8 @@ def test_to_unstacked_dataset(self, dtype):
).rename({elem.magnitude: elem for elem in x})
actual = func(data_array)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"func",
@@ -3438,9 +3537,7 @@ def test_to_unstacked_dataset(self, dtype):
method("transpose", "y", "x", "z"),
method("stack", a=("x", "y")),
method("set_index", x="x2"),
- pytest.param(
- method("shift", x=2), marks=pytest.mark.xfail(reason="strips units")
- ),
+ method("shift", x=2),
method("roll", x=2, roll_coords=False),
method("sortby", "x2"),
),
@@ -3466,7 +3563,8 @@ def test_stacking_reordering(self, func, dtype):
expected = attach_units(func(strip_units(data_array)), {None: unit_registry.m})
actual = func(data_array)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"func",
@@ -3476,16 +3574,13 @@ def test_stacking_reordering(self, func, dtype):
method("integrate", dim="x"),
pytest.param(
method("quantile", q=[0.25, 0.75]),
- marks=pytest.mark.xfail(reason="nanquantile not implemented"),
- ),
- method("reduce", func=np.sum, dim="x"),
- pytest.param(
- lambda x: x.dot(x),
- id="method_dot",
marks=pytest.mark.xfail(
- reason="pint does not implement the dot method"
+ LooseVersion(pint.__version__) <= "0.12",
+ reason="quantile / nanquantile not implemented yet",
),
),
+ method("reduce", func=np.sum, dim="x"),
+ pytest.param(lambda x: x.dot(x), id="method_dot"),
),
ids=repr,
)
@@ -3512,8 +3607,13 @@ def test_computation(self, func, dtype):
expected = attach_units(func(strip_units(data_array)), units)
actual = func(data_array)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.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(
"func",
(
@@ -3522,11 +3622,15 @@ def test_computation(self, func, dtype):
method("coarsen", y=2),
pytest.param(
method("rolling", y=3),
- marks=pytest.mark.xfail(reason="rolling strips units"),
+ marks=pytest.mark.xfail(
+ reason="numpy.lib.stride_tricks.as_strided converts to ndarray"
+ ),
),
pytest.param(
method("rolling_exp", y=3),
- marks=pytest.mark.xfail(reason="units not supported by numbagg"),
+ marks=pytest.mark.xfail(
+ reason="numbagg functions are not supported by pint"
+ ),
),
),
ids=repr,
@@ -3545,7 +3649,8 @@ def test_computation_objects(self, func, dtype):
expected = attach_units(func(strip_units(data_array)).mean(), units)
actual = func(data_array).mean()
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_allclose(expected, actual)
def test_resample(self, dtype):
array = np.linspace(0, 5, 10).astype(dtype) * unit_registry.m
@@ -3559,7 +3664,8 @@ def test_resample(self, dtype):
expected = attach_units(func(strip_units(data_array)).mean(), units)
actual = func(data_array).mean()
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
@pytest.mark.parametrize(
"func",
@@ -3569,7 +3675,10 @@ def test_resample(self, dtype):
method("last"),
pytest.param(
method("quantile", q=[0.25, 0.5, 0.75], dim="x"),
- marks=pytest.mark.xfail(reason="nanquantile not implemented"),
+ marks=pytest.mark.xfail(
+ LooseVersion(pint.__version__) <= "0.12",
+ reason="quantile / nanquantile not implemented yet",
+ ),
),
),
ids=repr,
@@ -3598,18 +3707,20 @@ def test_grouped_operations(self, func, dtype):
)
actual = func(data_array.groupby("y"))
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ xr.testing.assert_identical(expected, actual)
+@pytest.mark.filterwarnings("error::pint.UnitStrippedWarning")
class TestDataset:
@pytest.mark.parametrize(
"unit,error",
(
- pytest.param(1, DimensionalityError, id="no_unit"),
+ pytest.param(1, xr.MergeError, id="no_unit"),
pytest.param(
- unit_registry.dimensionless, DimensionalityError, id="dimensionless"
+ unit_registry.dimensionless, xr.MergeError, id="dimensionless"
),
- pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
+ pytest.param(unit_registry.s, xr.MergeError, id="incompatible_unit"),
pytest.param(unit_registry.mm, None, id="compatible_unit"),
pytest.param(unit_registry.m, None, id="same_unit"),
),
@@ -3618,11 +3729,10 @@ class TestDataset:
"shared",
(
"nothing",
- pytest.param("dims", marks=pytest.mark.xfail(reason="indexes strip units")),
pytest.param(
- "coords",
- marks=pytest.mark.xfail(reason="reindex does not work with pint yet"),
+ "dims", marks=pytest.mark.xfail(reason="indexes don't support units")
),
+ "coords",
),
)
def test_init(self, shared, unit, error, dtype):
@@ -3630,60 +3740,53 @@ def test_init(self, shared, unit, error, dtype):
scaled_unit = unit_registry.mm
a = np.linspace(0, 1, 10).astype(dtype) * unit_registry.Pa
- b = np.linspace(-1, 0, 12).astype(dtype) * unit_registry.Pa
-
- raw_x = np.arange(a.shape[0])
- x = raw_x * original_unit
- x2 = x.to(scaled_unit)
-
- raw_y = np.arange(b.shape[0])
- y = raw_y * unit
- y_units = unit if isinstance(y, unit_registry.Quantity) else None
- if isinstance(y, unit_registry.Quantity):
- if y.check(scaled_unit):
- y2 = y.to(scaled_unit)
- else:
- y2 = y * 1000
- y2_units = y2.units
- else:
- y2 = y * 1000
- y2_units = None
+ b = np.linspace(-1, 0, 10).astype(dtype) * unit_registry.degK
+
+ values_a = np.arange(a.shape[0])
+ dim_a = values_a * original_unit
+ coord_a = dim_a.to(scaled_unit)
+
+ values_b = np.arange(b.shape[0])
+ dim_b = values_b * unit
+ coord_b = (
+ dim_b.to(scaled_unit)
+ if unit_registry.is_compatible_with(dim_b, scaled_unit)
+ and unit != scaled_unit
+ else dim_b * 1000
+ )
variants = {
- "nothing": ({"x": x, "x2": ("x", x2)}, {"y": y, "y2": ("y", y2)}),
- "dims": (
- {"x": x, "x2": ("x", strip_units(x2))},
- {"x": y, "y2": ("x", strip_units(y2))},
+ "nothing": ({}, {}),
+ "dims": ({"x": dim_a}, {"x": dim_b}),
+ "coords": (
+ {"x": values_a, "y": ("x", coord_a)},
+ {"x": values_b, "y": ("x", coord_b)},
),
- "coords": ({"x": raw_x, "y": ("x", x2)}, {"x": raw_y, "y": ("x", y2)}),
}
coords_a, coords_b = variants.get(shared)
dims_a, dims_b = ("x", "y") if shared == "nothing" else ("x", "x")
- arr1 = xr.DataArray(data=a, coords=coords_a, dims=dims_a)
- arr2 = xr.DataArray(data=b, coords=coords_b, dims=dims_b)
+ a = xr.DataArray(data=a, coords=coords_a, dims=dims_a)
+ b = xr.DataArray(data=b, coords=coords_b, dims=dims_b)
+
if error is not None and shared != "nothing":
with pytest.raises(error):
- xr.Dataset(data_vars={"a": arr1, "b": arr2})
+ xr.Dataset(data_vars={"a": a, "b": b})
return
- actual = xr.Dataset(data_vars={"a": arr1, "b": arr2})
+ actual = xr.Dataset(data_vars={"a": a, "b": b})
- expected_units = {
- "a": a.units,
- "b": b.units,
- "x": x.units,
- "x2": x2.units,
- "y": y_units,
- "y2": y2_units,
- }
+ units = merge_mappings(
+ extract_units(a.rename("a")), extract_units(b.rename("b"))
+ )
expected = attach_units(
- xr.Dataset(data_vars={"a": strip_units(arr1), "b": strip_units(arr2)}),
- expected_units,
+ xr.Dataset(data_vars={"a": strip_units(a), "b": strip_units(b)}), units
)
- assert_equal_with_units(actual, expected)
+
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"func", (pytest.param(str, id="str"), pytest.param(repr, id="repr"))
@@ -3691,79 +3794,79 @@ def test_init(self, shared, unit, error, dtype):
@pytest.mark.parametrize(
"variant",
(
+ "data",
pytest.param(
- "with_dims",
+ "dims",
marks=pytest.mark.xfail(reason="units in indexes are not supported"),
),
- pytest.param("with_coords"),
- pytest.param("without_coords"),
+ "coords",
),
)
- @pytest.mark.filterwarnings("error:::pint[.*]")
def test_repr(self, func, variant, dtype):
- array1 = np.linspace(1, 2, 10, dtype=dtype) * unit_registry.Pa
- array2 = np.linspace(0, 1, 10, dtype=dtype) * unit_registry.degK
+ unit1, unit2 = (
+ (unit_registry.Pa, unit_registry.degK) if variant == "data" else (1, 1)
+ )
+
+ array1 = np.linspace(1, 2, 10, dtype=dtype) * unit1
+ array2 = np.linspace(0, 1, 10, dtype=dtype) * unit2
x = np.arange(len(array1)) * unit_registry.s
y = x.to(unit_registry.ms)
variants = {
- "with_dims": {"x": x},
- "with_coords": {"y": ("x", y)},
- "without_coords": {},
+ "dims": {"x": x},
+ "coords": {"y": ("x", y)},
+ "data": {},
}
- data_array = xr.Dataset(
+ ds = xr.Dataset(
data_vars={"a": ("x", array1), "b": ("x", array2)},
coords=variants.get(variant),
)
# FIXME: this just checks that the repr does not raise
# warnings or errors, but does not check the result
- func(data_array)
+ func(ds)
@pytest.mark.parametrize(
"func",
(
+ function("all"),
+ function("any"),
pytest.param(
- function("all"),
- marks=pytest.mark.xfail(reason="not implemented by pint"),
+ function("argmax"),
+ marks=pytest.mark.skip(
+ reason="calling np.argmax as a function on xarray objects is not "
+ "supported"
+ ),
),
pytest.param(
- function("any"),
- marks=pytest.mark.xfail(reason="not implemented by pint"),
+ function("argmin"),
+ marks=pytest.mark.skip(
+ reason="calling np.argmin as a function on xarray objects is not "
+ "supported"
+ ),
),
- function("argmax"),
- function("argmin"),
function("max"),
function("min"),
function("mean"),
pytest.param(
function("median"),
- marks=pytest.mark.xfail(
- reason="np.median does not work with dataset yet"
- ),
+ marks=pytest.mark.xfail(reason="median does not work with dataset yet"),
),
function("sum"),
pytest.param(
function("prod"),
- marks=pytest.mark.xfail(reason="not implemented by pint"),
+ marks=pytest.mark.xfail(reason="prod does not work with dataset yet"),
),
function("std"),
function("var"),
function("cumsum"),
- pytest.param(
- function("cumprod"),
- marks=pytest.mark.xfail(reason="fails within xarray"),
- ),
- pytest.param(
- method("all"), marks=pytest.mark.xfail(reason="not implemented by pint")
- ),
- pytest.param(
- method("any"), marks=pytest.mark.xfail(reason="not implemented by pint")
- ),
- method("argmax"),
- method("argmin"),
+ function("cumprod"),
+ method("all"),
+ method("any"),
+ method("argmax", dim="x"),
+ method("argmin", dim="x"),
method("max"),
method("min"),
method("mean"),
@@ -3771,68 +3874,64 @@ def test_repr(self, func, variant, dtype):
method("sum"),
pytest.param(
method("prod"),
- marks=pytest.mark.xfail(reason="not implemented by pint"),
+ marks=pytest.mark.xfail(reason="prod does not work with dataset yet"),
),
method("std"),
method("var"),
method("cumsum"),
- pytest.param(
- method("cumprod"), marks=pytest.mark.xfail(reason="fails within xarray")
- ),
+ method("cumprod"),
),
ids=repr,
)
def test_aggregation(self, func, dtype):
- unit_a = (
- unit_registry.Pa if func.name != "cumprod" else unit_registry.dimensionless
- )
- unit_b = (
- unit_registry.kg / unit_registry.m ** 3
+ unit_a, unit_b = (
+ (unit_registry.Pa, unit_registry.degK)
if func.name != "cumprod"
- else unit_registry.dimensionless
- )
- a = xr.DataArray(data=np.linspace(0, 1, 10).astype(dtype) * unit_a, dims="x")
- b = xr.DataArray(data=np.linspace(-1, 0, 10).astype(dtype) * unit_b, dims="x")
- x = xr.DataArray(data=np.arange(10).astype(dtype) * unit_registry.m, dims="x")
- y = xr.DataArray(
- data=np.arange(10, 20).astype(dtype) * unit_registry.s, dims="x"
+ else (unit_registry.dimensionless, unit_registry.dimensionless)
)
- ds = xr.Dataset(data_vars={"a": a, "b": b}, coords={"x": x, "y": y})
+ a = np.linspace(0, 1, 10).astype(dtype) * unit_a
+ b = np.linspace(-1, 0, 10).astype(dtype) * unit_b
+
+ ds = xr.Dataset({"a": ("x", a), "b": ("x", b)})
+
+ if "dim" in func.kwargs:
+ numpy_kwargs = func.kwargs.copy()
+ dim = numpy_kwargs.pop("dim")
+
+ axis_a = ds.a.get_axis_num(dim)
+ axis_b = ds.b.get_axis_num(dim)
+
+ numpy_kwargs_a = numpy_kwargs.copy()
+ numpy_kwargs_a["axis"] = axis_a
+ numpy_kwargs_b = numpy_kwargs.copy()
+ numpy_kwargs_b["axis"] = axis_b
+ else:
+ numpy_kwargs_a = {}
+ numpy_kwargs_b = {}
+
+ units_a = array_extract_units(func(a, **numpy_kwargs_a))
+ units_b = array_extract_units(func(b, **numpy_kwargs_b))
+ units = {"a": units_a, "b": units_b}
actual = func(ds)
- expected = attach_units(
- func(strip_units(ds)),
- {
- "a": extract_units(func(a)).get(None),
- "b": extract_units(func(b)).get(None),
- },
- )
+ expected = attach_units(func(strip_units(ds)), units)
- assert_equal_with_units(actual, expected)
+ assert_units_equal(expected, actual)
+ assert_allclose(expected, actual)
@pytest.mark.parametrize("property", ("imag", "real"))
def test_numpy_properties(self, property, dtype):
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(
- data=np.linspace(0, 1, 10) * unit_registry.Pa, dims="x"
- ),
- "b": xr.DataArray(
- data=np.linspace(-1, 0, 15) * unit_registry.Pa, dims="y"
- ),
- },
- coords={
- "x": np.arange(10) * unit_registry.m,
- "y": np.arange(15) * unit_registry.s,
- },
- )
+ a = np.linspace(0, 1, 10) * unit_registry.Pa
+ b = np.linspace(-1, 0, 15) * unit_registry.degK
+ ds = xr.Dataset({"a": ("x", a), "b": ("y", b)})
units = extract_units(ds)
actual = getattr(ds, property)
expected = attach_units(getattr(strip_units(ds), property), units)
- assert_equal_with_units(actual, expected)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"func",
@@ -3846,31 +3945,19 @@ def test_numpy_properties(self, property, dtype):
ids=repr,
)
def test_numpy_methods(self, func, dtype):
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(
- data=np.linspace(1, -1, 10) * unit_registry.Pa, dims="x"
- ),
- "b": xr.DataArray(
- data=np.linspace(-1, 1, 15) * unit_registry.Pa, dims="y"
- ),
- },
- coords={
- "x": np.arange(10) * unit_registry.m,
- "y": np.arange(15) * unit_registry.s,
- },
- )
- units = {
- "a": array_extract_units(func(ds.a)),
- "b": array_extract_units(func(ds.b)),
- "x": unit_registry.m,
- "y": unit_registry.s,
- }
+ a = np.linspace(1, -1, 10) * unit_registry.Pa
+ b = np.linspace(-1, 1, 15) * unit_registry.degK
+ ds = xr.Dataset({"a": ("x", a), "b": ("y", b)})
+
+ units_a = array_extract_units(func(a))
+ units_b = array_extract_units(func(b))
+ units = {"a": units_a, "b": units_b}
actual = func(ds)
expected = attach_units(func(strip_units(ds)), units)
- assert_equal_with_units(actual, expected)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize("func", (method("clip", min=3, max=8),), ids=repr)
@pytest.mark.parametrize(
@@ -3887,21 +3974,13 @@ def test_numpy_methods(self, func, dtype):
)
def test_numpy_methods_with_args(self, func, unit, error, dtype):
data_unit = unit_registry.m
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=np.arange(10) * data_unit, dims="x"),
- "b": xr.DataArray(data=np.arange(15) * data_unit, dims="y"),
- },
- coords={
- "x": np.arange(10) * unit_registry.m,
- "y": np.arange(15) * unit_registry.s,
- },
- )
+ a = np.linspace(0, 10, 15) * unit_registry.m
+ b = np.linspace(-2, 12, 20) * unit_registry.m
+ ds = xr.Dataset({"a": ("x", a), "b": ("y", b)})
units = extract_units(ds)
kwargs = {
- key: (value * unit if isinstance(value, (int, float)) else value)
- for key, value in func.kwargs.items()
+ key: array_attach_units(value, unit) for key, value in func.kwargs.items()
}
if error is not None:
@@ -3918,7 +3997,8 @@ def test_numpy_methods_with_args(self, func, unit, error, dtype):
actual = func(ds, **kwargs)
expected = attach_units(func(strip_units(ds), **stripped_kwargs), units)
- assert_equal_with_units(actual, expected)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"func", (method("isnull"), method("notnull"), method("count")), ids=repr
@@ -3948,22 +4028,13 @@ def test_missing_value_detection(self, func, dtype):
* unit_registry.Pa
)
- x = np.arange(array1.shape[0]) * unit_registry.m
- y = np.arange(array1.shape[1]) * unit_registry.m
- z = np.arange(array2.shape[0]) * unit_registry.m
-
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims=("x", "y")),
- "b": xr.DataArray(data=array2, dims=("z", "x")),
- },
- coords={"x": x, "y": y, "z": z},
- )
+ ds = xr.Dataset({"a": (("x", "y"), array1), "b": (("z", "x"), array2)})
expected = func(strip_units(ds))
actual = func(ds)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.xfail(reason="ffill and bfill lose the unit")
@pytest.mark.parametrize("func", (method("ffill"), method("bfill")), ids=repr)
@@ -3977,23 +4048,14 @@ def test_missing_value_filling(self, func, dtype):
* unit_registry.Pa
)
- x = np.arange(len(array1))
-
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims="x"),
- "b": xr.DataArray(data=array2, dims="x"),
- },
- coords={"x": x},
- )
+ ds = xr.Dataset({"a": ("x", array1), "b": ("y", array2)})
+ units = extract_units(ds)
- expected = attach_units(
- func(strip_units(ds), dim="x"),
- {"a": unit_registry.degK, "b": unit_registry.Pa},
- )
+ expected = attach_units(func(strip_units(ds), dim="x"), units)
actual = func(ds, dim="x")
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"unit,error",
@@ -4003,14 +4065,7 @@ def test_missing_value_filling(self, func, dtype):
unit_registry.dimensionless, DimensionalityError, id="dimensionless"
),
pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
- pytest.param(
- unit_registry.cm,
- None,
- id="compatible_unit",
- marks=pytest.mark.xfail(
- reason="where converts the array, not the fill value"
- ),
- ),
+ pytest.param(unit_registry.cm, None, id="compatible_unit",),
pytest.param(unit_registry.m, None, id="identical_unit"),
),
)
@@ -4031,30 +4086,26 @@ def test_fillna(self, fill_value, unit, error, dtype):
np.array([4.3, 9.8, 7.5, np.nan, 8.2, np.nan]).astype(dtype)
* unit_registry.m
)
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims="x"),
- "b": xr.DataArray(data=array2, dims="x"),
- }
- )
+ ds = xr.Dataset({"a": ("x", array1), "b": ("x", array2)})
+ value = fill_value * unit
+ units = extract_units(ds)
if error is not None:
with pytest.raises(error):
- ds.fillna(value=fill_value * unit)
+ ds.fillna(value=value)
return
- actual = ds.fillna(value=fill_value * unit)
+ actual = ds.fillna(value=value)
expected = attach_units(
strip_units(ds).fillna(
- value=strip_units(
- convert_units(fill_value * unit, {None: unit_registry.m})
- )
+ value=strip_units(convert_units(value, {None: unit_registry.m}))
),
- {"a": unit_registry.m, "b": unit_registry.m},
+ units,
)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
def test_dropna(self, dtype):
array1 = (
@@ -4065,22 +4116,14 @@ def test_dropna(self, dtype):
np.array([4.3, 9.8, 7.5, np.nan, 8.2, np.nan]).astype(dtype)
* unit_registry.Pa
)
- x = np.arange(len(array1))
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims="x"),
- "b": xr.DataArray(data=array2, dims="x"),
- },
- coords={"x": x},
- )
+ ds = xr.Dataset({"a": ("x", array1), "b": ("x", array2)})
+ units = extract_units(ds)
- expected = attach_units(
- strip_units(ds).dropna(dim="x"),
- {"a": unit_registry.degK, "b": unit_registry.Pa},
- )
+ expected = attach_units(strip_units(ds).dropna(dim="x"), units)
actual = ds.dropna(dim="x")
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"unit",
@@ -4101,34 +4144,28 @@ def test_isin(self, unit, dtype):
np.array([4.3, 9.8, 7.5, np.nan, 8.2, np.nan]).astype(dtype)
* unit_registry.m
)
- x = np.arange(len(array1))
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims="x"),
- "b": xr.DataArray(data=array2, dims="x"),
- },
- coords={"x": x},
- )
+ ds = xr.Dataset({"a": ("x", array1), "b": ("x", array2)})
raw_values = np.array([1.4, np.nan, 2.3]).astype(dtype)
values = raw_values * unit
- if (
- isinstance(values, unit_registry.Quantity)
- and values.check(unit_registry.m)
- and unit != unit_registry.m
- ):
- raw_values = values.to(unit_registry.m).magnitude
+ converted_values = (
+ convert_units(values, {None: unit_registry.m})
+ if is_compatible(unit, unit_registry.m)
+ else values
+ )
- expected = strip_units(ds).isin(raw_values)
- if not isinstance(values, unit_registry.Quantity) or not values.check(
- unit_registry.m
- ):
+ expected = strip_units(ds).isin(strip_units(converted_values))
+ # TODO: use `unit_registry.is_compatible_with(unit, unit_registry.m)` instead.
+ # Needs `pint>=0.12.1`, though, so we probably should wait until that is released.
+ if not is_compatible(unit, unit_registry.m):
expected.a[:] = False
expected.b[:] = False
+
actual = ds.isin(values)
- assert_equal_with_units(actual, expected)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"variant", ("masking", "replacing_scalar", "replacing_array", "dropping")
@@ -4150,13 +4187,8 @@ def test_where(self, variant, unit, error, dtype):
array1 = np.linspace(0, 1, 10).astype(dtype) * original_unit
array2 = np.linspace(-1, 0, 10).astype(dtype) * original_unit
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims="x"),
- "b": xr.DataArray(data=array2, dims="x"),
- },
- coords={"x": np.arange(len(array1))},
- )
+ ds = xr.Dataset({"a": ("x", array1), "b": ("x", array2)})
+ units = extract_units(ds)
condition = ds < 0.5 * original_unit
other = np.linspace(-2, -1, 10).astype(dtype) * unit
@@ -4178,15 +4210,13 @@ def test_where(self, variant, unit, error, dtype):
for key, value in kwargs.items()
}
- expected = attach_units(
- strip_units(ds).where(**kwargs_without_units),
- {"a": original_unit, "b": original_unit},
- )
+ expected = attach_units(strip_units(ds).where(**kwargs_without_units), units,)
actual = ds.where(**kwargs)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
- @pytest.mark.xfail(reason="interpolate strips units")
+ @pytest.mark.xfail(reason="interpolate_na uses numpy.vectorize")
def test_interpolate_na(self, dtype):
array1 = (
np.array([1.4, np.nan, 2.3, np.nan, np.nan, 9.1]).astype(dtype)
@@ -4196,24 +4226,15 @@ def test_interpolate_na(self, dtype):
np.array([4.3, 9.8, 7.5, np.nan, 8.2, np.nan]).astype(dtype)
* unit_registry.Pa
)
- x = np.arange(len(array1))
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims="x"),
- "b": xr.DataArray(data=array2, dims="x"),
- },
- coords={"x": x},
- )
+ ds = xr.Dataset({"a": ("x", array1), "b": ("x", array2)})
+ units = extract_units(ds)
- expected = attach_units(
- strip_units(ds).interpolate_na(dim="x"),
- {"a": unit_registry.degK, "b": unit_registry.Pa},
- )
+ expected = attach_units(strip_units(ds).interpolate_na(dim="x"), units,)
actual = ds.interpolate_na(dim="x")
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
- @pytest.mark.xfail(reason="wrong argument order for `where`")
@pytest.mark.parametrize(
"unit,error",
(
@@ -4226,31 +4247,40 @@ def test_interpolate_na(self, dtype):
pytest.param(unit_registry.m, None, id="same_unit"),
),
)
- def test_combine_first(self, unit, error, dtype):
+ @pytest.mark.parametrize(
+ "variant",
+ (
+ "data",
+ pytest.param(
+ "dims", marks=pytest.mark.xfail(reason="indexes don't support units"),
+ ),
+ ),
+ )
+ def test_combine_first(self, variant, unit, error, dtype):
+ variants = {
+ "data": (unit_registry.m, unit, 1, 1),
+ "dims": (1, 1, unit_registry.m, unit),
+ }
+ data_unit, other_data_unit, dims_unit, other_dims_unit = variants.get(variant)
+
array1 = (
- np.array([1.4, np.nan, 2.3, np.nan, np.nan, 9.1]).astype(dtype)
- * unit_registry.m
+ np.array([1.4, np.nan, 2.3, np.nan, np.nan, 9.1]).astype(dtype) * data_unit
)
array2 = (
- np.array([4.3, 9.8, 7.5, np.nan, 8.2, np.nan]).astype(dtype)
- * unit_registry.m
+ np.array([4.3, 9.8, 7.5, np.nan, 8.2, np.nan]).astype(dtype) * data_unit
)
- x = np.arange(len(array1))
+ x = np.arange(len(array1)) * dims_unit
ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims="x"),
- "b": xr.DataArray(data=array2, dims="x"),
- },
- coords={"x": x},
+ data_vars={"a": ("x", array1), "b": ("x", array2)}, coords={"x": x},
)
- other_array1 = np.ones_like(array1) * unit
- other_array2 = -1 * np.ones_like(array2) * unit
+ units = extract_units(ds)
+
+ other_array1 = np.ones_like(array1) * other_data_unit
+ other_array2 = np.full_like(array2, fill_value=-1) * other_data_unit
+ other_x = (np.arange(array1.shape[0]) + 5) * other_dims_unit
other = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=other_array1, dims="x"),
- "b": xr.DataArray(data=other_array2, dims="x"),
- },
- coords={"x": np.arange(array1.shape[0])},
+ data_vars={"a": ("x", other_array1), "b": ("x", other_array2)},
+ coords={"x": other_x},
)
if error is not None:
@@ -4260,16 +4290,13 @@ def test_combine_first(self, unit, error, dtype):
return
expected = attach_units(
- strip_units(ds).combine_first(
- strip_units(
- convert_units(other, {"a": unit_registry.m, "b": unit_registry.m})
- )
- ),
- {"a": unit_registry.m, "b": unit_registry.m},
+ strip_units(ds).combine_first(strip_units(convert_units(other, units))),
+ units,
)
actual = ds.combine_first(other)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"unit",
@@ -4282,7 +4309,7 @@ def test_combine_first(self, unit, error, dtype):
),
)
@pytest.mark.parametrize(
- "variation",
+ "variant",
(
"data",
pytest.param(
@@ -4291,50 +4318,67 @@ def test_combine_first(self, unit, error, dtype):
"coords",
),
)
- @pytest.mark.parametrize("func", (method("equals"), method("identical")), ids=repr)
- def test_comparisons(self, func, variation, unit, dtype):
- def is_compatible(a, b):
- a = a if a is not None else 1
- b = b if b is not None else 1
- quantity = np.arange(5) * a
-
- return a == b or quantity.check(b)
-
+ @pytest.mark.parametrize(
+ "func",
+ (
+ method("equals"),
+ pytest.param(
+ method("identical"),
+ marks=pytest.mark.skip("behaviour of identical is unclear"),
+ ),
+ ),
+ ids=repr,
+ )
+ def test_comparisons(self, func, variant, unit, dtype):
array1 = np.linspace(0, 5, 10).astype(dtype)
array2 = np.linspace(-5, 0, 10).astype(dtype)
coord = np.arange(len(array1)).astype(dtype)
- original_unit = unit_registry.m
- quantity1 = array1 * original_unit
- quantity2 = array2 * original_unit
- x = coord * original_unit
- y = coord * original_unit
+ 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)
- units = {"data": (unit, 1, 1), "dims": (1, unit, 1), "coords": (1, 1, unit)}
- data_unit, dim_unit, coord_unit = units.get(variation)
+ a = array1 * data_unit
+ b = array2 * data_unit
+ x = coord * dim_unit
+ y = coord * coord_unit
ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=quantity1, dims="x"),
- "b": xr.DataArray(data=quantity2, dims="x"),
- },
- coords={"x": x, "y": ("x", y)},
+ data_vars={"a": ("x", a), "b": ("x", b)}, coords={"x": x, "y": ("x", y)},
)
+ units = extract_units(ds)
+
+ other_variants = {
+ "data": (unit, 1, 1),
+ "dims": (1, unit, 1),
+ "coords": (1, 1, unit),
+ }
+ other_data_unit, other_dim_unit, other_coord_unit = other_variants.get(variant)
other_units = {
- "a": data_unit if quantity1.check(data_unit) else None,
- "b": data_unit if quantity2.check(data_unit) else None,
- "x": dim_unit if x.check(dim_unit) else None,
- "y": coord_unit if y.check(coord_unit) else None,
+ "a": other_data_unit,
+ "b": other_data_unit,
+ "x": other_dim_unit,
+ "y": other_coord_unit,
}
- other = attach_units(strip_units(convert_units(ds, other_units)), other_units)
- units = extract_units(ds)
+ to_convert = {
+ key: unit if is_compatible(unit, reference) else None
+ for key, (unit, reference) in zip_mappings(units, other_units)
+ }
+ # convert units where possible, then attach all units to the converted dataset
+ other = attach_units(strip_units(convert_units(ds, to_convert)), other_units)
other_units = extract_units(other)
+ # make sure all units are compatible and only then try to
+ # convert and compare values
equal_ds = all(
- is_compatible(units[name], other_units[name]) for name in units.keys()
+ is_compatible(unit, other_unit)
+ for _, (unit, other_unit) in zip_mappings(units, other_units)
) and (strip_units(ds).equals(strip_units(convert_units(other, units))))
equal_units = units == other_units
expected = equal_ds and (func.name != "identical" or equal_units)
@@ -4343,6 +4387,9 @@ def is_compatible(a, b):
assert expected == actual
+ # TODO: eventually use another decorator / wrapper function that
+ # applies a filter to the parametrize combinations:
+ # we only need a single test for data
@pytest.mark.parametrize(
"unit",
(
@@ -4353,14 +4400,29 @@ 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"),
+ ),
+ ),
+ )
+ def test_broadcast_like(self, variant, unit, dtype):
+ variants = {
+ "data": ((unit_registry.m, unit), (1, 1)),
+ "dims": ((1, 1), (unit_registry.m, unit)),
+ }
+ (data_unit1, data_unit2), (dim_unit1, dim_unit2) = variants.get(variant)
- x1 = np.arange(2) * unit_registry.m
- x2 = np.arange(2) * unit
- y1 = np.array([0]) * unit_registry.m
- y2 = np.arange(3) * unit
+ 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
ds1 = xr.Dataset(
data_vars={"a": (("x", "y"), array1)}, coords={"x": x1, "y": y1}
@@ -4374,7 +4436,8 @@ def test_broadcast_like(self, unit, dtype):
)
actual = ds1.broadcast_like(ds2)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"unit",
@@ -4387,32 +4450,25 @@ def test_broadcast_like(self, unit, dtype):
),
)
def test_broadcast_equals(self, unit, dtype):
+ # TODO: does this use indexes?
left_array1 = np.ones(shape=(2, 3), dtype=dtype) * unit_registry.m
left_array2 = np.zeros(shape=(3, 6), dtype=dtype) * unit_registry.m
right_array1 = np.ones(shape=(2,)) * unit
- right_array2 = np.ones(shape=(3,)) * unit
+ right_array2 = np.zeros(shape=(3,)) * unit
left = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=left_array1, dims=("x", "y")),
- "b": xr.DataArray(data=left_array2, dims=("y", "z")),
- }
- )
- right = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=right_array1, dims="x"),
- "b": xr.DataArray(data=right_array2, dims="y"),
- }
+ {"a": (("x", "y"), left_array1), "b": (("y", "z"), left_array2)},
)
+ right = xr.Dataset({"a": ("x", right_array1), "b": ("y", right_array2)})
- units = {
- **extract_units(left),
- **({} if left_array1.check(unit) else {"a": None, "b": None}),
- }
- expected = strip_units(left).broadcast_equals(
- strip_units(convert_units(right, units))
- ) & left_array1.check(unit)
+ units = merge_mappings(
+ extract_units(left),
+ {} if is_compatible(left_array1, unit) else {"a": None, "b": None},
+ )
+ expected = is_compatible(left_array1, unit) and strip_units(
+ left
+ ).broadcast_equals(strip_units(convert_units(right, units)))
actual = left.broadcast_equals(right)
assert expected == actual
@@ -4422,68 +4478,74 @@ def test_broadcast_equals(self, unit, dtype):
(method("unstack"), method("reset_index", "v"), method("reorder_levels")),
ids=repr,
)
- def test_stacking_stacked(self, func, dtype):
- array1 = (
- np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * unit_registry.m
- )
+ @pytest.mark.parametrize(
+ "variant",
+ (
+ "data",
+ pytest.param(
+ "dims", marks=pytest.mark.xfail(reason="indexes don't support units"),
+ ),
+ ),
+ )
+ def test_stacking_stacked(self, variant, func, dtype):
+ variants = {
+ "data": (unit_registry.m, 1),
+ "dims": (1, unit_registry.m),
+ }
+ data_unit, dim_unit = variants.get(variant)
+
+ array1 = np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * data_unit
array2 = (
np.linspace(-10, 0, 5 * 10 * 15).reshape(5, 10, 15).astype(dtype)
- * unit_registry.m
+ * data_unit
)
- x = np.arange(array1.shape[0])
- y = np.arange(array1.shape[1])
- z = np.arange(array2.shape[2])
+ x = np.arange(array1.shape[0]) * dim_unit
+ y = np.arange(array1.shape[1]) * dim_unit
+ z = np.arange(array2.shape[2]) * dim_unit
ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims=("x", "y")),
- "b": xr.DataArray(data=array2, dims=("x", "y", "z")),
- },
+ data_vars={"a": (("x", "y"), array1), "b": (("x", "y", "z"), array2)},
coords={"x": x, "y": y, "z": z},
)
+ units = extract_units(ds)
stacked = ds.stack(v=("x", "y"))
- expected = attach_units(
- func(strip_units(stacked)), {"a": unit_registry.m, "b": unit_registry.m}
- )
+ expected = attach_units(func(strip_units(stacked)), units)
actual = func(stacked)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
- @pytest.mark.xfail(reason="does not work with quantities yet")
+ @pytest.mark.xfail(
+ reason="stacked dimension's labels have to be hashable, but is a numpy.array"
+ )
def test_to_stacked_array(self, dtype):
- labels = np.arange(5).astype(dtype) * unit_registry.s
- arrays = {name: np.linspace(0, 1, 10) * unit_registry.m for name in labels}
+ labels = range(5) * unit_registry.s
+ arrays = {
+ name: np.linspace(0, 1, 10).astype(dtype) * unit_registry.m
+ for name in labels
+ }
- ds = xr.Dataset(
- data_vars={
- name: xr.DataArray(data=array, dims="x")
- for name, array in arrays.items()
- }
- )
+ ds = xr.Dataset({name: ("x", array) for name, array in arrays.items()})
+ units = {None: unit_registry.m, "y": unit_registry.s}
func = method("to_stacked_array", "z", variable_dim="y", sample_dims=["x"])
actual = func(ds).rename(None)
- expected = attach_units(
- func(strip_units(ds)).rename(None),
- {None: unit_registry.m, "y": unit_registry.s},
- )
+ expected = attach_units(func(strip_units(ds)).rename(None), units,)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"func",
(
method("transpose", "y", "x", "z1", "z2"),
- method("stack", a=("x", "y")),
+ method("stack", u=("x", "y")),
method("set_index", x="x2"),
- pytest.param(
- method("shift", x=2),
- marks=pytest.mark.xfail(reason="tries to concatenate nan arrays"),
- ),
+ method("shift", x=2),
method("roll", x=2, roll_coords=False),
method("sortby", "x2"),
),
@@ -4508,20 +4570,19 @@ def test_stacking_reordering(self, func, dtype):
ds = xr.Dataset(
data_vars={
- "a": xr.DataArray(data=array1, dims=("x", "y", "z1")),
- "b": xr.DataArray(data=array2, dims=("x", "y", "z2")),
+ "a": (("x", "y", "z1"), array1),
+ "b": (("x", "y", "z2"), array2),
},
coords={"x": x, "y": y, "z1": z1, "z2": z2, "x2": ("x", x2)},
)
+ units = extract_units(ds)
- expected = attach_units(
- func(strip_units(ds)), {"a": unit_registry.Pa, "b": unit_registry.degK}
- )
+ expected = attach_units(func(strip_units(ds)), units)
actual = func(ds)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
- @pytest.mark.xfail(reason="indexes strip units")
@pytest.mark.parametrize(
"indices",
(
@@ -4533,22 +4594,14 @@ def test_isel(self, indices, dtype):
array1 = np.arange(10).astype(dtype) * unit_registry.s
array2 = np.linspace(0, 1, 10).astype(dtype) * unit_registry.Pa
- x = np.arange(len(array1)) * unit_registry.m
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims="x"),
- "b": xr.DataArray(data=array2, dims="x"),
- },
- coords={"x": x},
- )
+ ds = xr.Dataset(data_vars={"a": ("x", array1), "b": ("x", array2)})
+ units = extract_units(ds)
- expected = attach_units(
- strip_units(ds).isel(x=indices),
- {"a": unit_registry.s, "b": unit_registry.Pa, "x": unit_registry.m},
- )
+ expected = attach_units(strip_units(ds).isel(x=indices), units)
actual = ds.isel(x=indices)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.xfail(reason="indexes don't support units")
@pytest.mark.parametrize(
@@ -4565,7 +4618,7 @@ def test_isel(self, indices, dtype):
pytest.param(1, KeyError, id="no_units"),
pytest.param(unit_registry.dimensionless, KeyError, id="dimensionless"),
pytest.param(unit_registry.degree, KeyError, id="incompatible_unit"),
- pytest.param(unit_registry.dm, KeyError, id="compatible_unit"),
+ pytest.param(unit_registry.mm, KeyError, id="compatible_unit"),
pytest.param(unit_registry.m, None, id="identical_unit"),
),
)
@@ -4584,20 +4637,24 @@ def test_sel(self, raw_values, unit, error, dtype):
values = raw_values * unit
- if error is not None and not (
- isinstance(raw_values, (int, float)) and x.check(unit)
- ):
+ # TODO: if we choose dm as compatible unit, single value keys
+ # can be found. Should we check that?
+ if error is not None:
with pytest.raises(error):
ds.sel(x=values)
return
expected = attach_units(
- strip_units(ds).sel(x=strip_units(convert_units(values, {None: x.units}))),
- {"a": array1.units, "b": array2.units, "x": x.units},
+ strip_units(ds).sel(
+ x=strip_units(convert_units(values, {None: unit_registry.m}))
+ ),
+ extract_units(ds),
)
actual = ds.sel(x=values)
- assert_equal_with_units(expected, actual)
+
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.xfail(reason="indexes don't support units")
@pytest.mark.parametrize(
@@ -4614,7 +4671,7 @@ def test_sel(self, raw_values, unit, error, dtype):
pytest.param(1, KeyError, id="no_units"),
pytest.param(unit_registry.dimensionless, KeyError, id="dimensionless"),
pytest.param(unit_registry.degree, KeyError, id="incompatible_unit"),
- pytest.param(unit_registry.dm, KeyError, id="compatible_unit"),
+ pytest.param(unit_registry.mm, KeyError, id="compatible_unit"),
pytest.param(unit_registry.m, None, id="identical_unit"),
),
)
@@ -4633,9 +4690,9 @@ def test_drop_sel(self, raw_values, unit, error, dtype):
values = raw_values * unit
- if error is not None and not (
- isinstance(raw_values, (int, float)) and x.check(unit)
- ):
+ # TODO: if we choose dm as compatible unit, single value keys
+ # can be found. Should we check that?
+ if error is not None:
with pytest.raises(error):
ds.drop_sel(x=values)
@@ -4643,12 +4700,14 @@ def test_drop_sel(self, raw_values, unit, error, dtype):
expected = attach_units(
strip_units(ds).drop_sel(
- x=strip_units(convert_units(values, {None: x.units}))
+ x=strip_units(convert_units(values, {None: unit_registry.m}))
),
extract_units(ds),
)
actual = ds.drop_sel(x=values)
- assert_equal_with_units(expected, actual)
+
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.xfail(reason="indexes don't support units")
@pytest.mark.parametrize(
@@ -4665,7 +4724,7 @@ def test_drop_sel(self, raw_values, unit, error, dtype):
pytest.param(1, KeyError, id="no_units"),
pytest.param(unit_registry.dimensionless, KeyError, id="dimensionless"),
pytest.param(unit_registry.degree, KeyError, id="incompatible_unit"),
- pytest.param(unit_registry.dm, KeyError, id="compatible_unit"),
+ pytest.param(unit_registry.mm, KeyError, id="compatible_unit"),
pytest.param(unit_registry.m, None, id="identical_unit"),
),
)
@@ -4684,9 +4743,9 @@ def test_loc(self, raw_values, unit, error, dtype):
values = raw_values * unit
- if error is not None and not (
- isinstance(raw_values, (int, float)) and x.check(unit)
- ):
+ # TODO: if we choose dm as compatible unit, single value keys
+ # can be found. Should we check that?
+ if error is not None:
with pytest.raises(error):
ds.loc[{"x": values}]
@@ -4694,12 +4753,14 @@ def test_loc(self, raw_values, unit, error, dtype):
expected = attach_units(
strip_units(ds).loc[
- {"x": strip_units(convert_units(values, {None: x.units}))}
+ {"x": strip_units(convert_units(values, {None: unit_registry.m}))}
],
- {"a": array1.units, "b": array2.units, "x": x.units},
+ extract_units(ds),
)
actual = ds.loc[{"x": values}]
- assert_equal_with_units(expected, actual)
+
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"func",
@@ -4710,14 +4771,34 @@ def test_loc(self, raw_values, unit, error, dtype):
),
ids=repr,
)
- def test_head_tail_thin(self, func, dtype):
- array1 = np.linspace(1, 2, 10 * 5).reshape(10, 5) * unit_registry.degK
- array2 = np.linspace(1, 2, 10 * 8).reshape(10, 8) * unit_registry.Pa
+ @pytest.mark.parametrize(
+ "variant",
+ (
+ "data",
+ pytest.param(
+ "dims", marks=pytest.mark.xfail(reason="indexes don't support units")
+ ),
+ "coords",
+ ),
+ )
+ def test_head_tail_thin(self, func, variant, dtype):
+ variants = {
+ "data": ((unit_registry.degK, unit_registry.Pa), 1, 1),
+ "dims": ((1, 1), unit_registry.m, 1),
+ "coords": ((1, 1), 1, unit_registry.m),
+ }
+ (unit_a, unit_b), dim_unit, coord_unit = variants.get(variant)
+
+ array1 = np.linspace(1, 2, 10 * 5).reshape(10, 5) * unit_a
+ array2 = np.linspace(1, 2, 10 * 8).reshape(10, 8) * unit_b
coords = {
- "x": np.arange(10) * unit_registry.m,
- "y": np.arange(5) * unit_registry.m,
- "z": np.arange(8) * unit_registry.m,
+ "x": np.arange(10) * dim_unit,
+ "y": np.arange(5) * dim_unit,
+ "z": np.arange(8) * dim_unit,
+ "u": ("x", np.linspace(0, 1, 10) * coord_unit),
+ "v": ("y", np.linspace(1, 2, 5) * coord_unit),
+ "w": ("z", np.linspace(-1, 0, 8) * coord_unit),
}
ds = xr.Dataset(
@@ -4731,8 +4812,10 @@ def test_head_tail_thin(self, func, dtype):
expected = attach_units(func(strip_units(ds)), extract_units(ds))
actual = func(ds)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
+ @pytest.mark.parametrize("dim", ("x", "y", "z", "t", "all"))
@pytest.mark.parametrize(
"shape",
(
@@ -4743,13 +4826,9 @@ def test_head_tail_thin(self, func, 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"
- coords = {
- name: np.arange(length).astype(dtype)
- * (unit_registry.m if name != "t" else unit_registry.s)
- for name, length in zip(names, shape)
- }
+ dim_lengths = dict(zip(names, shape))
array1 = (
np.linspace(0, 1, 10 * 20).astype(dtype).reshape(shape) * unit_registry.degK
)
@@ -4759,74 +4838,59 @@ def test_squeeze(self, shape, dtype):
ds = xr.Dataset(
data_vars={
- "a": xr.DataArray(data=array1, dims=tuple(names[: len(shape)])),
- "b": xr.DataArray(data=array2, dims=tuple(names[: len(shape)])),
+ "a": (tuple(names[: len(shape)]), array1),
+ "b": (tuple(names[: len(shape)]), array2),
},
- coords=coords,
)
units = extract_units(ds)
- expected = attach_units(strip_units(ds).squeeze(), units)
+ kwargs = {"dim": dim} if dim != "all" and dim_lengths.get(dim, 0) == 1 else {}
- actual = ds.squeeze()
- assert_equal_with_units(actual, expected)
+ expected = attach_units(strip_units(ds).squeeze(**kwargs), units)
- # try squeezing the dimensions separately
- names = tuple(dim for dim, coord in coords.items() if len(coord) == 1)
- for name in names:
- expected = attach_units(strip_units(ds).squeeze(dim=name), units)
- actual = ds.squeeze(dim=name)
- assert_equal_with_units(actual, expected)
+ actual = ds.squeeze(**kwargs)
+
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
- @pytest.mark.xfail(reason="ignores units")
+ @pytest.mark.parametrize("variant", ("data", "coords"))
@pytest.mark.parametrize(
- "unit,error",
+ "func",
(
- pytest.param(1, DimensionalityError, id="no_unit"),
pytest.param(
- unit_registry.dimensionless, DimensionalityError, id="dimensionless"
+ method("interp"), marks=pytest.mark.xfail(reason="uses scipy")
),
- pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
- pytest.param(unit_registry.cm, None, id="compatible_unit"),
- pytest.param(unit_registry.m, None, id="identical_unit"),
+ method("reindex"),
),
+ ids=repr,
)
- def test_interp(self, unit, error):
- array1 = np.linspace(1, 2, 10 * 5).reshape(10, 5) * unit_registry.degK
- array2 = np.linspace(1, 2, 10 * 8).reshape(10, 8) * unit_registry.Pa
-
- coords = {
- "x": np.arange(10) * unit_registry.m,
- "y": np.arange(5) * unit_registry.m,
- "z": np.arange(8) * unit_registry.s,
+ def test_interp_reindex(self, func, variant, dtype):
+ variants = {
+ "data": (unit_registry.m, 1),
+ "coords": (1, unit_registry.m),
}
+ data_unit, coord_unit = variants.get(variant)
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims=("x", "y")),
- "b": xr.DataArray(data=array2, dims=("x", "z")),
- },
- coords=coords,
- )
+ array1 = np.linspace(-1, 0, 10).astype(dtype) * data_unit
+ array2 = np.linspace(0, 1, 10).astype(dtype) * data_unit
- new_coords = (np.arange(10) + 0.5) * unit
+ y = np.arange(10) * coord_unit
- if error is not None:
- with pytest.raises(error):
- ds.interp(x=new_coords)
-
- return
+ x = np.arange(10)
+ new_x = np.arange(8) + 0.5
- units = extract_units(ds)
- expected = attach_units(
- strip_units(ds).interp(x=strip_units(convert_units(new_coords, units))),
- units,
+ ds = xr.Dataset(
+ {"a": ("x", array1), "b": ("x", array2)}, coords={"x": x, "y": ("x", y)}
)
- actual = ds.interp(x=new_coords)
+ units = extract_units(ds)
+
+ expected = attach_units(func(strip_units(ds), x=new_x), units)
+ actual = func(ds, x=new_x)
- assert_equal_with_units(actual, expected)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
- @pytest.mark.xfail(reason="ignores units")
+ @pytest.mark.xfail(reason="indexes don't support units")
@pytest.mark.parametrize(
"unit,error",
(
@@ -4839,106 +4903,67 @@ def test_interp(self, unit, error):
pytest.param(unit_registry.m, None, id="identical_unit"),
),
)
- def test_interp_like(self, unit, error, dtype):
- array1 = (
- np.linspace(0, 10, 10 * 5).reshape(10, 5).astype(dtype) * unit_registry.degK
- )
- array2 = (
- np.linspace(10, 20, 10 * 8).reshape(10, 8).astype(dtype) * unit_registry.Pa
- )
+ @pytest.mark.parametrize("func", (method("interp"), method("reindex")), ids=repr)
+ def test_interp_reindex_indexing(self, func, unit, error, dtype):
+ array1 = np.linspace(-1, 0, 10).astype(dtype)
+ array2 = np.linspace(0, 1, 10).astype(dtype)
- coords = {
- "x": np.arange(10) * unit_registry.m,
- "y": np.arange(5) * unit_registry.m,
- "z": np.arange(8) * unit_registry.m,
- }
-
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims=("x", "y")),
- "b": xr.DataArray(data=array2, dims=("x", "z")),
- },
- coords=coords,
- )
+ x = np.arange(10) * unit_registry.m
+ new_x = (np.arange(8) + 0.5) * unit
- other = xr.Dataset(
- data_vars={
- "c": xr.DataArray(data=np.empty((20, 10)), dims=("x", "y")),
- "d": xr.DataArray(data=np.empty((20, 15)), dims=("x", "z")),
- },
- coords={
- "x": (np.arange(20) + 0.3) * unit,
- "y": (np.arange(10) - 0.2) * unit,
- "z": (np.arange(15) + 0.4) * unit,
- },
- )
+ ds = xr.Dataset({"a": ("x", array1), "b": ("x", array2)}, coords={"x": x})
+ units = extract_units(ds)
if error is not None:
with pytest.raises(error):
- ds.interp_like(other)
+ func(ds, x=new_x)
return
- units = extract_units(ds)
- expected = attach_units(
- strip_units(ds).interp_like(strip_units(convert_units(other, units))), units
- )
- actual = ds.interp_like(other)
+ expected = attach_units(func(strip_units(ds), x=new_x), units)
+ actual = func(ds, x=new_x)
- assert_equal_with_units(actual, expected)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
- @pytest.mark.xfail(reason="indexes don't support units")
+ @pytest.mark.parametrize("variant", ("data", "coords"))
@pytest.mark.parametrize(
- "unit,error",
+ "func",
(
- pytest.param(1, DimensionalityError, id="no_unit"),
pytest.param(
- unit_registry.dimensionless, DimensionalityError, id="dimensionless"
+ method("interp_like"), marks=pytest.mark.xfail(reason="uses scipy")
),
- pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"),
- pytest.param(unit_registry.cm, None, id="compatible_unit"),
- pytest.param(unit_registry.m, None, id="identical_unit"),
+ method("reindex_like"),
),
+ ids=repr,
)
- def test_reindex(self, unit, error, dtype):
- array1 = (
- np.linspace(1, 2, 10 * 5).reshape(10, 5).astype(dtype) * unit_registry.degK
- )
- array2 = (
- np.linspace(1, 2, 10 * 8).reshape(10, 8).astype(dtype) * unit_registry.Pa
- )
-
- coords = {
- "x": np.arange(10) * unit_registry.m,
- "y": np.arange(5) * unit_registry.m,
- "z": np.arange(8) * unit_registry.s,
+ def test_interp_reindex_like(self, func, variant, dtype):
+ variants = {
+ "data": (unit_registry.m, 1),
+ "coords": (1, unit_registry.m),
}
+ data_unit, coord_unit = variants.get(variant)
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims=("x", "y")),
- "b": xr.DataArray(data=array2, dims=("x", "z")),
- },
- coords=coords,
- )
+ array1 = np.linspace(-1, 0, 10).astype(dtype) * data_unit
+ array2 = np.linspace(0, 1, 10).astype(dtype) * data_unit
- new_coords = (np.arange(10) + 0.5) * unit
+ y = np.arange(10) * coord_unit
- if error is not None:
- with pytest.raises(error):
- ds.reindex(x=new_coords)
-
- return
+ x = np.arange(10)
+ new_x = np.arange(8) + 0.5
- expected = attach_units(
- strip_units(ds).reindex(
- x=strip_units(convert_units(new_coords, {None: coords["x"].units}))
- ),
- extract_units(ds),
+ ds = xr.Dataset(
+ {"a": ("x", array1), "b": ("x", array2)}, coords={"x": x, "y": ("x", y)}
)
- actual = ds.reindex(x=new_coords)
+ units = extract_units(ds)
+
+ other = xr.Dataset({"a": ("x", np.empty_like(new_x))}, coords={"x": new_x})
+
+ expected = attach_units(func(strip_units(ds), other), units)
+ actual = func(ds, other)
- assert_equal_with_units(actual, expected)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.xfail(reason="indexes don't support units")
@pytest.mark.parametrize(
@@ -4953,54 +4978,32 @@ def test_reindex(self, unit, error, dtype):
pytest.param(unit_registry.m, None, id="identical_unit"),
),
)
- def test_reindex_like(self, unit, error, dtype):
- array1 = (
- np.linspace(0, 10, 10 * 5).reshape(10, 5).astype(dtype) * unit_registry.degK
- )
- array2 = (
- np.linspace(10, 20, 10 * 8).reshape(10, 8).astype(dtype) * unit_registry.Pa
- )
+ @pytest.mark.parametrize(
+ "func", (method("interp_like"), method("reindex_like")), ids=repr
+ )
+ def test_interp_reindex_like_indexing(self, func, unit, error, dtype):
+ array1 = np.linspace(-1, 0, 10).astype(dtype)
+ array2 = np.linspace(0, 1, 10).astype(dtype)
- coords = {
- "x": np.arange(10) * unit_registry.m,
- "y": np.arange(5) * unit_registry.m,
- "z": np.arange(8) * unit_registry.m,
- }
+ x = np.arange(10) * unit_registry.m
+ new_x = (np.arange(8) + 0.5) * unit
- ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims=("x", "y")),
- "b": xr.DataArray(data=array2, dims=("x", "z")),
- },
- coords=coords,
- )
+ ds = xr.Dataset({"a": ("x", array1), "b": ("x", array2)}, coords={"x": x})
+ units = extract_units(ds)
- other = xr.Dataset(
- data_vars={
- "c": xr.DataArray(data=np.empty((20, 10)), dims=("x", "y")),
- "d": xr.DataArray(data=np.empty((20, 15)), dims=("x", "z")),
- },
- coords={
- "x": (np.arange(20) + 0.3) * unit,
- "y": (np.arange(10) - 0.2) * unit,
- "z": (np.arange(15) + 0.4) * unit,
- },
- )
+ other = xr.Dataset({"a": ("x", np.empty_like(new_x))}, coords={"x": new_x})
if error is not None:
with pytest.raises(error):
- ds.reindex_like(other)
+ func(ds, other)
return
- units = extract_units(ds)
- expected = attach_units(
- strip_units(ds).reindex_like(strip_units(convert_units(other, units))),
- units,
- )
- actual = ds.reindex_like(other)
+ expected = attach_units(func(strip_units(ds), other), units)
+ actual = func(ds, other)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"func",
@@ -5010,30 +5013,46 @@ def test_reindex_like(self, unit, error, dtype):
method("integrate", coord="x"),
pytest.param(
method("quantile", q=[0.25, 0.75]),
- marks=pytest.mark.xfail(reason="nanquantile not implemented"),
+ marks=pytest.mark.xfail(
+ LooseVersion(pint.__version__) <= "0.12",
+ reason="nanquantile not implemented yet",
+ ),
),
method("reduce", func=np.sum, dim="x"),
method("map", np.fabs),
),
ids=repr,
)
- def test_computation(self, func, dtype):
- array1 = (
- np.linspace(-5, 5, 10 * 5).reshape(10, 5).astype(dtype) * unit_registry.degK
- )
- array2 = (
- np.linspace(10, 20, 10 * 8).reshape(10, 8).astype(dtype) * unit_registry.Pa
- )
- x = np.arange(10) * unit_registry.m
- y = np.arange(5) * unit_registry.m
- z = np.arange(8) * unit_registry.m
+ @pytest.mark.parametrize(
+ "variant",
+ (
+ "data",
+ pytest.param(
+ "dims", marks=pytest.mark.xfail(reason="indexes don't support units")
+ ),
+ "coords",
+ ),
+ )
+ def test_computation(self, func, variant, dtype):
+ variants = {
+ "data": ((unit_registry.degK, unit_registry.Pa), 1, 1),
+ "dims": ((1, 1), unit_registry.m, 1),
+ "coords": ((1, 1), 1, unit_registry.m),
+ }
+ (unit1, unit2), dim_unit, coord_unit = variants.get(variant)
+
+ array1 = np.linspace(-5, 5, 4 * 5).reshape(4, 5).astype(dtype) * unit1
+ array2 = np.linspace(10, 20, 4 * 3).reshape(4, 3).astype(dtype) * unit2
+ x = np.arange(4) * dim_unit
+ y = np.arange(5) * dim_unit
+ z = np.arange(3) * dim_unit
ds = xr.Dataset(
data_vars={
"a": xr.DataArray(data=array1, dims=("x", "y")),
"b": xr.DataArray(data=array2, dims=("x", "z")),
},
- coords={"x": x, "y": y, "z": z},
+ coords={"x": x, "y": y, "z": z, "y2": ("y", np.arange(5) * coord_unit)},
)
units = extract_units(ds)
@@ -5041,69 +5060,105 @@ def test_computation(self, func, dtype):
expected = attach_units(func(strip_units(ds)), units)
actual = func(ds)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"func",
(
method("groupby", "x"),
- method("groupby_bins", "x", bins=4),
+ 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("coarsen", x=2),
pytest.param(
method("rolling", x=3), marks=pytest.mark.xfail(reason="strips units")
),
pytest.param(
method("rolling_exp", x=3),
- marks=pytest.mark.xfail(reason="uses numbagg which strips units"),
+ marks=pytest.mark.xfail(
+ reason="numbagg functions are not supported by pint"
+ ),
),
),
ids=repr,
)
- def test_computation_objects(self, func, dtype):
- array1 = (
- np.linspace(-5, 5, 10 * 5).reshape(10, 5).astype(dtype) * unit_registry.degK
- )
- array2 = (
- np.linspace(10, 20, 10 * 5 * 8).reshape(10, 5, 8).astype(dtype)
- * unit_registry.Pa
- )
- x = np.arange(10) * unit_registry.m
- y = np.arange(5) * unit_registry.m
- z = np.arange(8) * unit_registry.m
+ @pytest.mark.parametrize(
+ "variant",
+ (
+ "data",
+ pytest.param(
+ "dims", marks=pytest.mark.xfail(reason="indexes don't support units")
+ ),
+ "coords",
+ ),
+ )
+ def test_computation_objects(self, func, variant, dtype):
+ variants = {
+ "data": ((unit_registry.degK, unit_registry.Pa), 1, 1),
+ "dims": ((1, 1), unit_registry.m, 1),
+ "coords": ((1, 1), 1, unit_registry.m),
+ }
+ (unit1, unit2), dim_unit, coord_unit = variants.get(variant)
+
+ array1 = np.linspace(-5, 5, 4 * 5).reshape(4, 5).astype(dtype) * unit1
+ array2 = np.linspace(10, 20, 4 * 3).reshape(4, 3).astype(dtype) * unit2
+ x = np.arange(4) * dim_unit
+ y = np.arange(5) * dim_unit
+ z = np.arange(3) * dim_unit
ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims=("x", "y")),
- "b": xr.DataArray(data=array2, dims=("x", "y", "z")),
- },
- coords={"x": x, "y": y, "z": z},
+ data_vars={"a": (("x", "y"), array1), "b": (("x", "z"), array2)},
+ coords={"x": x, "y": y, "z": z, "y2": ("y", np.arange(5) * coord_unit)},
)
units = extract_units(ds)
args = [] if func.name != "groupby" else ["y"]
- reduce_func = method("mean", *args)
- expected = attach_units(reduce_func(func(strip_units(ds))), units)
- actual = reduce_func(func(ds))
+ expected = attach_units(func(strip_units(ds)).mean(*args), units)
+ actual = func(ds).mean(*args)
- assert_equal_with_units(expected, actual)
+ 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)
+
+ @pytest.mark.parametrize(
+ "variant",
+ (
+ "data",
+ pytest.param(
+ "dims", marks=pytest.mark.xfail(reason="indexes don't support units")
+ ),
+ "coords",
+ ),
+ )
+ def test_resample(self, variant, dtype):
+ # TODO: move this to test_computation_objects
+ variants = {
+ "data": ((unit_registry.degK, unit_registry.Pa), 1, 1),
+ "dims": ((1, 1), unit_registry.m, 1),
+ "coords": ((1, 1), 1, unit_registry.m),
+ }
+ (unit1, unit2), dim_unit, coord_unit = variants.get(variant)
+
+ array1 = np.linspace(-5, 5, 10 * 5).reshape(10, 5).astype(dtype) * unit1
+ array2 = np.linspace(10, 20, 10 * 8).reshape(10, 8).astype(dtype) * unit2
- def test_resample(self, dtype):
- array1 = (
- np.linspace(-5, 5, 10 * 5).reshape(10, 5).astype(dtype) * unit_registry.degK
- )
- array2 = (
- np.linspace(10, 20, 10 * 8).reshape(10, 8).astype(dtype) * unit_registry.Pa
- )
t = pd.date_range("10-09-2010", periods=array1.shape[0], freq="1y")
- y = np.arange(5) * unit_registry.m
- z = np.arange(8) * unit_registry.m
+ y = np.arange(5) * dim_unit
+ z = np.arange(8) * dim_unit
+
+ u = np.linspace(-1, 0, 5) * coord_unit
ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims=("time", "y")),
- "b": xr.DataArray(data=array2, dims=("time", "z")),
- },
- coords={"time": t, "y": y, "z": z},
+ data_vars={"a": (("time", "y"), array1), "b": (("time", "z"), array2)},
+ coords={"time": t, "y": y, "z": z, "u": ("y", u)},
)
units = extract_units(ds)
@@ -5112,43 +5167,59 @@ def test_resample(self, dtype):
expected = attach_units(func(strip_units(ds)).mean(), units)
actual = func(ds).mean()
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"func",
(
method("assign", c=lambda ds: 10 * ds.b),
- method("assign_coords", v=("x", np.arange(10) * unit_registry.s)),
+ 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(reason="nanquantile not implemented"),
+ marks=pytest.mark.xfail(
+ LooseVersion(pint.__version__) <= "0.12",
+ reason="nanquantile not implemented",
+ ),
),
),
ids=repr,
)
- def test_grouped_operations(self, func, dtype):
- array1 = (
- np.linspace(-5, 5, 10 * 5).reshape(10, 5).astype(dtype) * unit_registry.degK
- )
- array2 = (
- np.linspace(10, 20, 10 * 5 * 8).reshape(10, 5, 8).astype(dtype)
- * unit_registry.Pa
- )
- x = np.arange(10) * unit_registry.m
- y = np.arange(5) * unit_registry.m
- z = np.arange(8) * unit_registry.m
+ @pytest.mark.parametrize(
+ "variant",
+ (
+ "data",
+ pytest.param(
+ "dims", marks=pytest.mark.xfail(reason="indexes don't support units")
+ ),
+ "coords",
+ ),
+ )
+ def test_grouped_operations(self, func, variant, dtype):
+ variants = {
+ "data": ((unit_registry.degK, unit_registry.Pa), 1, 1),
+ "dims": ((1, 1), unit_registry.m, 1),
+ "coords": ((1, 1), 1, unit_registry.m),
+ }
+ (unit1, unit2), dim_unit, coord_unit = variants.get(variant)
+
+ array1 = np.linspace(-5, 5, 5 * 4).reshape(5, 4).astype(dtype) * unit1
+ array2 = np.linspace(10, 20, 5 * 4 * 3).reshape(5, 4, 3).astype(dtype) * unit2
+ x = np.arange(5) * dim_unit
+ y = np.arange(4) * dim_unit
+ z = np.arange(3) * dim_unit
+
+ u = np.linspace(-1, 0, 4) * coord_unit
ds = xr.Dataset(
- data_vars={
- "a": xr.DataArray(data=array1, dims=("x", "y")),
- "b": xr.DataArray(data=array2, dims=("x", "y", "z")),
- },
- coords={"x": x, "y": y, "z": z},
+ data_vars={"a": (("x", "y"), array1), "b": (("x", "y", "z"), array2)},
+ coords={"x": x, "y": y, "z": z, "u": ("y", u)},
)
- units = extract_units(ds)
- units.update({"c": unit_registry.Pa, "v": unit_registry.s})
+
+ assigned_units = {"c": unit2, "v": unit_registry.s}
+ units = merge_mappings(extract_units(ds), assigned_units)
stripped_kwargs = {
name: strip_units(value) for name, value in func.kwargs.items()
@@ -5158,20 +5229,26 @@ def test_grouped_operations(self, func, dtype):
)
actual = func(ds.groupby("y"))
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"func",
(
method("pipe", lambda ds: ds * 10),
method("assign", d=lambda ds: ds.b * 10),
- method("assign_coords", y2=("y", np.arange(5) * unit_registry.mm)),
+ method("assign_coords", y2=("y", np.arange(4) * unit_registry.mm)),
method("assign_attrs", attr1="value"),
method("rename", x2="x_mm"),
method("rename_vars", c="temperature"),
method("rename_dims", x="offset_x"),
- method("swap_dims", {"x": "x2"}),
- method("expand_dims", v=np.linspace(10, 20, 12) * unit_registry.s, axis=1),
+ method("swap_dims", {"x": "u"}),
+ pytest.param(
+ method(
+ "expand_dims", v=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("drop_dims", "z"),
method("set_coords", names="c"),
@@ -5180,40 +5257,55 @@ def test_grouped_operations(self, func, dtype):
),
ids=repr,
)
- def test_content_manipulation(self, func, dtype):
- array1 = (
- np.linspace(-5, 5, 10 * 5).reshape(10, 5).astype(dtype)
- * unit_registry.m ** 3
- )
- array2 = (
- np.linspace(10, 20, 10 * 5 * 8).reshape(10, 5, 8).astype(dtype)
- * unit_registry.Pa
- )
- array3 = np.linspace(0, 10, 10).astype(dtype) * unit_registry.degK
+ @pytest.mark.parametrize(
+ "variant",
+ (
+ "data",
+ pytest.param(
+ "dims", marks=pytest.mark.xfail(reason="indexes don't support units")
+ ),
+ "coords",
+ ),
+ )
+ def test_content_manipulation(self, func, variant, dtype):
+ variants = {
+ "data": (
+ (unit_registry.m ** 3, unit_registry.Pa, unit_registry.degK),
+ 1,
+ 1,
+ ),
+ "dims": ((1, 1, 1), unit_registry.m, 1),
+ "coords": ((1, 1, 1), 1, unit_registry.m),
+ }
+ (unit1, unit2, unit3), dim_unit, coord_unit = variants.get(variant)
- x = np.arange(10) * unit_registry.m
- x2 = x.to(unit_registry.mm)
- y = np.arange(5) * unit_registry.m
- z = np.arange(8) * unit_registry.m
+ array1 = np.linspace(-5, 5, 5 * 4).reshape(5, 4).astype(dtype) * unit1
+ array2 = np.linspace(10, 20, 5 * 4 * 3).reshape(5, 4, 3).astype(dtype) * unit2
+ array3 = np.linspace(0, 10, 5).astype(dtype) * unit3
+
+ x = np.arange(5) * dim_unit
+ y = np.arange(4) * dim_unit
+ z = np.arange(3) * dim_unit
+
+ x2 = np.linspace(-1, 0, 5) * coord_unit
ds = xr.Dataset(
data_vars={
- "a": xr.DataArray(data=array1, dims=("x", "y")),
- "b": xr.DataArray(data=array2, dims=("x", "y", "z")),
- "c": xr.DataArray(data=array3, dims="x"),
+ "a": (("x", "y"), array1),
+ "b": (("x", "y", "z"), array2),
+ "c": ("x", array3),
},
coords={"x": x, "y": y, "z": z, "x2": ("x", x2)},
)
- units = {
- **extract_units(ds),
- **{
- "y2": unit_registry.mm,
- "x_mm": unit_registry.mm,
- "offset_x": unit_registry.m,
- "d": unit_registry.Pa,
- "temperature": unit_registry.degK,
- },
+
+ new_units = {
+ "y2": unit_registry.mm,
+ "x_mm": coord_unit,
+ "offset_x": unit_registry.m,
+ "d": unit2,
+ "temperature": unit3,
}
+ units = merge_mappings(extract_units(ds), new_units)
stripped_kwargs = {
key: strip_units(value) for key, value in func.kwargs.items()
@@ -5221,7 +5313,8 @@ def test_content_manipulation(self, func, dtype):
expected = attach_units(func(strip_units(ds), **stripped_kwargs), units)
actual = func(ds)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
@pytest.mark.parametrize(
"unit,error",
@@ -5246,25 +5339,29 @@ def test_content_manipulation(self, func, dtype):
),
)
def test_merge(self, variant, unit, error, dtype):
- original_data_unit = unit_registry.m
- original_dim_unit = unit_registry.m
- original_coord_unit = unit_registry.m
+ left_variants = {
+ "data": (unit_registry.m, 1, 1),
+ "dims": (1, unit_registry.m, 1),
+ "coords": (1, 1, unit_registry.m),
+ }
- variants = {
- "data": (unit, original_dim_unit, original_coord_unit),
- "dims": (original_data_unit, unit, original_coord_unit),
- "coords": (original_data_unit, original_dim_unit, unit),
+ left_data_unit, left_dim_unit, left_coord_unit = left_variants.get(variant)
+
+ right_variants = {
+ "data": (unit, 1, 1),
+ "dims": (1, unit, 1),
+ "coords": (1, 1, unit),
}
- data_unit, dim_unit, coord_unit = variants.get(variant)
+ right_data_unit, right_dim_unit, right_coord_unit = right_variants.get(variant)
- left_array = np.arange(10).astype(dtype) * original_data_unit
- right_array = np.arange(-5, 5).astype(dtype) * data_unit
+ left_array = np.arange(10).astype(dtype) * left_data_unit
+ right_array = np.arange(-5, 5).astype(dtype) * right_data_unit
- left_dim = np.arange(10, 20) * original_dim_unit
- right_dim = np.arange(5, 15) * dim_unit
+ left_dim = np.arange(10, 20) * left_dim_unit
+ right_dim = np.arange(5, 15) * right_dim_unit
- left_coord = np.arange(-10, 0) * original_coord_unit
- right_coord = np.arange(-15, -5) * coord_unit
+ left_coord = np.arange(-10, 0) * left_coord_unit
+ right_coord = np.arange(-15, -5) * right_coord_unit
left = xr.Dataset(
data_vars={"a": ("x", left_array)},
@@ -5287,4 +5384,5 @@ def test_merge(self, variant, unit, error, dtype):
expected = attach_units(strip_units(left).merge(strip_units(converted)), units)
actual = left.merge(right)
- assert_equal_with_units(expected, actual)
+ assert_units_equal(expected, actual)
+ assert_equal(expected, actual)
diff --git a/xarray/tests/test_variable.py b/xarray/tests/test_variable.py
index 78e3848b8fb..d79d40d67c0 100644
--- a/xarray/tests/test_variable.py
+++ b/xarray/tests/test_variable.py
@@ -1657,7 +1657,7 @@ def test_reduce_funcs(self):
assert_identical(v.all(dim="x"), Variable([], False))
v = Variable("t", pd.date_range("2000-01-01", periods=3))
- assert v.argmax(skipna=True) == 2
+ assert v.argmax(skipna=True, dim="t") == 2
assert_identical(v.max(), Variable([], pd.Timestamp("2000-01-03")))
@@ -2213,6 +2213,10 @@ def test_full_like(self):
assert expect.dtype == bool
assert_identical(expect, full_like(orig, True, dtype=bool))
+ # raise error on non-scalar fill_value
+ with raises_regex(ValueError, "must be scalar"):
+ full_like(orig, [1.0, 2.0])
+
@requires_dask
def test_full_like_dask(self):
orig = Variable(
diff --git a/xarray/tests/test_weighted.py b/xarray/tests/test_weighted.py
index 24531215dfb..1bf685cc95d 100644
--- a/xarray/tests/test_weighted.py
+++ b/xarray/tests/test_weighted.py
@@ -59,6 +59,18 @@ def test_weighted_sum_of_weights_nan(weights, expected):
assert_equal(expected, result)
+def test_weighted_sum_of_weights_bool():
+ # https://github.com/pydata/xarray/issues/4074
+
+ da = DataArray([1, 2])
+ weights = DataArray([True, True])
+ result = da.weighted(weights).sum_of_weights()
+
+ expected = DataArray(2)
+
+ assert_equal(expected, result)
+
+
@pytest.mark.parametrize("da", ([1.0, 2], [1, np.nan], [np.nan, np.nan]))
@pytest.mark.parametrize("factor", [0, 1, 3.14])
@pytest.mark.parametrize("skipna", (True, False))
@@ -158,6 +170,17 @@ def test_weighted_mean_nan(weights, expected, skipna):
assert_equal(expected, result)
+def test_weighted_mean_bool():
+ # https://github.com/pydata/xarray/issues/4074
+ da = DataArray([1, 1])
+ weights = DataArray([True, True])
+ expected = DataArray(1)
+
+ result = da.weighted(weights).mean()
+
+ assert_equal(expected, result)
+
+
def expected_weighted(da, weights, dim, skipna, operation):
"""
Generate expected result using ``*`` and ``sum``. This is checked against
diff --git a/xarray/util/print_versions.py b/xarray/util/print_versions.py
index 32051bb6843..96983c83aab 100755
--- a/xarray/util/print_versions.py
+++ b/xarray/util/print_versions.py
@@ -129,7 +129,7 @@ def show_versions(file=sys.stdout):
("sphinx", lambda mod: mod.__version__),
]
- deps_blob = list()
+ deps_blob = []
for (modname, ver_f) in deps:
try:
if modname in sys.modules: