From 5c674e6212ed5a5daeb3f289f023d6abc48417ae Mon Sep 17 00:00:00 2001 From: keewis Date: Mon, 9 Dec 2019 11:25:51 +0100 Subject: [PATCH] Fix pint integration tests (#3600) * don't use the index to construct the condition array * use drop_vars to drop variables and coordinates * update the xfail reason for astype and item * remove the xfails due to np.result_type not being implemented * make array_extract_units a bit more robust * add the missing dataset apply_ufunc test * use None as the dict key for quantities * convert to variable to avoid the merge / assign_coords issue relevant issue: #3483 * update the align tests * fix the bugs introduced by converting to variables * update a few more tests * update the aggregation tests * update a few more tests * don't test rank which was deprecated in numpy 1.9 * update the DataArray.fillna tests * update most of the remaining DataArray tests * remove a debug assert * fix the broadcast_equals tests * update the indexing tests * update the tests depending on einsum * update the squeeze tests * fix a bug in the head / tail / thin test design * use dictionaries to index in the loc tests * update the expected unit of the computation tests * update the grouped operations tests * update the where tests * update most of the remaining dataset tests * create new tests for drop_sel * final batch of updated tests * rename result to actual to match the other test files * fix some more test bugs * update the xfail marks --- xarray/tests/test_units.py | 1998 ++++++++++++++++++------------------ 1 file changed, 1019 insertions(+), 979 deletions(-) diff --git a/xarray/tests/test_units.py b/xarray/tests/test_units.py index 0be6f8af464..f8a8a259c1f 100644 --- a/xarray/tests/test_units.py +++ b/xarray/tests/test_units.py @@ -29,9 +29,11 @@ def array_extract_units(obj): - raw = obj.data if hasattr(obj, "data") else obj + if isinstance(obj, (xr.Variable, xr.DataArray, xr.Dataset)): + obj = obj.data + try: - return raw.units + return obj.units except AttributeError: return None @@ -112,7 +114,7 @@ def extract_units(obj): units = {**vars_units, **coords_units} elif isinstance(obj, Quantity): - vars_units = {"": array_extract_units(obj)} + vars_units = {None: array_extract_units(obj)} units = {**vars_units} else: @@ -203,21 +205,25 @@ def attach_units(obj, units): def convert_units(obj, to): if isinstance(obj, xr.Dataset): data_vars = { - name: convert_units(array, to) for name, array in obj.data_vars.items() + name: convert_units(array.variable, {None: to.get(name)}) + for name, array in obj.data_vars.items() + } + coords = { + name: convert_units(array.variable, {None: to.get(name)}) + for name, array in obj.coords.items() } - coords = {name: convert_units(array, to) for name, array in obj.coords.items()} new_obj = xr.Dataset(data_vars=data_vars, coords=coords, attrs=obj.attrs) elif isinstance(obj, xr.DataArray): name = obj.name new_units = ( - to.get(name, None) or to.get("data", None) or to.get(None, None) or 1 + to.get(name, None) or to.get("data", None) or to.get(None, None) or None ) - data = convert_units(obj.data, {None: new_units}) + data = convert_units(obj.variable, {None: new_units}) coords = { - name: (array.dims, convert_units(array.data, to)) + name: (array.dims, convert_units(array.variable, {None: to.get(name)})) for name, array in obj.coords.items() if name != obj.name } @@ -225,6 +231,9 @@ def convert_units(obj, to): new_obj = xr.DataArray( name=name, data=data, coords=coords, attrs=obj.attrs, dims=obj.dims ) + elif isinstance(obj, xr.Variable): + new_data = convert_units(obj.data, to) + new_obj = obj.copy(data=new_data) elif isinstance(obj, unit_registry.Quantity): units = to.get(None) new_obj = obj.to(units) if units is not None else obj @@ -344,14 +353,34 @@ def test_apply_ufunc_dataarray(dtype): data_array = xr.DataArray(data=array, dims="x", coords={"x": x}) expected = attach_units(func(strip_units(data_array)), extract_units(data_array)) - result = func(data_array) + actual = func(data_array) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) -@pytest.mark.xfail( - reason="pint does not implement `np.result_type` and align strips units" -) +def test_apply_ufunc_dataset(dtype): + func = function( + xr.apply_ufunc, np.mean, input_core_dims=[["x"]], kwargs={"axis": -1} + ) + + array1 = np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * unit_registry.m + array2 = np.linspace(0, 10, 5).astype(dtype) * unit_registry.m + + x = np.arange(5) * unit_registry.s + y = np.arange(10) * unit_registry.m + + ds = xr.Dataset( + data_vars={"a": (("x", "y"), array1), "b": ("x", array2)}, + coords={"x": x, "y": y}, + ) + + expected = attach_units(func(strip_units(ds)), extract_units(ds)) + actual = func(ds) + + assert_equal_with_units(expected, actual) + + +@pytest.mark.xfail(reason="blocked by `reindex` / `where`") @pytest.mark.parametrize( "unit,error", ( @@ -378,9 +407,9 @@ def test_align_dataarray(fill_value, variant, unit, error, dtype): original_unit = unit_registry.m variants = { - "data": (unit, original_unit, original_unit), - "dims": (original_unit, unit, original_unit), - "coords": (original_unit, original_unit, unit), + "data": (unit, 1, 1), + "dims": (original_unit, unit, 1), + "coords": (original_unit, 1, unit), } data_unit, dim_unit, coord_unit = variants.get(variant) @@ -410,32 +439,27 @@ def test_align_dataarray(fill_value, variant, unit, error, dtype): stripped_kwargs = { key: strip_units( - convert_units(value, {None: original_unit}) - if isinstance(value, unit_registry.Quantity) - else value + convert_units(value, {None: original_unit if data_unit != 1 else None}) ) for key, value in func.kwargs.items() } - units = extract_units(data_array1) - # FIXME: should the expected_b have the same units as data_array1 - # or data_array2? - expected_a, expected_b = tuple( - attach_units(elem, units) - for elem in func( - strip_units(data_array1), - strip_units(convert_units(data_array2, units)), - **stripped_kwargs, - ) + units_a = extract_units(data_array1) + units_b = extract_units(data_array2) + expected_a, expected_b = func( + strip_units(data_array1), + strip_units(convert_units(data_array2, units_a)), + **stripped_kwargs, ) - result_a, result_b = func(data_array1, data_array2) + expected_a = attach_units(expected_a, units_a) + expected_b = convert_units(attach_units(expected_b, units_a), units_b) - assert_equal_with_units(expected_a, result_a) - assert_equal_with_units(expected_b, result_b) + actual_a, actual_b = func(data_array1, data_array2) + assert_equal_with_units(expected_a, actual_a) + assert_equal_with_units(expected_b, actual_b) -@pytest.mark.xfail( - reason="pint does not implement `np.result_type` and align strips units" -) + +@pytest.mark.xfail(reason="blocked by `reindex` / `where`") @pytest.mark.parametrize( "unit,error", ( @@ -461,11 +485,7 @@ def test_align_dataarray(fill_value, variant, unit, error, dtype): def test_align_dataset(fill_value, unit, variant, error, dtype): original_unit = unit_registry.m - variants = { - "data": (unit, original_unit, original_unit), - "dims": (original_unit, unit, original_unit), - "coords": (original_unit, original_unit, unit), - } + variants = {"data": (unit, 1, 1), "dims": (1, unit, 1), "coords": (1, 1, unit)} data_unit, dim_unit, coord_unit = variants.get(variant) array1 = np.linspace(0, 10, 2 * 5).reshape(2, 5).astype(dtype) * original_unit @@ -497,24 +517,22 @@ def test_align_dataset(fill_value, unit, variant, error, dtype): stripped_kwargs = { key: strip_units( - convert_units(value, {None: original_unit}) - if isinstance(value, unit_registry.Quantity) - else value + convert_units(value, {None: original_unit if data_unit != 1 else None}) ) for key, value in func.kwargs.items() } - units = extract_units(ds1) - # FIXME: should the expected_b have the same units as ds1 or ds2? - expected_a, expected_b = tuple( - attach_units(elem, units) - for elem in func( - strip_units(ds1), strip_units(convert_units(ds2, units)), **stripped_kwargs - ) + units_a = extract_units(ds1) + units_b = extract_units(ds2) + expected_a, expected_b = func( + strip_units(ds1), strip_units(convert_units(ds2, units_a)), **stripped_kwargs ) - result_a, result_b = func(ds1, ds2) + expected_a = attach_units(expected_a, units_a) + expected_b = convert_units(attach_units(expected_b, units_a), units_b) - assert_equal_with_units(expected_a, result_a) - assert_equal_with_units(expected_b, result_b) + actual_a, actual_b = func(ds1, ds2) + + assert_equal_with_units(expected_a, actual_a) + assert_equal_with_units(expected_b, actual_b) def test_broadcast_dataarray(dtype): @@ -528,10 +546,10 @@ def test_broadcast_dataarray(dtype): attach_units(elem, extract_units(a)) for elem in xr.broadcast(strip_units(a), strip_units(b)) ) - result_a, result_b = xr.broadcast(a, b) + actual_a, actual_b = xr.broadcast(a, b) - assert_equal_with_units(expected_a, result_a) - assert_equal_with_units(expected_b, result_b) + assert_equal_with_units(expected_a, actual_a) + assert_equal_with_units(expected_b, actual_b) def test_broadcast_dataset(dtype): @@ -543,12 +561,11 @@ def test_broadcast_dataset(dtype): (expected,) = tuple( attach_units(elem, extract_units(ds)) for elem in xr.broadcast(strip_units(ds)) ) - (result,) = xr.broadcast(ds) + (actual,) = xr.broadcast(ds) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) -@pytest.mark.xfail(reason="`combine_by_coords` strips units") @pytest.mark.parametrize( "unit,error", ( @@ -614,12 +631,11 @@ def test_combine_by_coords(variant, unit, error, dtype): ), units, ) - result = xr.combine_by_coords([ds, other]) + actual = xr.combine_by_coords([ds, other]) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) -@pytest.mark.xfail(reason="blocked by `where`") @pytest.mark.parametrize( "unit,error", ( @@ -628,7 +644,12 @@ def test_combine_by_coords(variant, unit, 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="wrong order of arguments to `where`"), + ), pytest.param(unit_registry.m, None, id="identical_unit"), ), ids=repr, @@ -714,12 +735,11 @@ def test_combine_nested(variant, unit, error, dtype): ), units, ) - result = func([[ds1, ds2], [ds3, ds4]]) + actual = func([[ds1, ds2], [ds3, ds4]]) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) -@pytest.mark.xfail(reason="`concat` strips units") @pytest.mark.parametrize( "unit,error", ( @@ -760,15 +780,18 @@ def test_concat_dataarray(variant, unit, error, dtype): return + units = extract_units(arr1) expected = attach_units( - xr.concat([strip_units(arr1), strip_units(arr2)], dim="x"), extract_units(arr1) + xr.concat( + [strip_units(arr1), strip_units(convert_units(arr2, units))], dim="x" + ), + units, ) - result = xr.concat([arr1, arr2], dim="x") + actual = xr.concat([arr1, arr2], dim="x") - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) -@pytest.mark.xfail(reason="`concat` strips units") @pytest.mark.parametrize( "unit,error", ( @@ -809,15 +832,17 @@ def test_concat_dataset(variant, unit, error, dtype): return + units = extract_units(ds1) expected = attach_units( - xr.concat([strip_units(ds1), strip_units(ds2)], dim="x"), extract_units(ds1) + xr.concat([strip_units(ds1), strip_units(convert_units(ds2, units))], dim="x"), + units, ) - result = xr.concat([ds1, ds2], dim="x") + actual = xr.concat([ds1, ds2], dim="x") - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) -@pytest.mark.xfail(reason="blocked by `where`") +@pytest.mark.xfail(reason="blocked by `reindex` / `where`") @pytest.mark.parametrize( "unit,error", ( @@ -902,12 +927,12 @@ def test_merge_dataarray(variant, unit, error, dtype): func([strip_units(arr1), convert_and_strip(arr2), convert_and_strip(arr3)]), units, ) - result = func([arr1, arr2, arr3]) + actual = func([arr1, arr2, arr3]) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) -@pytest.mark.xfail(reason="blocked by `where`") +@pytest.mark.xfail(reason="blocked by `reindex` / `where`") @pytest.mark.parametrize( "unit,error", ( @@ -985,9 +1010,9 @@ def test_merge_dataset(variant, unit, error, dtype): expected = attach_units( func([strip_units(ds1), convert_and_strip(ds2), convert_and_strip(ds3)]), units ) - result = func([ds1, ds2, ds3]) + actual = func([ds1, ds2, ds3]) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize("func", (xr.zeros_like, xr.ones_like)) @@ -997,9 +1022,9 @@ def test_replication_dataarray(func, dtype): numpy_func = getattr(np, func.__name__) expected = xr.DataArray(data=numpy_func(array), dims="x") - result = func(data_array) + actual = func(data_array) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize("func", (xr.zeros_like, xr.ones_like)) @@ -1019,9 +1044,9 @@ def test_replication_dataset(func, dtype): expected = ds.copy( data={name: numpy_func(array.data) for name, array in ds.data_vars.items()} ) - result = func(ds) + actual = func(ds) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.xfail( @@ -1051,11 +1076,16 @@ def test_replication_full_like_dataarray(unit, error, dtype): if error is not None: with pytest.raises(error): xr.full_like(data_array, fill_value=fill_value) - else: - result = xr.full_like(data_array, fill_value=fill_value) - expected = np.full_like(array, fill_value=fill_value) - assert_equal_with_units(expected, result) + return + + units = {**extract_units(data_array), **{None: unit if unit != 1 else None}} + expected = attach_units( + xr.full_like(strip_units(data_array), fill_value=strip_units(fill_value)), units + ) + actual = xr.full_like(data_array, fill_value=fill_value) + + assert_equal_with_units(expected, actual) @pytest.mark.xfail( @@ -1096,18 +1126,18 @@ def test_replication_full_like_dataset(unit, error, dtype): return - expected = ds.copy( - data={ - name: np.full_like(array, fill_value=fill_value) - for name, array in ds.data_vars.items() - } + units = { + **extract_units(ds), + **{name: unit if unit != 1 else None for name in ds.data_vars}, + } + expected = attach_units( + xr.full_like(strip_units(ds), fill_value=strip_units(fill_value)), units ) - result = xr.full_like(ds, fill_value=fill_value) + actual = xr.full_like(ds, fill_value=fill_value) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) -@pytest.mark.xfail(reason="`where` strips units") @pytest.mark.parametrize( "unit,error", ( @@ -1127,30 +1157,29 @@ def test_where_dataarray(fill_value, unit, error, dtype): x = xr.DataArray(data=array, dims="x") cond = x < 5 * unit_registry.m - # FIXME: this should work without wrapping in array() - fill_value = np.array(fill_value) * unit + fill_value = fill_value * unit - if error is not None: + if error is not None and not ( + np.isnan(fill_value) and not isinstance(fill_value, Quantity) + ): with pytest.raises(error): xr.where(cond, x, fill_value) return - fill_value_ = ( - fill_value.to(unit_registry.m) - if isinstance(fill_value, unit_registry.Quantity) - and fill_value.check(unit_registry.m) - else fill_value - ) expected = attach_units( - xr.where(cond, strip_units(x), strip_units(fill_value_)), extract_units(x) + xr.where( + cond, + strip_units(x), + strip_units(convert_units(fill_value, {None: unit_registry.m})), + ), + extract_units(x), ) - result = xr.where(cond, x, fill_value) + actual = xr.where(cond, x, fill_value) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) -@pytest.mark.xfail(reason="`where` strips units") @pytest.mark.parametrize( "unit,error", ( @@ -1171,31 +1200,30 @@ def test_where_dataset(fill_value, unit, error, dtype): x = np.arange(10) * unit_registry.s ds = xr.Dataset(data_vars={"a": ("x", array1), "b": ("x", array2)}, coords={"x": x}) - cond = ds.x < 5 * unit_registry.s - # FIXME: this should work without wrapping in array() - fill_value = np.array(fill_value) * unit + cond = x < 5 * unit_registry.s + fill_value = fill_value * unit - if error is not None: + if error is not None and not ( + np.isnan(fill_value) and not isinstance(fill_value, Quantity) + ): with pytest.raises(error): xr.where(cond, ds, fill_value) return - fill_value_ = ( - fill_value.to(unit_registry.m) - if isinstance(fill_value, unit_registry.Quantity) - and fill_value.check(unit_registry.m) - else fill_value - ) expected = attach_units( - xr.where(cond, strip_units(ds), strip_units(fill_value_)), extract_units(ds) + xr.where( + cond, + strip_units(ds), + strip_units(convert_units(fill_value, {None: unit_registry.m})), + ), + extract_units(ds), ) - result = xr.where(cond, ds, fill_value) + actual = xr.where(cond, ds, fill_value) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) -@pytest.mark.xfail(reason="pint does not implement `np.einsum`") def test_dot_dataarray(dtype): array1 = ( np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) @@ -1206,13 +1234,15 @@ def test_dot_dataarray(dtype): np.linspace(10, 20, 10 * 20).reshape(10, 20).astype(dtype) * unit_registry.s ) - arr1 = xr.DataArray(data=array1, dims=("x", "y")) - arr2 = xr.DataArray(data=array2, dims=("y", "z")) + data_array = xr.DataArray(data=array1, dims=("x", "y")) + other = xr.DataArray(data=array2, dims=("y", "z")) - expected = array1.dot(array2) - result = xr.dot(arr1, arr2) + expected = attach_units( + xr.dot(strip_units(data_array), strip_units(other)), {None: unit_registry.m} + ) + actual = xr.dot(data_array, other) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) class TestDataArray: @@ -1295,37 +1325,20 @@ def test_repr(self, func, variant, dtype): function("any"), marks=pytest.mark.xfail(reason="not implemented by pint yet"), ), - pytest.param( - function("argmax"), - marks=pytest.mark.xfail( - reason="comparison of quantity with ndarrays in nanops not implemented" - ), - ), - pytest.param( - function("argmin"), - marks=pytest.mark.xfail( - reason="comparison of quantity with ndarrays in nanops not implemented" - ), - ), + function("argmax"), + function("argmin"), function("max"), function("mean"), pytest.param( function("median"), - marks=pytest.mark.xfail( - reason="np.median on DataArray strips the units" - ), + marks=pytest.mark.xfail(reason="not implemented by xarray"), ), function("min"), pytest.param( function("prod"), marks=pytest.mark.xfail(reason="not implemented by pint yet"), ), - pytest.param( - function("sum"), - marks=pytest.mark.xfail( - reason="comparison of quantity with ndarrays in nanops not implemented" - ), - ), + function("sum"), function("std"), function("var"), function("cumsum"), @@ -1341,18 +1354,8 @@ def test_repr(self, func, variant, dtype): method("any"), marks=pytest.mark.xfail(reason="not implemented by pint yet"), ), - pytest.param( - method("argmax"), - marks=pytest.mark.xfail( - reason="comparison of quantities with ndarrays in nanops not implemented" - ), - ), - pytest.param( - method("argmin"), - marks=pytest.mark.xfail( - reason="comparison of quantities with ndarrays in nanops not implemented" - ), - ), + method("argmax"), + method("argmin"), method("max"), method("mean"), method("median"), @@ -1363,12 +1366,7 @@ def test_repr(self, func, variant, dtype): reason="comparison of quantity with ndarrays in nanops not implemented" ), ), - pytest.param( - method("sum"), - marks=pytest.mark.xfail( - reason="comparison of quantity with ndarrays in nanops not implemented" - ), - ), + method("sum"), method("std"), method("var"), method("cumsum"), @@ -1380,34 +1378,36 @@ def test_repr(self, func, variant, dtype): ids=repr, ) def test_aggregation(self, func, dtype): - array = np.arange(10).astype(dtype) * unit_registry.m - data_array = xr.DataArray(data=array) + array = np.arange(10).astype(dtype) * ( + unit_registry.m if func.name != "cumprod" else unit_registry.dimensionless + ) + data_array = xr.DataArray(data=array, dims="x") - expected = xr.DataArray(data=func(array)) - result = func(data_array) + # 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, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", ( pytest.param(operator.neg, id="negate"), pytest.param(abs, id="absolute"), - pytest.param( - np.round, - id="round", - marks=pytest.mark.xfail(reason="pint does not implement round"), - ), + pytest.param(np.round, id="round"), ), ) def test_unary_operations(self, func, dtype): array = np.arange(10).astype(dtype) * unit_registry.m data_array = xr.DataArray(data=array) - expected = xr.DataArray(data=func(array)) - result = func(data_array) + units = extract_units(func(array)) + expected = attach_units(func(strip_units(data_array)), units) + actual = func(data_array) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", @@ -1415,23 +1415,18 @@ def test_unary_operations(self, func, dtype): pytest.param(lambda x: 2 * x, id="multiply"), pytest.param(lambda x: x + x, id="add"), pytest.param(lambda x: x[0] + x, id="add scalar"), - pytest.param( - lambda x: x.T @ x, - id="matrix multiply", - marks=pytest.mark.xfail( - reason="pint does not support matrix multiplication yet" - ), - ), + pytest.param(lambda x: x.T @ x, id="matrix multiply"), ), ) def test_binary_operations(self, func, dtype): array = np.arange(10).astype(dtype) * unit_registry.m data_array = xr.DataArray(data=array) - expected = xr.DataArray(data=func(array)) - result = func(data_array) + units = extract_units(func(array)) + expected = attach_units(func(strip_units(data_array)), units) + actual = func(data_array) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "comparison", @@ -1448,8 +1443,9 @@ def test_binary_operations(self, func, dtype): pytest.param( unit_registry.dimensionless, DimensionalityError, id="dimensionless" ), - pytest.param(unit_registry.s, DimensionalityError, id="incorrect_unit"), - pytest.param(unit_registry.m, None, id="correct_unit"), + pytest.param(unit_registry.s, DimensionalityError, id="incompatible_unit"), + pytest.param(unit_registry.mm, None, id="compatible_unit"), + pytest.param(unit_registry.m, None, id="identical_unit"), ), ) def test_comparison_operations(self, comparison, unit, error, dtype): @@ -1469,48 +1465,85 @@ def test_comparison_operations(self, comparison, unit, error, dtype): with pytest.raises(error): comparison(data_array, to_compare_with) - else: - result = comparison(data_array, to_compare_with) - # pint compares incompatible arrays to False, so we need to extend - # the multiplication works for both scalar and array results - expected = xr.DataArray( - data=comparison(array, to_compare_with) - * np.ones_like(array, dtype=bool) - ) - assert_equal_with_units(expected, result) + return + + actual = comparison(data_array, to_compare_with) + + expected_units = {None: unit_registry.m if array.check(unit) else None} + expected = array.check(unit) & comparison( + strip_units(data_array), + strip_units(convert_units(to_compare_with, expected_units)), + ) + + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "units,error", ( pytest.param(unit_registry.dimensionless, None, id="dimensionless"), - pytest.param(unit_registry.m, DimensionalityError, id="incorrect unit"), - pytest.param(unit_registry.degree, None, id="correct unit"), + pytest.param(unit_registry.m, DimensionalityError, id="incompatible_unit"), + pytest.param(unit_registry.degree, None, id="compatible_unit"), ), ) def test_univariate_ufunc(self, units, error, dtype): array = np.arange(10).astype(dtype) * units data_array = xr.DataArray(data=array) + func = function("sin") + if error is not None: with pytest.raises(error): np.sin(data_array) - else: - expected = xr.DataArray(data=np.sin(array)) - result = np.sin(data_array) - assert_equal_with_units(expected, result) + return + + expected = attach_units( + func(strip_units(convert_units(data_array, {None: unit_registry.radians}))), + {None: unit_registry.dimensionless}, + ) + actual = func(data_array) + + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="pint's implementation of `np.maximum` strips units") - def test_bivariate_ufunc(self, dtype): - unit = unit_registry.m - array = np.arange(10).astype(dtype) * unit + @pytest.mark.xfail(reason="xarray's `np.maximum` strips units") + @pytest.mark.parametrize( + "unit,error", + ( + pytest.param(1, DimensionalityError, id="without_unit"), + pytest.param( + 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.m, None, id="identical_unit"), + ), + ) + def test_bivariate_ufunc(self, unit, error, dtype): + original_unit = unit_registry.m + array = np.arange(10).astype(dtype) * original_unit data_array = xr.DataArray(data=array) - expected = xr.DataArray(np.maximum(array, 0 * unit)) + if error is not None: + with pytest.raises(error): + np.maximum(data_array, 0 * unit) + + return + + expected_units = {None: original_unit} + expected = attach_units( + np.maximum( + strip_units(data_array), + strip_units(convert_units(0 * unit, expected_units)), + ), + expected_units, + ) + + actual = np.maximum(data_array, 0 * unit) + assert_equal_with_units(expected, actual) - assert_equal_with_units(expected, np.maximum(data_array, 0 * unit)) - assert_equal_with_units(expected, np.maximum(0 * unit, data_array)) + actual = np.maximum(0 * unit, data_array) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize("property", ("T", "imag", "real")) def test_numpy_properties(self, property, dtype): @@ -1518,41 +1551,43 @@ def test_numpy_properties(self, property, dtype): np.arange(5 * 10).astype(dtype) + 1j * np.linspace(-1, 0, 5 * 10).astype(dtype) ).reshape(5, 10) * unit_registry.s + data_array = xr.DataArray(data=array, dims=("x", "y")) - expected = xr.DataArray( - data=getattr(array, property), - dims=("x", "y")[:: 1 if property != "T" else -1], + expected = attach_units( + getattr(strip_units(data_array), property), extract_units(data_array) ) - result = getattr(data_array, property) + actual = getattr(data_array, property) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", - ( - method("conj"), - method("argsort"), - method("conjugate"), - method("round"), - pytest.param( - method("rank", dim="x"), - marks=pytest.mark.xfail(reason="pint does not implement rank yet"), - ), - ), + (method("conj"), method("argsort"), method("conjugate"), method("round")), ids=repr, ) def test_numpy_methods(self, func, dtype): array = np.arange(10).astype(dtype) * unit_registry.m data_array = xr.DataArray(data=array, dims="x") - expected = xr.DataArray(func(array), dims="x") - result = func(data_array) + units = extract_units(func(array)) + expected = attach_units(strip_units(data_array), units) + actual = func(data_array) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( - "func", (method("clip", min=3, max=8), method("searchsorted", v=5)), ids=repr + "func", + ( + method("clip", min=3, max=8), + pytest.param( + method("searchsorted", v=5), + marks=pytest.mark.xfail( + reason="searchsorted somehow requires a undocumented `keys` argument" + ), + ), + ), + ids=repr, ) @pytest.mark.parametrize( "unit,error", @@ -1575,20 +1610,24 @@ def test_numpy_methods_with_args(self, func, unit, error, dtype): 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) - else: - expected = func(array, **kwargs) - if func.name not in ["searchsorted"]: - expected = xr.DataArray(data=expected) - result = func(data_array, **kwargs) - if func.name in ["searchsorted"]: - assert np.allclose(expected, result) - else: - assert_equal_with_units(expected, result) + return + + units = extract_units(data_array) + expected_units = extract_units(func(array, **kwargs)) + 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 + ) + actual = func(data_array, **kwargs) + + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", (method("isnull"), method("notnull"), method("count")), ids=repr @@ -1611,9 +1650,9 @@ def test_missing_value_detection(self, func, dtype): data_array = xr.DataArray(data=array, coords={"x": x, "y": y}, dims=("x", "y")) expected = func(strip_units(data_array)) - result = func(data_array) + actual = func(data_array) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.xfail(reason="ffill and bfill lose units in data") @pytest.mark.parametrize("func", (method("ffill"), method("bfill")), ids=repr) @@ -1623,48 +1662,67 @@ def test_missing_value_filling(self, func, dtype): * unit_registry.degK ) x = np.arange(len(array)) - data_array = xr.DataArray(data=array, coords={"x": x}, dims=["x"]) - - result_without_units = func(strip_units(data_array), dim="x") - result = xr.DataArray( - data=result_without_units.data * unit_registry.degK, - coords={"x": x}, - dims=["x"], - ) + data_array = xr.DataArray(data=array, coords={"x": x}, dims="x") expected = attach_units( - func(strip_units(data_array), dim="x"), {"data": unit_registry.degK} + func(strip_units(data_array), dim="x"), extract_units(data_array) ) - result = func(data_array, dim="x") + actual = func(data_array, dim="x") - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="fillna drops the unit") @pytest.mark.parametrize( - "fill_value", + "unit,error", ( + pytest.param(1, DimensionalityError, id="no_unit"), pytest.param( - -1, - id="python scalar", - marks=pytest.mark.xfail( - reason="python scalar cannot be converted using astype()" - ), + 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(np.array(-1), id="numpy scalar"), - pytest.param(np.array([-1]), id="numpy array"), + pytest.param(unit_registry.m, None, id="identical_unit"), + ), + ) + @pytest.mark.parametrize( + "fill_value", + ( + pytest.param(-1, id="python_scalar"), + pytest.param(np.array(-1), id="numpy_scalar"), + pytest.param(np.array([-1]), id="numpy_array"), ), ) - def test_fillna(self, fill_value, dtype): - unit = unit_registry.m - array = np.array([1.4, np.nan, 2.3, np.nan, np.nan, 9.1]).astype(dtype) * unit + def test_fillna(self, fill_value, unit, error, dtype): + original_unit = unit_registry.m + array = ( + np.array([1.4, np.nan, 2.3, np.nan, np.nan, 9.1]).astype(dtype) + * original_unit + ) data_array = xr.DataArray(data=array) + func = method("fillna") + + value = fill_value * unit + if error is not None: + with pytest.raises(error): + func(data_array, value=value) + + return + + units = extract_units(data_array) expected = attach_units( - strip_units(data_array).fillna(value=fill_value), {"data": unit} + func( + strip_units(data_array), value=strip_units(convert_units(value, units)) + ), + units, ) - result = data_array.fillna(value=fill_value * unit) + actual = func(data_array, value=value) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) def test_dropna(self, dtype): array = ( @@ -1674,22 +1732,26 @@ def test_dropna(self, dtype): x = np.arange(len(array)) data_array = xr.DataArray(data=array, coords={"x": x}, dims=["x"]) - expected = attach_units( - strip_units(data_array).dropna(dim="x"), {"data": unit_registry.m} - ) - result = data_array.dropna(dim="x") + units = extract_units(data_array) + expected = attach_units(strip_units(data_array).dropna(dim="x"), units) + actual = data_array.dropna(dim="x") - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="pint does not implement `numpy.isin`") @pytest.mark.parametrize( "unit", ( - pytest.param(1, id="no_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(unit_registry.dimensionless, id="dimensionless"), pytest.param(unit_registry.s, id="incompatible_unit"), pytest.param(unit_registry.cm, id="compatible_unit"), - pytest.param(unit_registry.m, id="same_unit"), + pytest.param(unit_registry.m, id="identical_unit"), ), ) def test_isin(self, unit, dtype): @@ -1702,33 +1764,26 @@ def test_isin(self, unit, dtype): raw_values = np.array([1.4, np.nan, 2.3]).astype(dtype) values = raw_values * unit - result_without_units = strip_units(data_array).isin(raw_values) - if unit != unit_registry.m: - result_without_units[:] = False - result_with_units = data_array.isin(values) + units = {None: unit_registry.m if array.check(unit) else None} + expected = strip_units(data_array).isin( + strip_units(convert_units(values, units)) + ) & array.check(unit) + actual = data_array.isin(values) - assert_equal_with_units(result_without_units, result_with_units) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "variant", ( pytest.param( "masking", - marks=pytest.mark.xfail(reason="nan not compatible with quantity"), - ), - pytest.param( - "replacing_scalar", - marks=pytest.mark.xfail(reason="scalar not convertible using astype"), - ), - pytest.param( - "replacing_array", - marks=pytest.mark.xfail( - reason="replacing using an array drops the units" - ), + marks=pytest.mark.xfail(reason="array(nan) is not a quantity"), ), + "replacing_scalar", + "replacing_array", pytest.param( "dropping", - marks=pytest.mark.xfail(reason="nan not compatible with quantity"), + marks=pytest.mark.xfail(reason="array(nan) is not a quantity"), ), ), ) @@ -1741,13 +1796,10 @@ def test_isin(self, unit, dtype): ), 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="same_unit"), + pytest.param(unit_registry.m, None, id="identical_unit"), ), ) def test_where(self, variant, unit, error, dtype): - def _strip_units(mapping): - return {key: array_strip_units(value) for key, value in mapping.items()} - original_unit = unit_registry.m array = np.linspace(0, 1, 10).astype(dtype) * original_unit @@ -1762,19 +1814,28 @@ def _strip_units(mapping): "dropping": {"cond": condition, "drop": True}, } kwargs = variant_kwargs.get(variant) - kwargs_without_units = _strip_units(kwargs) + kwargs_without_units = { + key: strip_units( + convert_units( + value, {None: original_unit if array.check(unit) else None} + ) + ) + for key, value in kwargs.items() + } if variant not in ("masking", "dropping") and error is not None: with pytest.raises(error): data_array.where(**kwargs) - else: - expected = attach_units( - strip_units(array).where(**kwargs_without_units), - {"data": original_unit}, - ) - result = data_array.where(**kwargs) - assert_equal_with_units(expected, result) + return + + expected = attach_units( + strip_units(data_array).where(**kwargs_without_units), + extract_units(data_array), + ) + actual = data_array.where(**kwargs) + + assert_equal_with_units(expected, actual) @pytest.mark.xfail(reason="interpolate strips units") def test_interpolate_na(self, dtype): @@ -1785,14 +1846,12 @@ def test_interpolate_na(self, dtype): x = np.arange(len(array)) data_array = xr.DataArray(data=array, coords={"x": x}, dims="x").astype(dtype) - expected = attach_units( - strip_units(data_array).interpolate_na(dim="x"), {"data": unit_registry.m} - ) - result = data_array.interpolate_na(dim="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, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="uses DataArray.where, which currently fails") @pytest.mark.parametrize( "unit,error", ( @@ -1801,8 +1860,18 @@ 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"), - pytest.param(unit_registry.m, None, id="identical_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"), + ), ), ) def test_combine_first(self, unit, error, dtype): @@ -1819,14 +1888,19 @@ def test_combine_first(self, unit, error, dtype): if error is not None: with pytest.raises(error): data_array.combine_first(other) - else: - expected = attach_units( - strip_units(data_array).combine_first(strip_units(other)), - {"data": unit_registry.m}, - ) - result = data_array.combine_first(other) - assert_equal_with_units(expected, result) + return + + units = extract_units(data_array) + expected = attach_units( + strip_units(data_array).combine_first( + strip_units(convert_units(other, units)) + ), + units, + ) + actual = data_array.combine_first(other) + + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "unit", @@ -1834,11 +1908,7 @@ def test_combine_first(self, unit, error, dtype): 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", - marks=pytest.mark.xfail(reason="identical does not check units yet"), - ), + pytest.param(unit_registry.cm, id="compatible_unit"), pytest.param(unit_registry.m, id="identical_unit"), ), ) @@ -1854,53 +1924,51 @@ def test_combine_first(self, unit, error, dtype): ) @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) + data = np.linspace(0, 5, 10).astype(dtype) coord = np.arange(len(data)).astype(dtype) base_unit = unit_registry.m - quantity = data * base_unit - x = coord * base_unit - y = coord * base_unit - - units = { - "data": (unit, base_unit, base_unit), - "dims": (base_unit, unit, base_unit), - "coords": (base_unit, base_unit, unit), + array = data * (base_unit if variation == "data" else 1) + x = coord * (base_unit if variation == "dims" else 1) + y = coord * (base_unit if variation == "coords" else 1) + + variations = { + "data": (unit, 1, 1), + "dims": (1, unit, 1), + "coords": (1, 1, unit), } - data_unit, dim_unit, coord_unit = units.get(variation) + data_unit, dim_unit, coord_unit = variations.get(variation) - data_array = xr.DataArray( - data=quantity, coords={"x": x, "y": ("x", y)}, dims="x" - ) + data_array = xr.DataArray(data=array, coords={"x": x, "y": ("x", y)}, dims="x") other = attach_units( - strip_units(data_array), - { - None: (data_unit, base_unit if quantity.check(data_unit) else None), - "x": (dim_unit, base_unit if x.check(dim_unit) else None), - "y": (coord_unit, base_unit if y.check(coord_unit) else None), - }, + strip_units(data_array), {None: data_unit, "x": dim_unit, "y": coord_unit} ) - # TODO: test dim coord once indexes leave units intact - # also, express this in terms of calls on the raw data array - # and then check the units - equal_arrays = ( - np.all(quantity == other.data) - and (np.all(x == other.x.data) or True) # dims can't be checked yet - and np.all(y == other.y.data) - ) - equal_units = ( - data_unit == unit_registry.m - and coord_unit == unit_registry.m - and dim_unit == unit_registry.m + units = extract_units(data_array) + other_units = extract_units(other) + + equal_arrays = all( + is_compatible(units[name], other_units[name]) for name in units.keys() + ) and ( + strip_units(data_array).equals( + strip_units(convert_units(other, extract_units(data_array))) + ) ) + equal_units = units == other_units expected = equal_arrays and (func.name != "identical" or equal_units) - result = func(data_array, other) - assert expected == result + actual = func(data_array, other) + + assert expected == actual - @pytest.mark.xfail(reason="blocked by `where`") @pytest.mark.parametrize( "unit", ( @@ -1926,9 +1994,9 @@ def test_broadcast_like(self, unit, dtype): expected = attach_units( strip_units(arr1).broadcast_like(strip_units(arr2)), extract_units(arr1) ) - result = arr1.broadcast_like(arr2) + actual = arr1.broadcast_like(arr2) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "unit", @@ -1942,19 +2010,21 @@ def test_broadcast_like(self, unit, dtype): ) def test_broadcast_equals(self, unit, dtype): left_array = np.ones(shape=(2, 2), dtype=dtype) * unit_registry.m - right_array = array_attach_units( - np.ones(shape=(2,), dtype=dtype), - unit, - convert_from=unit_registry.m if left_array.check(unit) else None, - ) + right_array = np.ones(shape=(2,), dtype=dtype) * unit left = xr.DataArray(data=left_array, dims=("x", "y")) right = xr.DataArray(data=right_array, dims="x") - expected = np.all(left_array == right_array[:, None]) - result = left.broadcast_equals(right) + units = { + **extract_units(left), + **({} if left_array.check(unit) else {None: None}), + } + expected = strip_units(left).broadcast_equals( + strip_units(convert_units(right, units)) + ) & left_array.check(unit) + actual = left.broadcast_equals(right) - assert expected == result + assert expected == actual @pytest.mark.parametrize( "func", @@ -1969,16 +2039,11 @@ def test_broadcast_equals(self, unit, dtype): dim={"z": np.linspace(10, 20, 12) * unit_registry.s}, axis=1, ), - method("drop_sel", labels="x"), + method("drop_vars", "x"), method("reset_coords", names="x2"), method("copy"), - pytest.param( - method("astype", np.float32), - marks=pytest.mark.xfail(reason="units get stripped"), - ), - pytest.param( - method("item", 1), marks=pytest.mark.xfail(reason="units get stripped") - ), + method("astype", np.float32), + method("item", 1), ), ids=repr, ) @@ -2001,67 +2066,38 @@ def test_content_manipulation(self, func, dtype): stripped_kwargs = { key: array_strip_units(value) for key, value in func.kwargs.items() } - expected = attach_units( - func(strip_units(data_array), **stripped_kwargs), - { - "data": quantity.units, - "x": x.units, - "x_mm": x2.units, - "x2": x2.units, - "y": y.units, - }, - ) - result = func(data_array) + units = {**{"x_mm": x2.units, "x2": x2.units}, **extract_units(data_array)} - assert_equal_with_units(expected, result) + expected = attach_units(func(strip_units(data_array), **stripped_kwargs), units) + actual = func(data_array) + + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( - "func", - ( - pytest.param( - method("drop_sel", labels=dict(x=np.array([1, 5]))), - marks=pytest.mark.xfail( - reason="selecting using incompatible units does not raise" - ), - ), - pytest.param(method("copy", data=np.arange(20))), - ), - ids=repr, + "func", (pytest.param(method("copy", data=np.arange(20))),), ids=repr ) @pytest.mark.parametrize( - "unit,error", + "unit", ( - 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, KeyError, id="compatible_unit"), - pytest.param(unit_registry.m, None, id="identical_unit"), + pytest.param(1, id="no_unit"), + pytest.param(unit_registry.dimensionless, id="dimensionless"), + pytest.param(unit_registry.degK, id="with_unit"), ), ) - def test_content_manipulation_with_units(self, func, unit, error, dtype): + def test_content_manipulation_with_units(self, func, unit, dtype): quantity = np.linspace(0, 10, 20, dtype=dtype) * unit_registry.pascal x = np.arange(len(quantity)) * unit_registry.m - data_array = xr.DataArray(name="data", data=quantity, coords={"x": x}, dims="x") + data_array = xr.DataArray(data=quantity, coords={"x": x}, dims="x") - kwargs = { - key: (value * unit if isinstance(value, np.ndarray) else value) - for key, value in func.kwargs.items() - } - stripped_kwargs = func.kwargs + kwargs = {key: value * unit for key, value in func.kwargs.items()} expected = attach_units( - func(strip_units(data_array), **stripped_kwargs), - {"data": quantity.units if func.name == "drop_sel" else unit, "x": x.units}, + func(strip_units(data_array)), {None: unit, "x": x.units} ) - if error is not None and func.name == "drop_sel": - with pytest.raises(error): - func(data_array, **kwargs) - else: - result = func(data_array, **kwargs) - assert_equal_with_units(expected, result) + + actual = func(data_array, **kwargs) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "indices", @@ -2074,95 +2110,152 @@ def test_isel(self, indices, dtype): array = np.arange(10).astype(dtype) * unit_registry.s x = np.arange(len(array)) * unit_registry.m - data_array = xr.DataArray(data=array, coords={"x": x}, dims=["x"]) + data_array = xr.DataArray(data=array, coords={"x": x}, dims="x") expected = attach_units( - strip_units(data_array).isel(x=indices), - {"data": unit_registry.s, "x": unit_registry.m}, + strip_units(data_array).isel(x=indices), extract_units(data_array) ) - result = data_array.isel(x=indices) + actual = data_array.isel(x=indices) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail( - reason="xarray does not support duck arrays in dimension coordinates" - ) + @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( - "values", + "raw_values", ( - pytest.param(12, id="single value"), - pytest.param([10, 5, 13], id="list of multiple values"), - pytest.param(np.array([9, 3, 7, 12]), id="array of multiple values"), + pytest.param(10, id="single_value"), + pytest.param([10, 5, 13], id="list_of_values"), + pytest.param(np.array([9, 3, 7, 12]), id="array_of_values"), ), ) @pytest.mark.parametrize( - "units,error", + "unit,error", ( - pytest.param(1, KeyError, id="no units"), + pytest.param(1, KeyError, id="no_units"), pytest.param(unit_registry.dimensionless, KeyError, id="dimensionless"), - pytest.param(unit_registry.degree, KeyError, id="incorrect unit"), - pytest.param(unit_registry.s, None, id="correct unit"), + pytest.param(unit_registry.degree, KeyError, id="incompatible_unit"), + pytest.param(unit_registry.dm, KeyError, id="compatible_unit"), + pytest.param(unit_registry.m, None, id="identical_unit"), ), ) - def test_sel(self, values, units, error, dtype): + def test_sel(self, raw_values, unit, error, dtype): array = np.linspace(5, 10, 20).astype(dtype) * unit_registry.m - x = np.arange(len(array)) * unit_registry.s - data_array = xr.DataArray(data=array, coords={"x": x}, dims=["x"]) + x = np.arange(len(array)) * unit_registry.m + data_array = xr.DataArray(data=array, coords={"x": x}, dims="x") - values_with_units = values * units + values = raw_values * unit - if error is not None: + if error is not None and not ( + isinstance(raw_values, (int, float)) and x.check(unit) + ): with pytest.raises(error): - data_array.sel(x=values_with_units) - else: - result_array = array[values] - result_data_array = data_array.sel(x=values_with_units) - assert_equal_with_units(result_array, result_data_array) + data_array.sel(x=values) + + return - @pytest.mark.xfail( - reason="xarray does not support duck arrays in dimension coordinates" + expected = attach_units( + strip_units(data_array).sel( + x=strip_units(convert_units(values, {None: array.units})) + ), + extract_units(data_array), + ) + actual = data_array.sel(x=values) + assert_equal_with_units(expected, actual) + + @pytest.mark.xfail(reason="indexes don't support units") + @pytest.mark.parametrize( + "raw_values", + ( + pytest.param(10, id="single_value"), + pytest.param([10, 5, 13], id="list_of_values"), + pytest.param(np.array([9, 3, 7, 12]), id="array_of_values"), + ), + ) + @pytest.mark.parametrize( + "unit,error", + ( + 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.m, None, id="identical_unit"), + ), ) + def test_loc(self, raw_values, unit, error, dtype): + array = np.linspace(5, 10, 20).astype(dtype) * unit_registry.m + x = np.arange(len(array)) * unit_registry.m + data_array = xr.DataArray(data=array, coords={"x": x}, dims="x") + + values = raw_values * unit + + if error is not None and not ( + isinstance(raw_values, (int, float)) and x.check(unit) + ): + with pytest.raises(error): + data_array.loc[{"x": values}] + + return + + expected = attach_units( + strip_units(data_array).loc[ + {"x": strip_units(convert_units(values, {None: array.units}))} + ], + extract_units(data_array), + ) + actual = data_array.loc[{"x": values}] + assert_equal_with_units(expected, actual) + + @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( - "values", + "raw_values", ( - pytest.param(12, id="single value"), - pytest.param([10, 5, 13], id="list of multiple values"), - pytest.param(np.array([9, 3, 7, 12]), id="array of multiple values"), + pytest.param(10, id="single_value"), + pytest.param([10, 5, 13], id="list_of_values"), + pytest.param(np.array([9, 3, 7, 12]), id="array_of_values"), ), ) @pytest.mark.parametrize( - "units,error", + "unit,error", ( - pytest.param(1, KeyError, id="no units"), + pytest.param(1, KeyError, id="no_units"), pytest.param(unit_registry.dimensionless, KeyError, id="dimensionless"), - pytest.param(unit_registry.degree, KeyError, id="incorrect unit"), - pytest.param(unit_registry.s, None, id="correct unit"), + pytest.param(unit_registry.degree, KeyError, id="incompatible_unit"), + pytest.param(unit_registry.dm, KeyError, id="compatible_unit"), + pytest.param(unit_registry.m, None, id="identical_unit"), ), ) - def test_loc(self, values, units, error, dtype): + def test_drop_sel(self, raw_values, unit, error, dtype): array = np.linspace(5, 10, 20).astype(dtype) * unit_registry.m - x = np.arange(len(array)) * unit_registry.s - data_array = xr.DataArray(data=array, coords={"x": x}, dims=["x"]) + x = np.arange(len(array)) * unit_registry.m + data_array = xr.DataArray(data=array, coords={"x": x}, dims="x") - values_with_units = values * units + values = raw_values * unit - if error is not None: + if error is not None and not ( + isinstance(raw_values, (int, float)) and x.check(unit) + ): with pytest.raises(error): - data_array.loc[values_with_units] - else: - result_array = array[values] - result_data_array = data_array.loc[values_with_units] - assert_equal_with_units(result_array, result_data_array) + data_array.drop_sel(x=values) + + return + + expected = attach_units( + strip_units(data_array).drop_sel( + x=strip_units(convert_units(values, {None: x.units})) + ), + extract_units(data_array), + ) + actual = data_array.drop_sel(x=values) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="tries to coerce using asarray") @pytest.mark.parametrize( "shape", ( - pytest.param((10, 20), id="nothing squeezable"), - pytest.param((10, 20, 1), id="last dimension squeezable"), - pytest.param((10, 1, 20), id="middle dimension squeezable"), - pytest.param((1, 10, 20), id="first dimension squeezable"), - pytest.param((1, 10, 1, 20), id="first and last dimension squeezable"), + pytest.param((10, 20), id="nothing_squeezable"), + pytest.param((10, 20, 1), id="last_dimension_squeezable"), + pytest.param((10, 1, 20), id="middle_dimension_squeezable"), + pytest.param((1, 10, 20), id="first_dimension_squeezable"), + pytest.param((1, 10, 1, 20), id="first_and_last_dimension_squeezable"), ), ) def test_squeeze(self, shape, dtype): @@ -2177,38 +2270,27 @@ def test_squeeze(self, shape, dtype): data=array, coords=coords, dims=tuple(names[: len(shape)]) ) - result_array = array.squeeze() - result_data_array = data_array.squeeze() - assert_equal_with_units(result_array, result_data_array) + expected = attach_units( + strip_units(data_array).squeeze(), extract_units(data_array) + ) + actual = data_array.squeeze() + assert_equal_with_units(expected, actual) # try squeezing the dimensions separately names = tuple(dim for dim, coord in coords.items() if len(coord) == 1) for index, name in enumerate(names): - assert_equal_with_units( - np.squeeze(array, axis=index), data_array.squeeze(dim=name) + expected = attach_units( + strip_units(data_array).squeeze(dim=name), extract_units(data_array) ) + actual = data_array.squeeze(dim=name) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail( - reason="indexes strip units and head / tail / thin only support integers" - ) - @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("head", x=7, y=3), method("tail", x=7, y=3), method("thin", x=7, y=3)), ids=repr, ) - def test_head_tail_thin(self, func, unit, error, dtype): + def test_head_tail_thin(self, func, dtype): array = np.linspace(1, 2, 10 * 5).reshape(10, 5) * unit_registry.degK coords = { @@ -2216,27 +2298,24 @@ def test_head_tail_thin(self, func, unit, error, dtype): "y": np.arange(5) * unit_registry.m, } - arr = xr.DataArray(data=array, coords=coords, dims=("x", "y")) + data_array = xr.DataArray(data=array, coords=coords, dims=("x", "y")) - kwargs = {name: value * unit for name, value in func.kwargs.items()} - - if error is not None: - with pytest.raises(error): - func(arr, **kwargs) - - return - - expected = attach_units(func(strip_units(arr)), extract_units(arr)) - result = func(arr, **kwargs) + expected = attach_units( + func(strip_units(data_array)), extract_units(data_array) + ) + actual = func(data_array) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) + @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( "unit,error", ( - pytest.param(1, None, id="no_unit"), - pytest.param(unit_registry.dimensionless, None, id="dimensionless"), - pytest.param(unit_registry.s, None, id="incompatible_unit"), + 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"), ), @@ -2254,24 +2333,29 @@ def test_interp(self, unit, error): if error is not None: with pytest.raises(error): data_array.interp(x=new_coords) - else: - new_coords_ = ( - new_coords.magnitude if hasattr(new_coords, "magnitude") else new_coords - ) - result_array = strip_units(data_array).interp( - x=new_coords_ * unit_registry.degK - ) - result_data_array = data_array.interp(x=new_coords) - assert_equal_with_units(result_array, result_data_array) + return + + 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) + + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="tries to coerce using asarray") + @pytest.mark.xfail(reason="indexes strip units") @pytest.mark.parametrize( "unit,error", ( - pytest.param(1, None, id="no_unit"), - pytest.param(unit_registry.dimensionless, None, id="dimensionless"), - pytest.param(unit_registry.s, None, id="incompatible_unit"), + 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"), ), @@ -2284,43 +2368,46 @@ def test_interp_like(self, unit, error): } data_array = xr.DataArray(array, coords=coords, dims=("x", "y")) - new_data_array = xr.DataArray( - data=np.empty((20, 10)), + 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"), ) if error is not None: with pytest.raises(error): - data_array.interp_like(new_data_array) - else: - result_array = ( - xr.DataArray( - data=array.magnitude, - coords={name: value.magnitude for name, value in coords.items()}, - dims=("x", "y"), - ).interp_like(strip_units(new_data_array)) - * unit_registry.degK - ) - result_data_array = data_array.interp_like(new_data_array) + data_array.interp_like(other) - assert_equal_with_units(result_array, result_data_array) + return - @pytest.mark.xfail( - reason="pint does not implement np.result_type in __array_function__ yet" - ) + units = extract_units(data_array) + expected = attach_units( + strip_units(data_array).interp_like( + strip_units(convert_units(other, units)) + ), + units, + ) + actual = data_array.interp_like(other) + + assert_equal_with_units(expected, actual) + + @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( "unit,error", ( - pytest.param(1, None, id="no_unit"), - pytest.param(unit_registry.dimensionless, None, id="dimensionless"), - pytest.param(unit_registry.s, None, id="incompatible_unit"), + 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"), ), ) - def test_reindex(self, unit, error): - array = np.linspace(1, 2, 10 * 5).reshape(10, 5) * unit_registry.degK + 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, @@ -2328,65 +2415,70 @@ def test_reindex(self, unit, error): } data_array = xr.DataArray(array, coords=coords, dims=("x", "y")) + func = method("reindex") if error is not None: with pytest.raises(error): - data_array.interp(x=new_coords) - else: - result_array = strip_units(data_array).reindex( - x=( - new_coords.magnitude - if hasattr(new_coords, "magnitude") - else new_coords - ) - * unit_registry.degK - ) - result_data_array = data_array.reindex(x=new_coords) + func(data_array, x=new_coords) - assert_equal_with_units(result_array, result_data_array) + return - @pytest.mark.xfail( - reason="pint does not implement np.result_type in __array_function__ yet" - ) + 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) + + assert_equal_with_units(expected, actual) + + @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( "unit,error", ( - pytest.param(1, None, id="no_unit"), - pytest.param(unit_registry.dimensionless, None, id="dimensionless"), - pytest.param(unit_registry.s, None, id="incompatible_unit"), + 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"), ), ) - def test_reindex_like(self, unit, error): - array = np.linspace(1, 2, 10 * 5).reshape(10, 5) * unit_registry.degK + 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, } data_array = xr.DataArray(array, coords=coords, dims=("x", "y")) - new_data_array = xr.DataArray( - data=np.empty((20, 10)), + 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"), ) if error is not None: with pytest.raises(error): - data_array.reindex_like(new_data_array) - else: - expected = attach_units( - strip_units(data_array).reindex_like(strip_units(new_data_array)), - { - "data": unit_registry.degK, - "x": unit_registry.m, - "y": unit_registry.m, - }, - ) - result = data_array.reindex_like(new_data_array) + data_array.reindex_like(other) - assert_equal_with_units(expected, result) + return + + units = extract_units(data_array) + expected = attach_units( + strip_units(data_array).reindex_like( + strip_units(convert_units(other, units)) + ), + units, + ) + actual = data_array.reindex_like(other) + + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", @@ -2406,11 +2498,11 @@ def test_stacking_stacked(self, func, dtype): stacked = data_array.stack(z=("x", "y")) expected = attach_units(func(strip_units(stacked)), {"data": unit_registry.m}) - result = func(stacked) + actual = func(stacked) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="indexes strip the label units") + @pytest.mark.xfail(reason="indexes don't support units") def test_to_unstacked_dataset(self, dtype): array = ( np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) @@ -2429,13 +2521,9 @@ def test_to_unstacked_dataset(self, dtype): func(strip_units(data_array)), {"y": y.units, **dict(zip(x.magnitude, [array.units] * len(y)))}, ).rename({elem.magnitude: elem for elem in x}) - result = func(data_array) - - print(data_array, expected, result, sep="\n") + actual = func(data_array) - assert_equal_with_units(expected, result) - - assert False + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", @@ -2446,10 +2534,7 @@ def test_to_unstacked_dataset(self, dtype): pytest.param( method("shift", x=2), marks=pytest.mark.xfail(reason="strips units") ), - pytest.param( - method("roll", x=2, roll_coords=False), - marks=pytest.mark.xfail(reason="strips units"), - ), + method("roll", x=2, roll_coords=False), method("sortby", "x2"), ), ids=repr, @@ -2471,12 +2556,10 @@ def test_stacking_reordering(self, func, dtype): dims=("x", "y", "z"), ) - expected = attach_units( - func(strip_units(data_array)), {"data": unit_registry.m} - ) - result = func(data_array) + expected = attach_units(func(strip_units(data_array)), {None: unit_registry.m}) + actual = func(data_array) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", @@ -2484,20 +2567,14 @@ def test_stacking_reordering(self, func, dtype): method("diff", dim="x"), method("differentiate", coord="x"), method("integrate", dim="x"), - pytest.param( - method("quantile", q=[0.25, 0.75]), - marks=pytest.mark.xfail( - reason="pint does not implement nanpercentile yet" - ), - ), - pytest.param( - method("reduce", func=np.sum, dim="x"), - marks=pytest.mark.xfail(reason="strips units"), - ), + method("quantile", q=[0.25, 0.75]), + 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 einsum"), + marks=pytest.mark.xfail( + reason="pint does not implement the dot method" + ), ), ), ids=repr, @@ -2511,30 +2588,35 @@ def test_computation(self, func, dtype): y = np.arange(array.shape[1]) * unit_registry.s data_array = xr.DataArray(data=array, coords={"x": x, "y": y}, dims=("x", "y")) - units = extract_units(data_array) + + # we want to make sure the output unit is correct + units = { + **extract_units(data_array), + **( + {} + if isinstance(func, (function, method)) + else extract_units(func(array.reshape(-1))) + ), + } expected = attach_units(func(strip_units(data_array)), units) - result = func(data_array) + actual = func(data_array) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", ( - pytest.param( - method("groupby", "y"), marks=pytest.mark.xfail(reason="strips units") - ), - pytest.param( - method("groupby_bins", "y", bins=4), - marks=pytest.mark.xfail(reason="strips units"), - ), + method("groupby", "x"), + method("groupby_bins", "y", bins=4), method("coarsen", y=2), pytest.param( - method("rolling", y=3), marks=pytest.mark.xfail(reason="strips units") + method("rolling", y=3), + marks=pytest.mark.xfail(reason="rolling strips units"), ), pytest.param( method("rolling_exp", y=3), - marks=pytest.mark.xfail(reason="strips units"), + marks=pytest.mark.xfail(reason="units not supported by numbagg"), ), ), ids=repr, @@ -2544,18 +2626,17 @@ def test_computation_objects(self, func, dtype): np.linspace(0, 10, 5 * 10).reshape(5, 10).astype(dtype) * unit_registry.m ) - x = np.arange(array.shape[0]) * unit_registry.m + x = np.array([0, 0, 1, 2, 2]) * unit_registry.m y = np.arange(array.shape[1]) * 3 * unit_registry.s data_array = xr.DataArray(data=array, coords={"x": x, "y": y}, dims=("x", "y")) units = extract_units(data_array) expected = attach_units(func(strip_units(data_array)).mean(), units) - result = func(data_array).mean() + actual = func(data_array).mean() - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="strips units") def test_resample(self, dtype): array = np.linspace(0, 5, 10).astype(dtype) * unit_registry.m @@ -2566,23 +2647,17 @@ def test_resample(self, dtype): func = method("resample", time="6m") expected = attach_units(func(strip_units(data_array)).mean(), units) - result = func(data_array).mean() + actual = func(data_array).mean() - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", ( - pytest.param( - method("assign_coords", {"z": (["x"], np.arange(5) * unit_registry.s)}), - marks=pytest.mark.xfail(reason="strips units"), - ), - pytest.param(method("first")), - pytest.param(method("last")), - pytest.param( - method("quantile", q=[0.25, 0.5, 0.75], dim="x"), - marks=pytest.mark.xfail(reason="strips units"), - ), + method("assign_coords", z=(["x"], np.arange(5) * unit_registry.s)), + method("first"), + method("last"), + method("quantile", q=np.array([0.25, 0.5, 0.75]), dim="x"), ), ids=repr, ) @@ -2595,12 +2670,22 @@ def test_grouped_operations(self, func, dtype): y = np.arange(array.shape[1]) * 3 * unit_registry.s data_array = xr.DataArray(data=array, coords={"x": x, "y": y}, dims=("x", "y")) - units = extract_units(data_array) + units = {**extract_units(data_array), **{"z": unit_registry.s, "q": None}} - expected = attach_units(func(strip_units(data_array).groupby("y")), units) - result = func(data_array.groupby("y")) + stripped_kwargs = { + key: ( + strip_units(value) + if not isinstance(value, tuple) + else tuple(strip_units(elem) for elem in value) + ) + for key, value in func.kwargs.items() + } + expected = attach_units( + func(strip_units(data_array).groupby("y"), **stripped_kwargs), units + ) + actual = func(data_array.groupby("y")) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) class TestDataset: @@ -2620,10 +2705,7 @@ class TestDataset: "shared", ( "nothing", - pytest.param( - "dims", - marks=pytest.mark.xfail(reason="reindex does not work with pint yet"), - ), + 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"), @@ -2674,7 +2756,7 @@ def test_init(self, shared, unit, error, dtype): return - result = xr.Dataset(data_vars={"a": arr1, "b": arr2}) + actual = xr.Dataset(data_vars={"a": arr1, "b": arr2}) expected_units = { "a": a.units, @@ -2688,7 +2770,7 @@ def test_init(self, shared, unit, error, dtype): xr.Dataset(data_vars={"a": strip_units(arr1), "b": strip_units(arr2)}), expected_units, ) - assert_equal_with_units(result, expected) + assert_equal_with_units(actual, expected) @pytest.mark.parametrize( "func", (pytest.param(str, id="str"), pytest.param(repr, id="repr")) @@ -2749,12 +2831,7 @@ def test_repr(self, func, variant, dtype): reason="np.median does not work with dataset yet" ), ), - pytest.param( - function("sum"), - marks=pytest.mark.xfail( - reason="np.result_type not implemented by pint" - ), - ), + function("sum"), pytest.param( function("prod"), marks=pytest.mark.xfail(reason="not implemented by pint"), @@ -2764,9 +2841,7 @@ def test_repr(self, func, variant, dtype): function("cumsum"), pytest.param( function("cumprod"), - marks=pytest.mark.xfail( - reason="pint does not support cumprod on non-dimensionless yet" - ), + marks=pytest.mark.xfail(reason="fails within xarray"), ), pytest.param( method("all"), marks=pytest.mark.xfail(reason="not implemented by pint") @@ -2780,12 +2855,7 @@ def test_repr(self, func, variant, dtype): method("min"), method("mean"), method("median"), - pytest.param( - method("sum"), - marks=pytest.mark.xfail( - reason="np.result_type not implemented by pint" - ), - ), + method("sum"), pytest.param( method("prod"), marks=pytest.mark.xfail(reason="not implemented by pint"), @@ -2794,17 +2864,20 @@ def test_repr(self, func, variant, dtype): method("var"), method("cumsum"), pytest.param( - method("cumprod"), - marks=pytest.mark.xfail( - reason="pint does not support cumprod on non-dimensionless yet" - ), + method("cumprod"), marks=pytest.mark.xfail(reason="fails within xarray") ), ), ids=repr, ) def test_aggregation(self, func, dtype): - unit_a = unit_registry.Pa - unit_b = unit_registry.kg / unit_registry.m ** 3 + unit_a = ( + unit_registry.Pa if func.name != "cumprod" else unit_registry.dimensionless + ) + unit_b = ( + unit_registry.kg / unit_registry.m ** 3 + 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") @@ -2814,13 +2887,16 @@ def test_aggregation(self, func, dtype): ds = xr.Dataset(data_vars={"a": a, "b": b}, coords={"x": x, "y": y}) - result = func(ds) + actual = func(ds) expected = attach_units( func(strip_units(ds)), - {"a": array_extract_units(func(a)), "b": array_extract_units(func(b))}, + { + "a": extract_units(func(a)).get(None), + "b": extract_units(func(b)).get(None), + }, ) - assert_equal_with_units(result, expected) + assert_equal_with_units(actual, expected) @pytest.mark.parametrize("property", ("imag", "real")) def test_numpy_properties(self, property, dtype): @@ -2840,10 +2916,10 @@ def test_numpy_properties(self, property, dtype): ) units = extract_units(ds) - result = getattr(ds, property) + actual = getattr(ds, property) expected = attach_units(getattr(strip_units(ds), property), units) - assert_equal_with_units(result, expected) + assert_equal_with_units(actual, expected) @pytest.mark.parametrize( "func", @@ -2853,10 +2929,6 @@ def test_numpy_properties(self, property, dtype): method("argsort"), method("conjugate"), method("round"), - pytest.param( - method("rank", dim="x"), - marks=pytest.mark.xfail(reason="pint does not implement rank yet"), - ), ), ids=repr, ) @@ -2882,10 +2954,10 @@ def test_numpy_methods(self, func, dtype): "y": unit_registry.s, } - result = func(ds) + actual = func(ds) expected = attach_units(func(strip_units(ds)), units) - assert_equal_with_units(result, expected) + assert_equal_with_units(actual, expected) @pytest.mark.parametrize("func", (method("clip", min=3, max=8),), ids=repr) @pytest.mark.parametrize( @@ -2914,37 +2986,26 @@ def test_numpy_methods_with_args(self, func, unit, error, dtype): ) units = extract_units(ds) - def strip(value): - return ( - value.magnitude if isinstance(value, unit_registry.Quantity) else value - ) - - def convert(value, to): - if isinstance(value, unit_registry.Quantity) and value.check(to): - return value.to(to) - - return value - - scalar_types = (int, float) kwargs = { - key: (value * unit if isinstance(value, scalar_types) else value) + key: (value * unit if isinstance(value, (int, float)) else value) for key, value in func.kwargs.items() } - stripped_kwargs = { - key: strip(convert(value, data_unit)) for key, value in kwargs.items() - } - if error is not None: with pytest.raises(error): func(ds, **kwargs) return - result = func(ds, **kwargs) + stripped_kwargs = { + key: strip_units(convert_units(value, {None: data_unit})) + for key, value in kwargs.items() + } + + actual = func(ds, **kwargs) expected = attach_units(func(strip_units(ds), **stripped_kwargs), units) - assert_equal_with_units(result, expected) + assert_equal_with_units(actual, expected) @pytest.mark.parametrize( "func", (method("isnull"), method("notnull"), method("count")), ids=repr @@ -2987,9 +3048,9 @@ def test_missing_value_detection(self, func, dtype): ) expected = func(strip_units(ds)) - result = func(ds) + actual = func(ds) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.xfail(reason="ffill and bfill lose the unit") @pytest.mark.parametrize("func", (method("ffill"), method("bfill")), ids=repr) @@ -3017,40 +3078,35 @@ def test_missing_value_filling(self, func, dtype): func(strip_units(ds), dim="x"), {"a": unit_registry.degK, "b": unit_registry.Pa}, ) - result = func(ds, dim="x") + actual = func(ds, dim="x") - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="fillna drops the unit") @pytest.mark.parametrize( "unit,error", ( - pytest.param( - 1, - DimensionalityError, - id="no_unit", - marks=pytest.mark.xfail(reason="blocked by the failing `where`"), - ), + 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.cm, + None, + id="compatible_unit", + marks=pytest.mark.xfail( + reason="where converts the array, not the fill value" + ), + ), pytest.param(unit_registry.m, None, id="identical_unit"), ), ) @pytest.mark.parametrize( "fill_value", ( - pytest.param( - -1, - id="python scalar", - marks=pytest.mark.xfail( - reason="python scalar cannot be converted using astype()" - ), - ), - pytest.param(np.array(-1), id="numpy scalar"), - pytest.param(np.array([-1]), id="numpy array"), + pytest.param(-1, id="python_scalar"), + pytest.param(np.array(-1), id="numpy_scalar"), + pytest.param(np.array([-1]), id="numpy_array"), ), ) def test_fillna(self, fill_value, unit, error, dtype): @@ -3075,13 +3131,17 @@ def test_fillna(self, fill_value, unit, error, dtype): return - result = ds.fillna(value=fill_value * unit) + actual = ds.fillna(value=fill_value * unit) expected = attach_units( - strip_units(ds).fillna(value=fill_value), + strip_units(ds).fillna( + value=strip_units( + convert_units(fill_value * unit, {None: unit_registry.m}) + ) + ), {"a": unit_registry.m, "b": unit_registry.m}, ) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) def test_dropna(self, dtype): array1 = ( @@ -3105,11 +3165,10 @@ def test_dropna(self, dtype): strip_units(ds).dropna(dim="x"), {"a": unit_registry.degK, "b": unit_registry.Pa}, ) - result = ds.dropna(dim="x") + actual = ds.dropna(dim="x") - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="pint does not implement `numpy.isin`") @pytest.mark.parametrize( "unit", ( @@ -3154,36 +3213,12 @@ def test_isin(self, unit, dtype): ): expected.a[:] = False expected.b[:] = False - result = ds.isin(values) + actual = ds.isin(values) - assert_equal_with_units(result, expected) + assert_equal_with_units(actual, expected) @pytest.mark.parametrize( - "variant", - ( - pytest.param( - "masking", - marks=pytest.mark.xfail( - reason="np.result_type not implemented by quantity" - ), - ), - pytest.param( - "replacing_scalar", - marks=pytest.mark.xfail( - reason="python scalar not convertible using astype" - ), - ), - pytest.param( - "replacing_array", - marks=pytest.mark.xfail( - reason="replacing using an array drops the units" - ), - ), - pytest.param( - "dropping", - marks=pytest.mark.xfail(reason="nan not compatible with quantity"), - ), - ), + "variant", ("masking", "replacing_scalar", "replacing_array", "dropping") ) @pytest.mark.parametrize( "unit,error", @@ -3198,9 +3233,6 @@ def test_isin(self, unit, dtype): ), ) def test_where(self, variant, unit, error, dtype): - def _strip_units(mapping): - return {key: array_strip_units(value) for key, value in mapping.items()} - original_unit = unit_registry.m array1 = np.linspace(0, 1, 10).astype(dtype) * original_unit array2 = np.linspace(-1, 0, 10).astype(dtype) * original_unit @@ -3222,21 +3254,24 @@ def _strip_units(mapping): "dropping": {"cond": condition, "drop": True}, } kwargs = variant_kwargs.get(variant) - kwargs_without_units = _strip_units(kwargs) - if variant not in ("masking", "dropping") and error is not None: with pytest.raises(error): ds.where(**kwargs) return + kwargs_without_units = { + key: strip_units(convert_units(value, {None: original_unit})) + for key, value in kwargs.items() + } + expected = attach_units( strip_units(ds).where(**kwargs_without_units), {"a": original_unit, "b": original_unit}, ) - result = ds.where(**kwargs) + actual = ds.where(**kwargs) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.xfail(reason="interpolate strips units") def test_interpolate_na(self, dtype): @@ -3261,11 +3296,11 @@ def test_interpolate_na(self, dtype): strip_units(ds).interpolate_na(dim="x"), {"a": unit_registry.degK, "b": unit_registry.Pa}, ) - result = ds.interpolate_na(dim="x") + actual = ds.interpolate_na(dim="x") - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="uses Dataset.where, which currently fails") + @pytest.mark.xfail(reason="wrong argument order for `where`") @pytest.mark.parametrize( "unit,error", ( @@ -3281,11 +3316,11 @@ def test_interpolate_na(self, dtype): def test_combine_first(self, unit, error, dtype): array1 = ( np.array([1.4, np.nan, 2.3, np.nan, np.nan, 9.1]).astype(dtype) - * unit_registry.degK + * unit_registry.m ) array2 = ( np.array([4.3, 9.8, 7.5, np.nan, 8.2, np.nan]).astype(dtype) - * unit_registry.Pa + * unit_registry.m ) x = np.arange(len(array1)) ds = xr.Dataset( @@ -3312,12 +3347,16 @@ def test_combine_first(self, unit, error, dtype): return expected = attach_units( - strip_units(ds).combine_first(strip_units(other)), + 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}, ) - result = ds.combine_first(other) + actual = ds.combine_first(other) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "unit", @@ -3325,11 +3364,7 @@ def test_combine_first(self, unit, error, dtype): 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", - marks=pytest.mark.xfail(reason="identical does not check units yet"), - ), + pytest.param(unit_registry.cm, id="compatible_unit"), pytest.param(unit_registry.m, id="identical_unit"), ), ) @@ -3345,6 +3380,13 @@ def test_combine_first(self, unit, error, dtype): ) @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) + array1 = np.linspace(0, 5, 10).astype(dtype) array2 = np.linspace(-5, 0, 10).astype(dtype) @@ -3356,11 +3398,7 @@ def test_comparisons(self, func, variation, unit, dtype): x = coord * original_unit y = coord * original_unit - units = { - "data": (unit, original_unit, original_unit), - "dims": (original_unit, unit, original_unit), - "coords": (original_unit, original_unit, unit), - } + units = {"data": (unit, 1, 1), "dims": (1, unit, 1), "coords": (1, 1, unit)} data_unit, dim_unit, coord_unit = units.get(variation) ds = xr.Dataset( @@ -3371,36 +3409,27 @@ def test_comparisons(self, func, variation, unit, dtype): coords={"x": x, "y": ("x", y)}, ) - other = attach_units( - strip_units(ds), - { - "a": (data_unit, original_unit if quantity1.check(data_unit) else None), - "b": (data_unit, original_unit if quantity2.check(data_unit) else None), - "x": (dim_unit, original_unit if x.check(dim_unit) else None), - "y": (coord_unit, original_unit if y.check(coord_unit) else None), - }, - ) + 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, + } + other = attach_units(strip_units(convert_units(ds, other_units)), other_units) - # TODO: test dim coord once indexes leave units intact - # also, express this in terms of calls on the raw data array - # and then check the units - equal_arrays = ( - np.all(ds.a.data == other.a.data) - and np.all(ds.b.data == other.b.data) - and (np.all(x == other.x.data) or True) # dims can't be checked yet - and np.all(y == other.y.data) - ) - equal_units = ( - data_unit == original_unit - and coord_unit == original_unit - and dim_unit == original_unit - ) - expected = equal_arrays and (func.name != "identical" or equal_units) - result = func(ds, other) + units = extract_units(ds) + other_units = extract_units(other) + + equal_ds = all( + is_compatible(units[name], other_units[name]) for name in units.keys() + ) 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) + + actual = func(ds, other) - assert expected == result + assert expected == actual - @pytest.mark.xfail(reason="blocked by `where`") @pytest.mark.parametrize( "unit", ( @@ -3430,9 +3459,9 @@ def test_broadcast_like(self, unit, dtype): expected = attach_units( strip_units(ds1).broadcast_like(strip_units(ds2)), extract_units(ds1) ) - result = ds1.broadcast_like(ds2) + actual = ds1.broadcast_like(ds2) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "unit", @@ -3446,38 +3475,34 @@ def test_broadcast_like(self, unit, dtype): ) def test_broadcast_equals(self, unit, dtype): left_array1 = np.ones(shape=(2, 3), dtype=dtype) * unit_registry.m - left_array2 = np.zeros(shape=(2, 6), dtype=dtype) * unit_registry.m + left_array2 = np.zeros(shape=(3, 6), dtype=dtype) * unit_registry.m - right_array1 = array_attach_units( - np.ones(shape=(2,), dtype=dtype), - unit, - convert_from=unit_registry.m if left_array1.check(unit) else None, - ) - right_array2 = array_attach_units( - np.ones(shape=(2,), dtype=dtype), - unit, - convert_from=unit_registry.m if left_array2.check(unit) else None, - ) + right_array1 = np.ones(shape=(2,)) * unit + right_array2 = np.ones(shape=(3,)) * unit left = xr.Dataset( data_vars={ "a": xr.DataArray(data=left_array1, dims=("x", "y")), - "b": xr.DataArray(data=left_array2, dims=("x", "z")), + "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="x"), + "b": xr.DataArray(data=right_array2, dims="y"), } ) - expected = np.all(left_array1 == right_array1[:, None]) and np.all( - left_array2 == right_array2[:, None] - ) - result = left.broadcast_equals(right) + 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) + actual = left.broadcast_equals(right) - assert expected == result + assert expected == actual @pytest.mark.parametrize( "func", @@ -3510,11 +3535,10 @@ def test_stacking_stacked(self, func, dtype): expected = attach_units( func(strip_units(stacked)), {"a": unit_registry.m, "b": unit_registry.m} ) - result = func(stacked) + actual = func(stacked) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="tries to subscript scalar quantities") 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} @@ -3528,13 +3552,13 @@ def test_to_stacked_array(self, dtype): func = method("to_stacked_array", "z", variable_dim="y", sample_dims=["x"]) - result = func(ds).rename(None) + actual = func(ds).rename(None) expected = attach_units( func(strip_units(ds)).rename(None), {None: unit_registry.m, "y": unit_registry.s}, ) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", @@ -3543,12 +3567,10 @@ def test_to_stacked_array(self, dtype): method("stack", a=("x", "y")), method("set_index", x="x2"), pytest.param( - method("shift", x=2), marks=pytest.mark.xfail(reason="sets all to nan") - ), - pytest.param( - method("roll", x=2, roll_coords=False), - marks=pytest.mark.xfail(reason="strips units"), + method("shift", x=2), + marks=pytest.mark.xfail(reason="tries to concatenate nan arrays"), ), + method("roll", x=2, roll_coords=False), method("sortby", "x2"), ), ids=repr, @@ -3581,9 +3603,9 @@ def test_stacking_reordering(self, func, dtype): expected = attach_units( func(strip_units(ds)), {"a": unit_registry.Pa, "b": unit_registry.degK} ) - result = func(ds) + actual = func(ds) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.xfail(reason="indexes strip units") @pytest.mark.parametrize( @@ -3610,35 +3632,33 @@ def test_isel(self, indices, dtype): strip_units(ds).isel(x=indices), {"a": unit_registry.s, "b": unit_registry.Pa, "x": unit_registry.m}, ) - result = ds.isel(x=indices) + actual = ds.isel(x=indices) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail( - reason="xarray does not support duck arrays in dimension coordinates" - ) + @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( - "values", + "raw_values", ( - pytest.param(12, id="single_value"), + pytest.param(10, id="single_value"), pytest.param([10, 5, 13], id="list_of_values"), pytest.param(np.array([9, 3, 7, 12]), id="array_of_values"), ), ) @pytest.mark.parametrize( - "units,error", + "unit,error", ( 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.ms, KeyError, id="compatible_unit"), - pytest.param(unit_registry.s, None, id="same_unit"), + pytest.param(unit_registry.dm, KeyError, id="compatible_unit"), + pytest.param(unit_registry.m, None, id="identical_unit"), ), ) - def test_sel(self, values, units, error, dtype): + def test_sel(self, raw_values, unit, error, dtype): array1 = np.linspace(5, 10, 20).astype(dtype) * unit_registry.degK array2 = np.linspace(0, 5, 20).astype(dtype) * unit_registry.Pa - x = np.arange(len(array1)) * unit_registry.s + x = np.arange(len(array1)) * unit_registry.m ds = xr.Dataset( data_vars={ @@ -3648,46 +3668,46 @@ def test_sel(self, values, units, error, dtype): coords={"x": x}, ) - values_with_units = values * units + values = raw_values * unit - if error is not None: + if error is not None and not ( + isinstance(raw_values, (int, float)) and x.check(unit) + ): with pytest.raises(error): - ds.sel(x=values_with_units) + ds.sel(x=values) return expected = attach_units( - strip_units(ds).sel(x=values), - {"a": unit_registry.degK, "b": unit_registry.Pa, "x": unit_registry.s}, + strip_units(ds).sel(x=strip_units(convert_units(values, {None: x.units}))), + {"a": array1.units, "b": array2.units, "x": x.units}, ) - result = ds.sel(x=values_with_units) - assert_equal_with_units(expected, result) + actual = ds.sel(x=values) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail( - reason="xarray does not support duck arrays in dimension coordinates" - ) + @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( - "values", + "raw_values", ( - pytest.param(12, id="single value"), - pytest.param([10, 5, 13], id="list of multiple values"), - pytest.param(np.array([9, 3, 7, 12]), id="array of multiple values"), + pytest.param(10, id="single_value"), + pytest.param([10, 5, 13], id="list_of_values"), + pytest.param(np.array([9, 3, 7, 12]), id="array_of_values"), ), ) @pytest.mark.parametrize( - "units,error", + "unit,error", ( 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.ms, KeyError, id="compatible_unit"), - pytest.param(unit_registry.s, None, id="same_unit"), + pytest.param(unit_registry.dm, KeyError, id="compatible_unit"), + pytest.param(unit_registry.m, None, id="identical_unit"), ), ) - def test_loc(self, values, units, error, dtype): + def test_drop_sel(self, raw_values, unit, error, dtype): array1 = np.linspace(5, 10, 20).astype(dtype) * unit_registry.degK array2 = np.linspace(0, 5, 20).astype(dtype) * unit_registry.Pa - x = np.arange(len(array1)) * unit_registry.s + x = np.arange(len(array1)) * unit_registry.m ds = xr.Dataset( data_vars={ @@ -3697,36 +3717,76 @@ def test_loc(self, values, units, error, dtype): coords={"x": x}, ) - values_with_units = values * units + values = raw_values * unit - if error is not None: + if error is not None and not ( + isinstance(raw_values, (int, float)) and x.check(unit) + ): with pytest.raises(error): - ds.loc[{"x": values_with_units}] + ds.drop_sel(x=values) return expected = attach_units( - strip_units(ds).loc[{"x": values}], - {"a": unit_registry.degK, "b": unit_registry.Pa, "x": unit_registry.s}, + strip_units(ds).drop_sel( + x=strip_units(convert_units(values, {None: x.units})) + ), + extract_units(ds), ) - result = ds.loc[{"x": values_with_units}] - assert_equal_with_units(expected, result) + actual = ds.drop_sel(x=values) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail( - reason="indexes strip units and head / tail / thin only support integers" + @pytest.mark.xfail(reason="indexes don't support units") + @pytest.mark.parametrize( + "raw_values", + ( + pytest.param(10, id="single_value"), + pytest.param([10, 5, 13], id="list_of_values"), + pytest.param(np.array([9, 3, 7, 12]), id="array_of_values"), + ), ) @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(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.m, None, id="identical_unit"), ), ) + def test_loc(self, raw_values, unit, error, dtype): + array1 = np.linspace(5, 10, 20).astype(dtype) * unit_registry.degK + array2 = np.linspace(0, 5, 20).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}, + ) + + values = raw_values * unit + + if error is not None and not ( + isinstance(raw_values, (int, float)) and x.check(unit) + ): + with pytest.raises(error): + ds.loc[{"x": values}] + + return + + expected = attach_units( + strip_units(ds).loc[ + {"x": strip_units(convert_units(values, {None: x.units}))} + ], + {"a": array1.units, "b": array2.units, "x": x.units}, + ) + actual = ds.loc[{"x": values}] + assert_equal_with_units(expected, actual) + @pytest.mark.parametrize( "func", ( @@ -3736,7 +3796,7 @@ def test_loc(self, values, units, error, dtype): ), ids=repr, ) - def test_head_tail_thin(self, func, unit, error, dtype): + 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 @@ -3754,18 +3814,10 @@ def test_head_tail_thin(self, func, unit, error, dtype): coords=coords, ) - kwargs = {name: value * unit for name, value in func.kwargs.items()} - - if error is not None: - with pytest.raises(error): - func(ds, **kwargs) - - return - expected = attach_units(func(strip_units(ds)), extract_units(ds)) - result = func(ds, **kwargs) + actual = func(ds) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "shape", @@ -3802,15 +3854,15 @@ def test_squeeze(self, shape, dtype): expected = attach_units(strip_units(ds).squeeze(), units) - result = ds.squeeze() - assert_equal_with_units(result, expected) + actual = ds.squeeze() + assert_equal_with_units(actual, expected) # 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) - result = ds.squeeze(dim=name) - assert_equal_with_units(result, expected) + actual = ds.squeeze(dim=name) + assert_equal_with_units(actual, expected) @pytest.mark.xfail(reason="ignores units") @pytest.mark.parametrize( @@ -3851,12 +3903,14 @@ def test_interp(self, unit, error): return + units = extract_units(ds) expected = attach_units( - strip_units(ds).interp(x=strip_units(new_coords)), extract_units(ds) + strip_units(ds).interp(x=strip_units(convert_units(new_coords, units))), + units, ) - result = ds.interp(x=new_coords) + actual = ds.interp(x=new_coords) - assert_equal_with_units(result, expected) + assert_equal_with_units(actual, expected) @pytest.mark.xfail(reason="ignores units") @pytest.mark.parametrize( @@ -3911,16 +3965,15 @@ def test_interp_like(self, unit, error, dtype): return + units = extract_units(ds) expected = attach_units( - strip_units(ds).interp_like(strip_units(other)), extract_units(ds) + strip_units(ds).interp_like(strip_units(convert_units(other, units))), units ) - result = ds.interp_like(other) + actual = ds.interp_like(other) - assert_equal_with_units(result, expected) + assert_equal_with_units(actual, expected) - @pytest.mark.xfail( - reason="pint does not implement np.result_type in __array_function__ yet" - ) + @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( "unit,error", ( @@ -3933,9 +3986,13 @@ def test_interp_like(self, unit, error, dtype): pytest.param(unit_registry.m, None, id="identical_unit"), ), ) - def test_reindex(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 + 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, @@ -3955,20 +4012,21 @@ def test_reindex(self, unit, error): if error is not None: with pytest.raises(error): - ds.interp(x=new_coords) + ds.reindex(x=new_coords) return expected = attach_units( - strip_units(ds).reindex(x=strip_units(new_coords)), extract_units(ds) + strip_units(ds).reindex( + x=strip_units(convert_units(new_coords, {None: coords["x"].units})) + ), + extract_units(ds), ) - result = ds.reindex(x=new_coords) + actual = ds.reindex(x=new_coords) - assert_equal_with_units(result, expected) + assert_equal_with_units(actual, expected) - @pytest.mark.xfail( - reason="pint does not implement np.result_type in __array_function__ yet" - ) + @pytest.mark.xfail(reason="indexes don't support units") @pytest.mark.parametrize( "unit,error", ( @@ -4021,12 +4079,14 @@ def test_reindex_like(self, unit, error, dtype): return + units = extract_units(ds) expected = attach_units( - strip_units(ds).reindex_like(strip_units(other)), extract_units(ds) + strip_units(ds).reindex_like(strip_units(convert_units(other, units))), + units, ) - result = ds.reindex_like(other) + actual = ds.reindex_like(other) - assert_equal_with_units(result, expected) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", @@ -4034,20 +4094,9 @@ def test_reindex_like(self, unit, error, dtype): method("diff", dim="x"), method("differentiate", coord="x"), method("integrate", coord="x"), - pytest.param( - method("quantile", q=[0.25, 0.75]), - marks=pytest.mark.xfail( - reason="pint does not implement nanpercentile yet" - ), - ), - pytest.param( - method("reduce", func=np.sum, dim="x"), - marks=pytest.mark.xfail(reason="strips units"), - ), - pytest.param( - method("map", np.fabs), - marks=pytest.mark.xfail(reason="fabs strips units"), - ), + method("quantile", q=[0.25, 0.75]), + method("reduce", func=np.sum, dim="x"), + method("map", np.fabs), ), ids=repr, ) @@ -4073,27 +4122,22 @@ def test_computation(self, func, dtype): units = extract_units(ds) expected = attach_units(func(strip_units(ds)), units) - result = func(ds) + actual = func(ds) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", ( - pytest.param( - method("groupby", "x"), marks=pytest.mark.xfail(reason="strips units") - ), - pytest.param( - method("groupby_bins", "x", bins=4), - marks=pytest.mark.xfail(reason="strips units"), - ), + method("groupby", "x"), + method("groupby_bins", "x", bins=4), 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="strips units"), + marks=pytest.mark.xfail(reason="uses numbagg which strips units"), ), ), ids=repr, @@ -4122,11 +4166,10 @@ def test_computation_objects(self, func, dtype): args = [] if func.name != "groupby" else ["y"] reduce_func = method("mean", *args) expected = attach_units(reduce_func(func(strip_units(ds))), units) - result = reduce_func(func(ds)) + actual = reduce_func(func(ds)) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="strips units") def test_resample(self, dtype): array1 = ( np.linspace(-5, 5, 10 * 5).reshape(10, 5).astype(dtype) * unit_registry.degK @@ -4150,29 +4193,18 @@ def test_resample(self, dtype): func = method("resample", time="6m") expected = attach_units(func(strip_units(ds)).mean(), units) - result = func(ds).mean() + actual = func(ds).mean() - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", ( - pytest.param( - method("assign", c=lambda ds: 10 * ds.b), - marks=pytest.mark.xfail(reason="strips units"), - ), - pytest.param( - method("assign_coords", v=("x", np.arange(10) * unit_registry.s)), - marks=pytest.mark.xfail(reason="strips units"), - ), - pytest.param(method("first")), - pytest.param(method("last")), - pytest.param( - method("quantile", q=[0.25, 0.5, 0.75], dim="x"), - marks=pytest.mark.xfail( - reason="dataset groupby does not implement quantile" - ), - ), + method("assign", c=lambda ds: 10 * ds.b), + method("assign_coords", v=("x", np.arange(10) * unit_registry.s)), + method("first"), + method("last"), + method("quantile", q=[0.25, 0.5, 0.75], dim="x"), ), ids=repr, ) @@ -4204,9 +4236,9 @@ def test_grouped_operations(self, func, dtype): expected = attach_units( func(strip_units(ds).groupby("y"), **stripped_kwargs), units ) - result = func(ds.groupby("y")) + actual = func(ds.groupby("y")) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) @pytest.mark.parametrize( "func", @@ -4220,7 +4252,7 @@ def test_grouped_operations(self, func, dtype): 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("drop_sel", labels="x"), + method("drop_vars", "x"), method("drop_dims", "z"), method("set_coords", names="c"), method("reset_coords", names="x2"), @@ -4252,26 +4284,25 @@ def test_content_manipulation(self, func, dtype): }, coords={"x": x, "y": y, "z": z, "x2": ("x", x2)}, ) - units = extract_units(ds) - units.update( - { + 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, - } - ) + }, + } stripped_kwargs = { key: strip_units(value) for key, value in func.kwargs.items() } expected = attach_units(func(strip_units(ds), **stripped_kwargs), units) - result = func(ds) + actual = func(ds) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual) - @pytest.mark.xfail(reason="blocked by reindex") @pytest.mark.parametrize( "unit,error", ( @@ -4284,7 +4315,16 @@ def test_content_manipulation(self, func, dtype): pytest.param(unit_registry.m, None, id="identical_unit"), ), ) - @pytest.mark.parametrize("variant", ("data", "dims", "coords")) + @pytest.mark.parametrize( + "variant", + ( + "data", + pytest.param( + "dims", marks=pytest.mark.xfail(reason="indexes don't support units") + ), + "coords", + ), + ) def test_merge(self, variant, unit, error, dtype): original_data_unit = unit_registry.m original_dim_unit = unit_registry.m @@ -4325,6 +4365,6 @@ def test_merge(self, variant, unit, error, dtype): converted = convert_units(right, units) expected = attach_units(strip_units(left).merge(strip_units(converted)), units) - result = left.merge(right) + actual = left.merge(right) - assert_equal_with_units(expected, result) + assert_equal_with_units(expected, actual)