From b9baa9c8c93ed6c7d7e20460b438b7142ce5906c Mon Sep 17 00:00:00 2001 From: Keewis Date: Sat, 26 Mar 2022 00:45:31 +0100 Subject: [PATCH 01/56] add a `PintMetaIndex` that for now can only `sel` --- pint_xarray/index.py | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) create mode 100644 pint_xarray/index.py diff --git a/pint_xarray/index.py b/pint_xarray/index.py new file mode 100644 index 00000000..99f1bfb8 --- /dev/null +++ b/pint_xarray/index.py @@ -0,0 +1,30 @@ +from xarray.core.indexes import Index + +from . import conversion + + +class PintMetaIndex(Index): + # TODO: inherit from MetaIndex once that exists + def __init__(self, *, index, units): + """create a unit-aware MetaIndex + + Parameters + ---------- + index : xarray.Index + The wrapped index object. + units : mapping of hashable to unit-like + The units of the indexed coordinates + """ + self.index = index + self.units = units + + # don't need `from_variables`: we're always *wrapping* an existing index + + def sel(self, labels): + converted_labels = conversion.convert_indexer_units(labels, self.units) + stripped_labels = { + name: conversion.strip_indexer_units(indexer) + for name, indexer in converted_labels.items() + } + + return self.index.sel(stripped_labels) From 89c5e2a0d69cbac813af624ff46ba12ea7591c50 Mon Sep 17 00:00:00 2001 From: Keewis Date: Sun, 27 Mar 2022 22:22:41 +0200 Subject: [PATCH 02/56] add a function to compare indexers --- pint_xarray/tests/test_conversion.py | 32 ++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/pint_xarray/tests/test_conversion.py b/pint_xarray/tests/test_conversion.py index 6386b60c..2078c5ed 100644 --- a/pint_xarray/tests/test_conversion.py +++ b/pint_xarray/tests/test_conversion.py @@ -1,3 +1,5 @@ +from textwrap import indent + import numpy as np import pint import pytest @@ -48,6 +50,36 @@ def strip_quantity(q): return q +def assert_indexers_equal(first, second): + __tracebackhide__ = True + # same keys + assert first.keys() == second.keys(), "different keys" + + errors = {} + for name in first: + first_value = first[name] + second_value = second[name] + + try: + if isinstance(first_value, DataArray): + assert_identical(first_value, second_value) + else: + assert_array_equal(first_value, second_value) + except AssertionError as e: + errors[name] = e + + if errors: + message = "\n".join( + ["indexers are not equal:"] + + [ + f" - {name}:\n{indent(str(error), ' ' * 4)}" + for name, error in errors.items() + ] + ) + + raise AssertionError(message) + + class TestArrayFunctions: @pytest.mark.parametrize( ["unit", "data", "expected", "match"], From 17e9aec58e672a0d112e85e019fccda16c2f4e6a Mon Sep 17 00:00:00 2001 From: Keewis Date: Sun, 27 Mar 2022 22:24:18 +0200 Subject: [PATCH 03/56] expect indexer dicts for strip_indexer_units --- pint_xarray/accessors.py | 55 ++++++---------------------- pint_xarray/conversion.py | 29 ++++++++------- pint_xarray/tests/test_conversion.py | 54 ++++++++++++++++----------- 3 files changed, 59 insertions(+), 79 deletions(-) diff --git a/pint_xarray/accessors.py b/pint_xarray/accessors.py index b7de2435..56f6a822 100644 --- a/pint_xarray/accessors.py +++ b/pint_xarray/accessors.py @@ -173,10 +173,7 @@ def __getitem__(self, indexers): raise KeyError(*e.args) from e # index - stripped_indexers = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in indexers.items() - } + stripped_indexers = conversion.strip_indexer_units(indexers) return converted.loc[stripped_indexers] @@ -204,10 +201,7 @@ def __getitem__(self, indexers): raise KeyError(*e.args) from e # index - stripped_indexers = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in indexers.items() - } + stripped_indexers = conversion.strip_indexer_units(indexers) return converted.loc[stripped_indexers] def __setitem__(self, indexers, values): @@ -227,10 +221,7 @@ def __setitem__(self, indexers, values): raise KeyError(*e.args) from e # index - stripped_indexers = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in converted.items() - } + stripped_indexers = conversion.strip_indexer_units(converted) self.da.loc[stripped_indexers] = values @@ -632,10 +623,7 @@ def reindex( converted = conversion.convert_units(self.da, indexer_units) # index - stripped_indexers = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in indexers.items() - } + stripped_indexers = conversion.strip_indexer_units(indexers) indexed = converted.reindex( stripped_indexers, method=method, @@ -714,10 +702,7 @@ def interp( stripped = conversion.strip_units(converted) # index - stripped_indexers = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in indexers.items() - } + stripped_indexers = conversion.strip_indexer_units(indexers) interpolated = stripped.interp( stripped_indexers, method=method, @@ -790,10 +775,7 @@ def sel( raise KeyError(*e.args) from e # index - stripped_indexers = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in indexers.items() - } + stripped_indexers = conversion.strip_indexer_units(indexers) indexed = converted.sel( stripped_indexers, method=method, @@ -843,10 +825,7 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs): raise KeyError(*e.args) from e # index - stripped_indexers = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in converted_indexers.items() - } + stripped_indexers = conversion.strip_indexer_units(converted_indexers) indexed = self.da.drop_sel( stripped_indexers, errors=errors, @@ -1355,10 +1334,7 @@ def reindex( converted = conversion.convert_units(self.ds, indexer_units) # index - stripped_indexers = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in indexers.items() - } + stripped_indexers = conversion.strip_indexer_units(indexers) indexed = converted.reindex( stripped_indexers, method=method, @@ -1437,10 +1413,7 @@ def interp( stripped = conversion.strip_units(converted) # index - stripped_indexers = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in indexers.items() - } + stripped_indexers = conversion.strip_indexer_units(indexers) interpolated = stripped.interp( stripped_indexers, method=method, @@ -1513,10 +1486,7 @@ def sel( raise KeyError(*e.args) from e # index - stripped_indexers = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in indexers.items() - } + stripped_indexers = conversion.strip_indexer_units(indexers) indexed = converted.sel( stripped_indexers, method=method, @@ -1568,10 +1538,7 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs): raise KeyError(*e.args) from e # index - stripped_indexers = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in converted_indexers.items() - } + stripped_indexers = conversion.strip_indexer_units(converted_indexers) indexed = self.ds.drop_sel( stripped_indexers, errors=errors, diff --git a/pint_xarray/conversion.py b/pint_xarray/conversion.py index 1940e492..e3d3ceb6 100644 --- a/pint_xarray/conversion.py +++ b/pint_xarray/conversion.py @@ -390,16 +390,19 @@ def extract_indexer_units(indexer): return array_extract_units(indexer) -def strip_indexer_units(indexer): - if isinstance(indexer, slice): - return slice( - array_strip_units(indexer.start), - array_strip_units(indexer.stop), - array_strip_units(indexer.step), - ) - elif isinstance(indexer, DataArray): - return strip_units(indexer) - elif isinstance(indexer, Variable): - return strip_units_variable(indexer) - else: - return array_strip_units(indexer) +def strip_indexer_units(indexers): + def strip(indexer): + if isinstance(indexer, slice): + return slice( + array_strip_units(indexer.start), + array_strip_units(indexer.stop), + array_strip_units(indexer.step), + ) + elif isinstance(indexer, DataArray): + return strip_units(indexer) + elif isinstance(indexer, Variable): + return strip_units_variable(indexer) + else: + return array_strip_units(indexer) + + return {name: strip(indexer) for name, indexer in indexers.items()} diff --git a/pint_xarray/tests/test_conversion.py b/pint_xarray/tests/test_conversion.py index 2078c5ed..a49790f5 100644 --- a/pint_xarray/tests/test_conversion.py +++ b/pint_xarray/tests/test_conversion.py @@ -760,42 +760,52 @@ def test_extract_indexer_units(self, indexer, expected): assert actual == expected @pytest.mark.parametrize( - ["indexer", "expected"], + ["indexers", "expected"], ( - pytest.param(1, 1, id="scalar-no units"), - pytest.param(Quantity(1, "m"), 1, id="scalar-units"), - pytest.param(np.array([1, 2]), np.array([1, 2]), id="array-no units"), - pytest.param(Quantity([1, 2], "s"), np.array([1, 2]), id="array-units"), + pytest.param({"x": 1}, {"x": 1}, id="scalar-no units"), + pytest.param({"x": Quantity(1, "m")}, {"x": 1}, id="scalar-units"), pytest.param( - Variable("x", [1, 2]), Variable("x", [1, 2]), id="Variable-no units" + {"x": np.array([1, 2])}, + {"x": np.array([1, 2])}, + id="array-no units", ), pytest.param( - Variable("x", Quantity([1, 2], "m")), - Variable("x", [1, 2]), + {"x": Quantity([1, 2], "s")}, {"x": np.array([1, 2])}, id="array-units" + ), + pytest.param( + {"x": Variable("x", [1, 2])}, + {"x": Variable("x", [1, 2])}, + id="Variable-no units", + ), + pytest.param( + {"x": Variable("x", Quantity([1, 2], "m"))}, + {"x": Variable("x", [1, 2])}, id="Variable-units", ), pytest.param( - DataArray([1, 2], dims="x"), - DataArray([1, 2], dims="x"), + {"x": DataArray([1, 2], dims="x")}, + {"x": DataArray([1, 2], dims="x")}, id="DataArray-no units", ), pytest.param( - DataArray(Quantity([1, 2], "s"), dims="x"), - DataArray([1, 2], dims="x"), + {"x": DataArray(Quantity([1, 2], "s"), dims="x")}, + {"x": DataArray([1, 2], dims="x")}, id="DataArray-units", ), - pytest.param(slice(None), slice(None), id="empty slice-no units"), - pytest.param(slice(1, None), slice(1, None), id="slice-no units"), pytest.param( - slice(Quantity(1, "m"), Quantity(2, "m")), - slice(1, 2), + {"x": slice(None)}, {"x": slice(None)}, id="empty slice-no units" + ), + pytest.param( + {"x": slice(1, None)}, {"x": slice(1, None)}, id="slice-no units" + ), + pytest.param( + {"x": slice(Quantity(1, "m"), Quantity(2, "m"))}, + {"x": slice(1, 2)}, id="slice-units", ), ), ) - def test_strip_indexer_units(self, indexer, expected): - actual = conversion.strip_indexer_units(indexer) - if isinstance(indexer, DataArray): - assert_identical(actual, expected) - else: - assert_array_equal(actual, expected) + def test_strip_indexer_units(self, indexers, expected): + actual = conversion.strip_indexer_units(indexers) + + assert_indexers_equal(actual, expected) From 30a1d80ee0a02431adfd8fa6f5f9707e632f1f54 Mon Sep 17 00:00:00 2001 From: Keewis Date: Sun, 27 Mar 2022 22:36:21 +0200 Subject: [PATCH 04/56] move the indexer comparison function to the utils --- pint_xarray/tests/test_conversion.py | 33 +--------------------------- pint_xarray/tests/utils.py | 28 +++++++++++++++++++++++ 2 files changed, 29 insertions(+), 32 deletions(-) diff --git a/pint_xarray/tests/test_conversion.py b/pint_xarray/tests/test_conversion.py index a49790f5..4a55f7ae 100644 --- a/pint_xarray/tests/test_conversion.py +++ b/pint_xarray/tests/test_conversion.py @@ -1,5 +1,3 @@ -from textwrap import indent - import numpy as np import pint import pytest @@ -13,6 +11,7 @@ assert_identical, assert_indexer_equal, assert_indexer_units_equal, + assert_indexers_equal, ) unit_registry = pint.UnitRegistry() @@ -50,36 +49,6 @@ def strip_quantity(q): return q -def assert_indexers_equal(first, second): - __tracebackhide__ = True - # same keys - assert first.keys() == second.keys(), "different keys" - - errors = {} - for name in first: - first_value = first[name] - second_value = second[name] - - try: - if isinstance(first_value, DataArray): - assert_identical(first_value, second_value) - else: - assert_array_equal(first_value, second_value) - except AssertionError as e: - errors[name] = e - - if errors: - message = "\n".join( - ["indexers are not equal:"] - + [ - f" - {name}:\n{indent(str(error), ' ' * 4)}" - for name, error in errors.items() - ] - ) - - raise AssertionError(message) - - class TestArrayFunctions: @pytest.mark.parametrize( ["unit", "data", "expected", "match"], diff --git a/pint_xarray/tests/utils.py b/pint_xarray/tests/utils.py index eb18b179..3559112c 100644 --- a/pint_xarray/tests/utils.py +++ b/pint_xarray/tests/utils.py @@ -1,5 +1,6 @@ import re from contextlib import contextmanager +from textwrap import indent import numpy as np import pytest @@ -97,6 +98,33 @@ def assert_indexer_equal(a, b): assert a_ == b_, f"different values: {a_!r} ←→ {b_!r}" +def assert_indexers_equal(first, second): + __tracebackhide__ = True + # same keys + assert first.keys() == second.keys(), "different keys" + + errors = {} + for name in first: + first_value = first[name] + second_value = second[name] + + try: + assert_indexer_equal(first_value, second_value) + except AssertionError as e: + errors[name] = e + + if errors: + message = "\n".join( + ["indexers are not equal:"] + + [ + f" - {name}:\n{indent(str(error), ' ' * 4)}" + for name, error in errors.items() + ] + ) + + raise AssertionError(message) + + def assert_indexer_units_equal(a, b): __tracebackhide__ = True From 34caf09e530c5c1dd947d48c47605fe53557434d Mon Sep 17 00:00:00 2001 From: Keewis Date: Sun, 27 Mar 2022 22:59:25 +0200 Subject: [PATCH 05/56] change extract_indexer_units to expect a dict --- pint_xarray/conversion.py | 17 +++++---- pint_xarray/tests/test_conversion.py | 56 ++++++++++++++++------------ 2 files changed, 42 insertions(+), 31 deletions(-) diff --git a/pint_xarray/conversion.py b/pint_xarray/conversion.py index e3d3ceb6..8005c179 100644 --- a/pint_xarray/conversion.py +++ b/pint_xarray/conversion.py @@ -381,13 +381,16 @@ def convert(indexer, units): return converted -def extract_indexer_units(indexer): - if isinstance(indexer, slice): - return slice_extract_units(indexer) - elif isinstance(indexer, (DataArray, Variable)): - return array_extract_units(indexer.data) - else: - return array_extract_units(indexer) +def extract_indexer_units(indexers): + def extract(indexer): + if isinstance(indexer, slice): + return slice_extract_units(indexer) + elif isinstance(indexer, (DataArray, Variable)): + return array_extract_units(indexer.data) + else: + return array_extract_units(indexer) + + return {name: extract(indexer) for name, indexer in indexers.items()} def strip_indexer_units(indexers): diff --git a/pint_xarray/tests/test_conversion.py b/pint_xarray/tests/test_conversion.py index 4a55f7ae..35cdd569 100644 --- a/pint_xarray/tests/test_conversion.py +++ b/pint_xarray/tests/test_conversion.py @@ -675,57 +675,65 @@ def test_convert_indexer_units(self, indexers, units, expected, error, match): assert_indexer_units_equal(actual["x"], expected["x"]) @pytest.mark.parametrize( - ["indexer", "expected"], + ["indexers", "expected"], ( - pytest.param(1, None, id="scalar-no units"), - pytest.param(Quantity(1, "m"), Unit("m"), id="scalar-units"), - pytest.param(np.array([1, 2]), None, id="array-no units"), - pytest.param(Quantity([1, 2], "s"), Unit("s"), id="array-units"), - pytest.param(Variable("x", [1, 2]), None, id="Variable-no units"), + pytest.param({"x": 1}, {"x": None}, id="scalar-no units"), + pytest.param({"x": Quantity(1, "m")}, {"x": Unit("m")}, id="scalar-units"), + pytest.param({"x": np.array([1, 2])}, {"x": None}, id="array-no units"), + pytest.param( + {"x": Quantity([1, 2], "s")}, {"x": Unit("s")}, id="array-units" + ), + pytest.param( + {"x": Variable("x", [1, 2])}, {"x": None}, id="Variable-no units" + ), + pytest.param( + {"x": Variable("x", Quantity([1, 2], "m"))}, + {"x": Unit("m")}, + id="Variable-units", + ), pytest.param( - Variable("x", Quantity([1, 2], "m")), Unit("m"), id="Variable-units" + {"x": DataArray([1, 2], dims="x")}, {"x": None}, id="DataArray-no units" ), - pytest.param(DataArray([1, 2], dims="x"), None, id="DataArray-no units"), pytest.param( - DataArray(Quantity([1, 2], "s"), dims="x"), - Unit("s"), + {"x": DataArray(Quantity([1, 2], "s"), dims="x")}, + {"x": Unit("s")}, id="DataArray-units", ), - pytest.param(slice(None), None, id="empty slice-no units"), - pytest.param(slice(1, None), None, id="slice-no units"), + pytest.param({"x": slice(None)}, {"x": None}, id="empty slice-no units"), + pytest.param({"x": slice(1, None)}, {"x": None}, id="slice-no units"), pytest.param( - slice(Quantity(1, "m"), Quantity(2, "m")), - Unit("m"), + {"x": slice(Quantity(1, "m"), Quantity(2, "m"))}, + {"x": Unit("m")}, id="slice-identical units", ), pytest.param( - slice(Quantity(1, "m"), Quantity(2000, "mm")), - Unit("m"), + {"x": slice(Quantity(1, "m"), Quantity(2000, "mm"))}, + {"x": Unit("m")}, id="slice-compatible units", ), pytest.param( - slice(Quantity(1, "m"), Quantity(2, "ms")), + {"x": slice(Quantity(1, "m"), Quantity(2, "ms"))}, ValueError, id="slice-incompatible units", ), pytest.param( - slice(1, Quantity(2, "ms")), + {"x": slice(1, Quantity(2, "ms"))}, ValueError, id="slice-incompatible units-mixed", ), pytest.param( - slice(1, Quantity(2, "rad")), - Unit("rad"), + {"x": slice(1, Quantity(2, "rad"))}, + {"x": Unit("rad")}, id="slice-incompatible units-mixed-dimensionless", ), ), ) - def test_extract_indexer_units(self, indexer, expected): - if expected is not None and not isinstance(expected, Unit): + def test_extract_indexer_units(self, indexers, expected): + if isinstance(expected, type) and issubclass(expected, Exception): with pytest.raises(expected): - conversion.extract_indexer_units(indexer) + conversion.extract_indexer_units(indexers) else: - actual = conversion.extract_indexer_units(indexer) + actual = conversion.extract_indexer_units(indexers) assert actual == expected @pytest.mark.parametrize( From 05aa5a6a995f72d79f077673a8c96b310527937c Mon Sep 17 00:00:00 2001 From: Keewis Date: Sun, 27 Mar 2022 23:19:55 +0200 Subject: [PATCH 06/56] fix a few calls to extract_indexer_units --- pint_xarray/accessors.py | 35 +++++++++++----------------- pint_xarray/tests/test_conversion.py | 5 ++-- 2 files changed, 16 insertions(+), 24 deletions(-) diff --git a/pint_xarray/accessors.py b/pint_xarray/accessors.py index 56f6a822..40e08049 100644 --- a/pint_xarray/accessors.py +++ b/pint_xarray/accessors.py @@ -160,10 +160,9 @@ def __getitem__(self, indexers): raise NotImplementedError("pandas-style indexing is not supported, yet") dims = self.ds.dims + indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { - name: conversion.extract_indexer_units(indexer) - for name, indexer in indexers.items() - if name in dims + name: indexer for name, indexer in indexer_units.items() if name in dims } # convert the indexes to the indexer's units @@ -188,10 +187,9 @@ def __getitem__(self, indexers): raise NotImplementedError("pandas-style indexing is not supported, yet") dims = self.da.dims + indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { - name: conversion.extract_indexer_units(indexer) - for name, indexer in indexers.items() - if name in dims + name: indexer for name, indexer in indexer_units.items() if name in dims } # convert the indexes to the indexer's units @@ -610,10 +608,9 @@ def reindex( indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "reindex") dims = self.da.dims + indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { - name: conversion.extract_indexer_units(indexer) - for name, indexer in indexers.items() - if name in dims + name: indexer for name, indexer in indexer_units.items() if name in dims } # TODO: handle tolerance @@ -760,10 +757,9 @@ def sel( indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel") dims = self.da.dims + indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { - name: conversion.extract_indexer_units(indexer) - for name, indexer in indexers.items() - if name in dims + name: indexer for name, indexer in indexer_units.items() if name in dims } # TODO: handle tolerance @@ -1321,10 +1317,9 @@ def reindex( indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "reindex") dims = self.ds.dims + indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { - name: conversion.extract_indexer_units(indexer) - for name, indexer in indexers.items() - if name in dims + name: indexer for name, indexer in indexer_units.items() if name in dims } # TODO: handle tolerance @@ -1401,10 +1396,9 @@ def interp( indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp") dims = self.ds.dims + indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { - name: conversion.extract_indexer_units(indexer) - for name, indexer in indexers.items() - if name in dims + name: indexer for name, indexer in indexer_units.items() if name in dims } # convert the indexes to the indexer's units @@ -1471,10 +1465,9 @@ def sel( indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel") dims = self.ds.dims + indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { - name: conversion.extract_indexer_units(indexer) - for name, indexer in indexers.items() - if name in dims + name: indexer for name, indexer in indexer_units.items() if name in dims } # TODO: handle tolerance diff --git a/pint_xarray/tests/test_conversion.py b/pint_xarray/tests/test_conversion.py index 35cdd569..f636042b 100644 --- a/pint_xarray/tests/test_conversion.py +++ b/pint_xarray/tests/test_conversion.py @@ -9,7 +9,6 @@ assert_array_equal, assert_array_units_equal, assert_identical, - assert_indexer_equal, assert_indexer_units_equal, assert_indexers_equal, ) @@ -671,8 +670,8 @@ def test_convert_indexer_units(self, indexers, units, expected, error, match): conversion.convert_indexer_units(indexers, units) else: actual = conversion.convert_indexer_units(indexers, units) - assert_indexer_equal(actual["x"], expected["x"]) - assert_indexer_units_equal(actual["x"], expected["x"]) + assert_indexers_equal(actual, expected) + assert_indexer_units_equal(actual, expected) @pytest.mark.parametrize( ["indexers", "expected"], From 2e0e5bd9890cbb7a5df830437300264f8686ae0c Mon Sep 17 00:00:00 2001 From: Keewis Date: Sun, 27 Mar 2022 23:27:45 +0200 Subject: [PATCH 07/56] one more call --- pint_xarray/accessors.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pint_xarray/accessors.py b/pint_xarray/accessors.py index 40e08049..3a34aa74 100644 --- a/pint_xarray/accessors.py +++ b/pint_xarray/accessors.py @@ -687,10 +687,9 @@ def interp( indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp") dims = self.da.dims + indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { - name: conversion.extract_indexer_units(indexer) - for name, indexer in indexers.items() - if name in dims + name: indexer for name, indexer in indexer_units.items() if name in dims } # convert the indexes to the indexer's units From a049d038ec9c9e585e17ec4da2be8d66c6d44406 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 13 Sep 2023 11:00:05 +0200 Subject: [PATCH 08/56] implement `create_variables` and `from_variables` Co-authored-by: Benoit Bovy --- pint_xarray/index.py | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/pint_xarray/index.py b/pint_xarray/index.py index 99f1bfb8..b234e783 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -1,4 +1,5 @@ -from xarray.core.indexes import Index +from xarray import Variable +from xarray.core.indexes import Index, PandasIndex from . import conversion @@ -18,7 +19,22 @@ def __init__(self, *, index, units): self.index = index self.units = units - # don't need `from_variables`: we're always *wrapping* an existing index + def create_variables(self, variables=None): + index_vars = self.index.create_variables(variables) + + index_vars_units = {} + for name, var in index_vars.items(): + data = conversion.array_attach_units(var.data, self.units[name]) + var_units = Variable(var.dims, data, attrs=var.attrs, encoding=var.encoding) + index_vars_units[name] = var_units + + return index_vars_units + + @classmethod + def from_variables(cls, variables, options): + index = PandasIndex.from_variables(variables) + units_dict = {index.index.name: options.get("units")} + return cls(index, units_dict) def sel(self, labels): converted_labels = conversion.convert_indexer_units(labels, self.units) From 0706bc66bff51a5a6a1a2fd3e79a5e790a445334 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 13 Sep 2023 11:01:45 +0200 Subject: [PATCH 09/56] use the new index to attach units to dimension coordinates --- pint_xarray/conversion.py | 23 +++++++++++++++++++++-- 1 file changed, 21 insertions(+), 2 deletions(-) diff --git a/pint_xarray/conversion.py b/pint_xarray/conversion.py index 9c18786a..48a0268d 100644 --- a/pint_xarray/conversion.py +++ b/pint_xarray/conversion.py @@ -1,10 +1,11 @@ import itertools import pint -from xarray import DataArray, Dataset, IndexVariable, Variable +from xarray import Coordinates, DataArray, Dataset, IndexVariable, Variable from .compat import call_on_dataset from .errors import format_error_message +from .index import PintMetaIndex no_unit_values = ("none", None) unit_attribute_name = "units" @@ -122,7 +123,12 @@ def dataset_from_variables(variables, coords, attrs): def attach_units_dataset(obj, units): attached = {} rejected_vars = {} + + indexed_variables = obj.xindexes.variables for name, var in obj.variables.items(): + if name in indexed_variables: + continue + unit = units.get(name) try: converted = attach_units_variable(var, unit) @@ -130,10 +136,23 @@ def attach_units_dataset(obj, units): except ValueError as e: rejected_vars[name] = (unit, e) + ds_xindexes = obj.xindexes + new_indexes, new_index_vars = ds_xindexes.copy_indexes() + + for idx, idx_vars in ds_xindexes.group_by_index(): + idx_units = {name: units.get(name) for name in idx_vars.keys()} + new_idx = PintMetaIndex(index=idx, units=idx_units) + new_indexes.update({k: new_idx for k in idx_vars}) + new_index_vars.update(new_idx.create_variables(idx_vars)) + + new_coords = Coordinates(new_index_vars, new_indexes) + if rejected_vars: raise ValueError(rejected_vars) - return dataset_from_variables(attached, obj._coord_names, obj.attrs) + return dataset_from_variables(attached, obj._coord_names, obj.attrs).assign_coords( + new_coords + ) def attach_units(obj, units): From a603860db08bd60623f2aef0cf5d928105a3b6e0 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 13 Sep 2023 11:07:49 +0200 Subject: [PATCH 10/56] pass the dictionary of indexers instead iterating manually --- pint_xarray/index.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/pint_xarray/index.py b/pint_xarray/index.py index b234e783..3e74738b 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -38,9 +38,6 @@ def from_variables(cls, variables, options): def sel(self, labels): converted_labels = conversion.convert_indexer_units(labels, self.units) - stripped_labels = { - name: conversion.strip_indexer_units(indexer) - for name, indexer in converted_labels.items() - } + stripped_labels = conversion.strip_indexer_units(converted_labels) return self.index.sel(stripped_labels) From dea881c96a3129f586143647e03eabf4aec37f49 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 13 Sep 2023 12:55:28 +0200 Subject: [PATCH 11/56] use `Coordinates._construct_direct` --- pint_xarray/conversion.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pint_xarray/conversion.py b/pint_xarray/conversion.py index 48a0268d..1e244efb 100644 --- a/pint_xarray/conversion.py +++ b/pint_xarray/conversion.py @@ -145,7 +145,7 @@ def attach_units_dataset(obj, units): new_indexes.update({k: new_idx for k in idx_vars}) new_index_vars.update(new_idx.create_variables(idx_vars)) - new_coords = Coordinates(new_index_vars, new_indexes) + new_coords = Coordinates._construct_direct(new_index_vars, new_indexes) if rejected_vars: raise ValueError(rejected_vars) From a54b94a28f043c031c1fbcde4f9a08d507106853 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Thu, 14 Sep 2023 20:43:52 +0200 Subject: [PATCH 12/56] delegate `isel` to the wrapped index and wrap the result --- pint_xarray/index.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/pint_xarray/index.py b/pint_xarray/index.py index 3e74738b..a57fd92b 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -41,3 +41,10 @@ def sel(self, labels): stripped_labels = conversion.strip_indexer_units(converted_labels) return self.index.sel(stripped_labels) + + def isel(self, indexers): + subset = self.index.isel(indexers) + if subset is None: + return None + + return type(self)(index=subset, units=self.units) From f0d08907102c7b70076647c5e4868d646906a5a7 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Wed, 6 Dec 2023 21:38:20 +0100 Subject: [PATCH 13/56] add a inline `repr` for the index --- pint_xarray/index.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pint_xarray/index.py b/pint_xarray/index.py index a57fd92b..92390bd5 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -5,7 +5,6 @@ class PintMetaIndex(Index): - # TODO: inherit from MetaIndex once that exists def __init__(self, *, index, units): """create a unit-aware MetaIndex @@ -48,3 +47,6 @@ def isel(self, indexers): return None return type(self)(index=subset, units=self.units) + + def _repr_inline_(self, max_width): + return f"{self.__class__.__name__}({self.index.__class__.__name__})" From fa9f1b3426c97c742b86f2d80beb74df6934565d Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 9 Dec 2023 20:53:45 +0100 Subject: [PATCH 14/56] stubs for the remaining methods --- pint_xarray/index.py | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/pint_xarray/index.py b/pint_xarray/index.py index 92390bd5..fbb7495f 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -35,6 +35,17 @@ def from_variables(cls, variables, options): units_dict = {index.index.name: options.get("units")} return cls(index, units_dict) + @classmethod + def concat(cls, indexes, dim, positions): + raise NotImplementedError() + + @classmethod + def stack(cls, variables, dim): + raise NotImplementedError() + + def unstack(self): + raise NotImplementedError() + def sel(self, labels): converted_labels = conversion.convert_indexer_units(labels, self.units) stripped_labels = conversion.strip_indexer_units(converted_labels) @@ -48,5 +59,23 @@ def isel(self, indexers): return type(self)(index=subset, units=self.units) + def join(self, other, how="inner"): + raise NotImplementedError() + + def reindex_like(self, other): + raise NotImplementedError() + + def equals(self, other): + raise NotImplementedError() + + def roll(self, shifts): + return None + + def rename(self, name_dict, dims_dict): + return self + + def __getitem__(self, indexer): + raise NotImplementedError() + def _repr_inline_(self, max_width): return f"{self.__class__.__name__}({self.index.__class__.__name__})" From fb01e32ab18350a3b3cf21ddca44f2ead10601a2 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 9 Dec 2023 23:22:54 +0100 Subject: [PATCH 15/56] rename the index class to `PintIndex` --- pint_xarray/conversion.py | 4 ++-- pint_xarray/index.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pint_xarray/conversion.py b/pint_xarray/conversion.py index 1e244efb..7bdd2058 100644 --- a/pint_xarray/conversion.py +++ b/pint_xarray/conversion.py @@ -5,7 +5,7 @@ from .compat import call_on_dataset from .errors import format_error_message -from .index import PintMetaIndex +from .index import PintIndex no_unit_values = ("none", None) unit_attribute_name = "units" @@ -141,7 +141,7 @@ def attach_units_dataset(obj, units): for idx, idx_vars in ds_xindexes.group_by_index(): idx_units = {name: units.get(name) for name in idx_vars.keys()} - new_idx = PintMetaIndex(index=idx, units=idx_units) + new_idx = PintIndex(index=idx, units=idx_units) new_indexes.update({k: new_idx for k in idx_vars}) new_index_vars.update(new_idx.create_variables(idx_vars)) diff --git a/pint_xarray/index.py b/pint_xarray/index.py index fbb7495f..8efb0e43 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -4,7 +4,7 @@ from . import conversion -class PintMetaIndex(Index): +class PintIndex(Index): def __init__(self, *, index, units): """create a unit-aware MetaIndex From 9278a2d9bb954587e913bf3ec8bf73c0b48f0646 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 9 Dec 2023 23:41:43 +0100 Subject: [PATCH 16/56] add a utility method to wrap the output of the wrapped index's methods --- pint_xarray/index.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pint_xarray/index.py b/pint_xarray/index.py index 8efb0e43..4679586c 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -18,6 +18,9 @@ def __init__(self, *, index, units): self.index = index self.units = units + def _replace(self, new_index): + return self.__class__(index=new_index, units=self.units) + def create_variables(self, variables=None): index_vars = self.index.create_variables(variables) From 3200bc89a0c4cfbef54717a3471c7e5a47c6bc18 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 9 Dec 2023 23:42:17 +0100 Subject: [PATCH 17/56] implement `equals` --- pint_xarray/index.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/pint_xarray/index.py b/pint_xarray/index.py index 4679586c..0b5b994f 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -69,7 +69,15 @@ def reindex_like(self, other): raise NotImplementedError() def equals(self, other): - raise NotImplementedError() + if not isinstance(other, PintIndex): + return False + + # for now we require exactly matching units to avoid the potentially expensive computation + if self.units != other.units: + return False + + # last to avoid the potentially expensive comparison + return self.index.equals(other.index) def roll(self, shifts): return None From 281f03c8c706996d3cd8293f0537eb5c15dd0561 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 9 Dec 2023 23:42:28 +0100 Subject: [PATCH 18/56] implement `roll`, `rename`, and `__getitem__` by forwarding --- pint_xarray/index.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pint_xarray/index.py b/pint_xarray/index.py index 0b5b994f..bda6cfd9 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -80,13 +80,13 @@ def equals(self, other): return self.index.equals(other.index) def roll(self, shifts): - return None + return self._replace(self.index.roll(shifts)) def rename(self, name_dict, dims_dict): - return self + return self._replace(self.index.rename(name_dict, dims_dict)) def __getitem__(self, indexer): - raise NotImplementedError() + return self._replace(self.index[indexer]) def _repr_inline_(self, max_width): return f"{self.__class__.__name__}({self.index.__class__.__name__})" From c5e902232410509e9b35e418ba2d069551becda0 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 10 Dec 2023 16:38:39 +0100 Subject: [PATCH 19/56] start adding tests --- pint_xarray/index.py | 2 +- pint_xarray/tests/test_index.py | 35 +++++++++++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) create mode 100644 pint_xarray/tests/test_index.py diff --git a/pint_xarray/index.py b/pint_xarray/index.py index bda6cfd9..b774decf 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -72,7 +72,7 @@ def equals(self, other): if not isinstance(other, PintIndex): return False - # for now we require exactly matching units to avoid the potentially expensive computation + # for now we require exactly matching units to avoid the potentially expensive conversion if self.units != other.units: return False diff --git a/pint_xarray/tests/test_index.py b/pint_xarray/tests/test_index.py new file mode 100644 index 00000000..fa1efc37 --- /dev/null +++ b/pint_xarray/tests/test_index.py @@ -0,0 +1,35 @@ +import pandas as pd +import pytest +from xarray.core.indexes import PandasIndex + +from pint_xarray import unit_registry as ureg +from pint_xarray.index import PintIndex + + +@pytest.mark.parametrize( + "base_index", + [ + PandasIndex(pd.Index([1, 2, 3]), dim="x"), + PandasIndex(pd.Index([0.1, 0.2, 0.3]), dim="x"), + PandasIndex(pd.Index([1j, 2j, 3j]), dim="y"), + ], +) +@pytest.mark.parametrize("units", [ureg.Unit("m"), ureg.Unit("s")]) +def test_init(base_index, units): + index = PintIndex(index=base_index, units=units) + + assert index.index.equals(base_index) + assert index.units == units + + +def test_replace(): + old_index = PandasIndex([1, 2, 3], dim="y") + new_index = PandasIndex([0.1, 0.2, 0.3], dim="x") + + old = PintIndex(index=old_index, units=ureg.Unit("m")) + new = old._replace(new_index) + + assert new.index.equals(new_index) + assert new.units == old.units + # no mutation + assert old.index.equals(old_index) From 2c2c8148799e5fe28bdfeab621cb7eb5a9a4cf62 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 10 Dec 2023 20:02:18 +0100 Subject: [PATCH 20/56] add tests for `create_variables` --- pint_xarray/tests/test_index.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/pint_xarray/tests/test_index.py b/pint_xarray/tests/test_index.py index fa1efc37..8d83449f 100644 --- a/pint_xarray/tests/test_index.py +++ b/pint_xarray/tests/test_index.py @@ -1,5 +1,6 @@ import pandas as pd import pytest +import xarray as xr from xarray.core.indexes import PandasIndex from pint_xarray import unit_registry as ureg @@ -33,3 +34,27 @@ def test_replace(): assert new.units == old.units # no mutation assert old.index.equals(old_index) + + +@pytest.mark.parametrize( + ["wrapped_index", "units", "expected"], + ( + pytest.param( + PandasIndex(pd.Index([1, 2, 3]), dim="x"), + {"x": ureg.Unit("m")}, + {"x": xr.Variable("x", ureg.Quantity([1, 2, 3], "m"))}, + ), + pytest.param( + PandasIndex(pd.Index([1j, 2j, 3j]), dim="y"), + {"y": ureg.Unit("ms")}, + {"y": xr.Variable("y", ureg.Quantity([1j, 2j, 3j], "ms"))}, + ), + ), +) +def test_create_variables(wrapped_index, units, expected): + index = PintIndex(index=wrapped_index, units=units) + + actual = index.create_variables() + + assert list(actual.keys()) == list(expected.keys()) + assert all([actual[k].equals(expected[k]) for k in expected.keys()]) From e5d8369b770dc573d1e048d8a3fce08b54afd4bc Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 10 Dec 2023 20:51:57 +0100 Subject: [PATCH 21/56] add tests for `sel` --- pint_xarray/tests/test_index.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/pint_xarray/tests/test_index.py b/pint_xarray/tests/test_index.py index 8d83449f..abd45471 100644 --- a/pint_xarray/tests/test_index.py +++ b/pint_xarray/tests/test_index.py @@ -1,7 +1,7 @@ import pandas as pd import pytest import xarray as xr -from xarray.core.indexes import PandasIndex +from xarray.core.indexes import IndexSelResult, PandasIndex from pint_xarray import unit_registry as ureg from pint_xarray.index import PintIndex @@ -58,3 +58,21 @@ def test_create_variables(wrapped_index, units, expected): assert list(actual.keys()) == list(expected.keys()) assert all([actual[k].equals(expected[k]) for k in expected.keys()]) + + +@pytest.mark.parametrize( + ["labels", "expected"], + ( + ({"x": ureg.Quantity(1, "m")}, IndexSelResult(dim_indexers={"x": 0})), + ({"x": ureg.Quantity(3000, "mm")}, IndexSelResult(dim_indexers={"x": 2})), + ({"x": ureg.Quantity(0.002, "km")}, IndexSelResult(dim_indexers={"x": 1})), + ), +) +def test_sel(labels, expected): + index = PintIndex( + index=PandasIndex(pd.Index([1, 2, 3, 4]), dim="x"), units={"x": ureg.Unit("m")} + ) + + actual = index.sel(labels) + + assert actual == expected From 50a728764e3d4f4d6d19dd95ba0130a500483f18 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 10 Dec 2023 21:02:07 +0100 Subject: [PATCH 22/56] add tests for `isel` --- pint_xarray/tests/test_index.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/pint_xarray/tests/test_index.py b/pint_xarray/tests/test_index.py index abd45471..fd476b0e 100644 --- a/pint_xarray/tests/test_index.py +++ b/pint_xarray/tests/test_index.py @@ -76,3 +76,24 @@ def test_sel(labels, expected): actual = index.sel(labels) assert actual == expected + + +@pytest.mark.parametrize( + "indexers", + ({"y": 0}, {"y": [1, 2]}, {"y": slice(0, None, 2)}, {"y": xr.Variable("y", [1])}), +) +def test_isel(indexers): + wrapped_index = PandasIndex(pd.Index([1, 2, 3, 4]), dim="y") + index = PintIndex(index=wrapped_index, units={"y": ureg.Unit("s")}) + + actual = index.isel(indexers) + + wrapped_ = wrapped_index.isel(indexers) + if wrapped_ is not None: + expected = PintIndex( + index=wrapped_index.isel(indexers), units={"y": ureg.Unit("s")} + ) + else: + expected = None + + assert (actual is None and expected is None) or actual.equals(expected) From 2b3c5bb495a1098b0dba4c1457705bc4d0b2664f Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 10 Dec 2023 21:11:12 +0100 Subject: [PATCH 23/56] improve the tests for `sel` --- pint_xarray/tests/test_index.py | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/pint_xarray/tests/test_index.py b/pint_xarray/tests/test_index.py index fd476b0e..3a04000b 100644 --- a/pint_xarray/tests/test_index.py +++ b/pint_xarray/tests/test_index.py @@ -1,3 +1,4 @@ +import numpy as np import pandas as pd import pytest import xarray as xr @@ -7,6 +8,16 @@ from pint_xarray.index import PintIndex +def indexer_equal(first, second): + if type(first) is not type(second): + return False + + if isinstance(first, np.ndarray): + return np.all(first == second) + else: + return first == second + + @pytest.mark.parametrize( "base_index", [ @@ -66,6 +77,14 @@ def test_create_variables(wrapped_index, units, expected): ({"x": ureg.Quantity(1, "m")}, IndexSelResult(dim_indexers={"x": 0})), ({"x": ureg.Quantity(3000, "mm")}, IndexSelResult(dim_indexers={"x": 2})), ({"x": ureg.Quantity(0.002, "km")}, IndexSelResult(dim_indexers={"x": 1})), + ( + {"x": ureg.Quantity([0.002, 0.004], "km")}, + IndexSelResult(dim_indexers={"x": np.array([1, 3])}), + ), + ( + {"x": slice(ureg.Quantity(2, "m"), ureg.Quantity(3000, "mm"))}, + IndexSelResult(dim_indexers={"x": slice(1, 3)}), + ), ), ) def test_sel(labels, expected): @@ -75,7 +94,14 @@ def test_sel(labels, expected): actual = index.sel(labels) - assert actual == expected + assert isinstance(actual, IndexSelResult) + assert list(actual.dim_indexers.keys()) == list(expected.dim_indexers.keys()) + assert all( + [ + indexer_equal(actual.dim_indexers[k], expected.dim_indexers[k]) + for k in expected.dim_indexers.keys() + ] + ) @pytest.mark.parametrize( From 57ea8e5a26c01134ff2243acb3927bba55d1163d Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 10 Dec 2023 21:15:27 +0100 Subject: [PATCH 24/56] add tests for `equals` --- pint_xarray/tests/test_index.py | 44 +++++++++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/pint_xarray/tests/test_index.py b/pint_xarray/tests/test_index.py index 3a04000b..98dbe2c8 100644 --- a/pint_xarray/tests/test_index.py +++ b/pint_xarray/tests/test_index.py @@ -123,3 +123,47 @@ def test_isel(indexers): expected = None assert (actual is None and expected is None) or actual.equals(expected) + + +@pytest.mark.parametrize( + ["other", "expected"], + ( + ( + PintIndex( + index=PandasIndex(pd.Index([1, 2, 3, 4]), dim="x"), + units={"x": ureg.Unit("cm")}, + ), + True, + ), + (PandasIndex(pd.Index([1, 2, 3, 4]), dim="x"), False), + ( + PintIndex( + index=PandasIndex(pd.Index([1, 2, 3, 4]), dim="x"), + units={"x": ureg.Unit("m")}, + ), + False, + ), + ( + PintIndex( + index=PandasIndex(pd.Index([1, 2, 3, 4]), dim="y"), + units={"y": ureg.Unit("cm")}, + ), + False, + ), + ( + PintIndex( + index=PandasIndex(pd.Index([1, 3, 3, 4]), dim="x"), + units={"x": ureg.Unit("cm")}, + ), + False, + ), + ), +) +def test_equals(other, expected): + index = PintIndex( + index=PandasIndex(pd.Index([1, 2, 3, 4]), dim="x"), units={"x": ureg.Unit("cm")} + ) + + actual = index.equals(other) + + assert actual == expected From 3eed8c94bc3da3d96b357d446d05d5f4c06c741f Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 10 Dec 2023 21:20:09 +0100 Subject: [PATCH 25/56] add tests for `roll` --- pint_xarray/tests/test_index.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/pint_xarray/tests/test_index.py b/pint_xarray/tests/test_index.py index 98dbe2c8..68ba4d94 100644 --- a/pint_xarray/tests/test_index.py +++ b/pint_xarray/tests/test_index.py @@ -167,3 +167,25 @@ def test_equals(other, expected): actual = index.equals(other) assert actual == expected + + +@pytest.mark.parametrize( + ["shifts", "expected_index"], + ( + ({"x": 0}, PandasIndex(pd.Index([-2, -1, 0, 1, 2]), dim="x")), + ({"x": 1}, PandasIndex(pd.Index([2, -2, -1, 0, 1]), dim="x")), + ({"x": 2}, PandasIndex(pd.Index([1, 2, -2, -1, 0]), dim="x")), + ({"x": -1}, PandasIndex(pd.Index([-1, 0, 1, 2, -2]), dim="x")), + ({"x": -2}, PandasIndex(pd.Index([0, 1, 2, -2, -1]), dim="x")), + ), +) +def test_roll(shifts, expected_index): + index = PintIndex( + index=PandasIndex(pd.Index([-2, -1, 0, 1, 2]), dim="x"), + units={"x": ureg.Unit("m")}, + ) + + actual = index.roll(shifts) + expected = index._replace(expected_index) + + assert actual.equals(expected) From aebaf37f9aa272175898b30fee37a6167b93ebb9 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 10 Dec 2023 21:25:20 +0100 Subject: [PATCH 26/56] add tests for `rename` --- pint_xarray/tests/test_index.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/pint_xarray/tests/test_index.py b/pint_xarray/tests/test_index.py index 68ba4d94..e9fe5300 100644 --- a/pint_xarray/tests/test_index.py +++ b/pint_xarray/tests/test_index.py @@ -189,3 +189,17 @@ def test_roll(shifts, expected_index): expected = index._replace(expected_index) assert actual.equals(expected) + + +@pytest.mark.parametrize("dims_dict", ({"y": "x"}, {"y": "z"})) +@pytest.mark.parametrize("name_dict", ({"y2": "y3"}, {"y2": "y1"})) +def test_rename(name_dict, dims_dict): + wrapped_index = PandasIndex(pd.Index([1, 2], name="y2"), dim="y") + index = PintIndex(index=wrapped_index, units={"y": ureg.Unit("m")}) + + actual = index.rename(name_dict, dims_dict) + expected = PintIndex( + index=wrapped_index.rename(name_dict, dims_dict), units=index.units + ) + + assert actual.equals(expected) From 58c540faa902c05a182117b41ce0ac6f571c2473 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 10 Dec 2023 21:29:32 +0100 Subject: [PATCH 27/56] add tests for `__getitem__` --- pint_xarray/tests/test_index.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/pint_xarray/tests/test_index.py b/pint_xarray/tests/test_index.py index e9fe5300..1b16de0f 100644 --- a/pint_xarray/tests/test_index.py +++ b/pint_xarray/tests/test_index.py @@ -203,3 +203,14 @@ def test_rename(name_dict, dims_dict): ) assert actual.equals(expected) + + +@pytest.mark.parametrize("indexer", ([0], slice(0, 2))) +def test_getitem(indexer): + wrapped_index = PandasIndex(pd.Index([1, 2], name="y2"), dim="y") + index = PintIndex(index=wrapped_index, units={"y": ureg.Unit("m")}) + + actual = index[indexer] + expected = PintIndex(index=wrapped_index[indexer], units=index.units) + + assert actual.equals(expected) From 9cb7e912327b97f5ee9e5ecdbf376c01ad0c343e Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 10 Dec 2023 21:32:01 +0100 Subject: [PATCH 28/56] add tests for `_repr_inline_` --- pint_xarray/tests/test_index.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/pint_xarray/tests/test_index.py b/pint_xarray/tests/test_index.py index 1b16de0f..13f49064 100644 --- a/pint_xarray/tests/test_index.py +++ b/pint_xarray/tests/test_index.py @@ -214,3 +214,14 @@ def test_getitem(indexer): expected = PintIndex(index=wrapped_index[indexer], units=index.units) assert actual.equals(expected) + + +@pytest.mark.parametrize("wrapped_index", (PandasIndex(pd.Index([1, 2]), dim="x"),)) +def test_repr_inline(wrapped_index): + index = PintIndex(index=wrapped_index, units=ureg.Unit("m")) + + # TODO: parametrize + actual = index._repr_inline_(90) + + assert "PintIndex" in actual + assert wrapped_index.__class__.__name__ in actual From 9822520b0aff8403013d612ace04b2241223ab32 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 11 Dec 2023 15:51:51 +0100 Subject: [PATCH 29/56] configure coverage, just in case --- pyproject.toml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 87627dfd..8561e002 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,3 +52,11 @@ skip_gitignore = "true" force_to_top = "true" default_section = "THIRDPARTY" known_first_party = "pint_xarray" + +[tool.coverage.run] +source = ["pint_xarray"] +branch = true + +[tool.coverage.report] +show_missing = true +exclude_lines = ["pragma: no cover", "if TYPE_CHECKING"] From 55ccb00cdb0ef295c9ae66a2204a8f2d7b0fc40b Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 11 Dec 2023 15:52:16 +0100 Subject: [PATCH 30/56] use `_replace` instead of manually constructing the new index --- pint_xarray/index.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pint_xarray/index.py b/pint_xarray/index.py index b774decf..570ced3a 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -60,7 +60,7 @@ def isel(self, indexers): if subset is None: return None - return type(self)(index=subset, units=self.units) + return self._replace(subset) def join(self, other, how="inner"): raise NotImplementedError() From 6bd672612c1243d8e6661cc584e87c4a66df1ff7 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Fri, 22 Dec 2023 16:33:48 +0100 Subject: [PATCH 31/56] explicitly check that the pint index gets created --- pint_xarray/tests/test_conversion.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/pint_xarray/tests/test_conversion.py b/pint_xarray/tests/test_conversion.py index 8ed12b68..378897c1 100644 --- a/pint_xarray/tests/test_conversion.py +++ b/pint_xarray/tests/test_conversion.py @@ -1,9 +1,11 @@ import numpy as np import pint import pytest -from xarray import DataArray, Dataset, Variable +from xarray import Coordinates, DataArray, Dataset, Variable +from xarray.core.indexes import PandasIndex from pint_xarray import conversion +from pint_xarray.index import PintIndex from .utils import ( assert_array_equal, @@ -245,17 +247,22 @@ def test_attach_units(self, type, units): q_a = to_quantity(a, units.get("a")) q_b = to_quantity(b, units.get("b")) + q_x = to_quantity(x, units.get("x")) q_u = to_quantity(u, units.get("u")) - units_x = units.get("x") + index = PandasIndex(x, dim="x") + if units.get("x") is not None: + index = PintIndex(index=index, units=units.get("x")) obj = Dataset({"a": ("x", a), "b": ("x", b)}, coords={"u": ("x", u), "x": x}) + coords = Coordinates._construct_direct( + coords={"u": Variable("x", q_u), "x": Variable("x", q_x)}, + indexes={"x": index}, + ) expected = Dataset( {"a": ("x", q_a), "b": ("x", q_b)}, - coords={"u": ("x", q_u), "x": x}, + coords=coords, ) - if units_x is not None: - expected.x.attrs["units"] = units_x if type == "DataArray": obj = obj["a"] @@ -264,6 +271,11 @@ def test_attach_units(self, type, units): actual = conversion.attach_units(obj, units) assert_identical(actual, expected) + if units.get("x") is not None: + index = actual.xindexes["x"] + + assert isinstance(index, PintIndex) and index.units == {"x": units.get("x")} + @pytest.mark.parametrize("type", ("DataArray", "Dataset")) def test_attach_unit_attributes(self, type): units = {"a": "K", "b": "hPa", "u": "m", "x": "s"} From c7d523b7df0d550980ea0c0f181cddb5380a755d Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Fri, 22 Dec 2023 16:41:31 +0100 Subject: [PATCH 32/56] also verify that non-quantity variables don't become `PintIndex`ed --- pint_xarray/conversion.py | 3 +++ pint_xarray/tests/test_conversion.py | 9 +++++---- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/pint_xarray/conversion.py b/pint_xarray/conversion.py index 7bdd2058..5d375091 100644 --- a/pint_xarray/conversion.py +++ b/pint_xarray/conversion.py @@ -141,6 +141,9 @@ def attach_units_dataset(obj, units): for idx, idx_vars in ds_xindexes.group_by_index(): idx_units = {name: units.get(name) for name in idx_vars.keys()} + if all(unit is None for unit in idx_units.values()): + # skip non-quantity indexed variables + continue new_idx = PintIndex(index=idx, units=idx_units) new_indexes.update({k: new_idx for k in idx_vars}) new_index_vars.update(new_idx.create_variables(idx_vars)) diff --git a/pint_xarray/tests/test_conversion.py b/pint_xarray/tests/test_conversion.py index 378897c1..9bd02c12 100644 --- a/pint_xarray/tests/test_conversion.py +++ b/pint_xarray/tests/test_conversion.py @@ -271,10 +271,11 @@ def test_attach_units(self, type, units): actual = conversion.attach_units(obj, units) assert_identical(actual, expected) - if units.get("x") is not None: - index = actual.xindexes["x"] - - assert isinstance(index, PintIndex) and index.units == {"x": units.get("x")} + if units.get("x") is None: + assert not isinstance(actual.xindexes["x"], PintIndex) + else: + assert isinstance(actual.xindexes["x"], PintIndex) + assert actual.xindexes["x"].units == {"x": units.get("x")} @pytest.mark.parametrize("type", ("DataArray", "Dataset")) def test_attach_unit_attributes(self, type): From 9dde67dcf6d574e994cad43a7a431292294f9b39 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 23 Jun 2024 18:16:57 +0200 Subject: [PATCH 33/56] don't use `.pint.sel` --- docs/examples/plotting.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/examples/plotting.ipynb b/docs/examples/plotting.ipynb index 5a4a0c7e..3955fb0f 100644 --- a/docs/examples/plotting.ipynb +++ b/docs/examples/plotting.ipynb @@ -108,7 +108,7 @@ "metadata": {}, "outputs": [], "source": [ - "monthly_means.pint.sel(\n", + "monthly_means.sel(\n", " lat=ureg.Quantity(4350, \"angular_minute\"),\n", " lon=ureg.Quantity(12000, \"angular_minute\"),\n", ")" From b927436a599fc7a7c7a4854c2a6b01102dd60298 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 23 Jun 2024 22:11:39 +0200 Subject: [PATCH 34/56] fix `PintIndex.from_variables` and properly test it --- pint_xarray/index.py | 9 ++++++--- pint_xarray/tests/test_accessors.py | 13 ++++++++++++- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/pint_xarray/index.py b/pint_xarray/index.py index 570ced3a..5d13ff1f 100644 --- a/pint_xarray/index.py +++ b/pint_xarray/index.py @@ -34,9 +34,12 @@ def create_variables(self, variables=None): @classmethod def from_variables(cls, variables, options): - index = PandasIndex.from_variables(variables) - units_dict = {index.index.name: options.get("units")} - return cls(index, units_dict) + if len(variables) != 1: + raise ValueError("can only create a default index from single variables") + + units = options.pop("units", None) + index = PandasIndex.from_variables(variables, options=options) + return cls(index=index, units={index.index.name: units}) @classmethod def concat(cls, indexes, dim, positions): diff --git a/pint_xarray/tests/test_accessors.py b/pint_xarray/tests/test_accessors.py index 46d4ed61..fb4c95bd 100644 --- a/pint_xarray/tests/test_accessors.py +++ b/pint_xarray/tests/test_accessors.py @@ -7,6 +7,7 @@ from pint import Unit, UnitRegistry from .. import accessors, conversion +from ..index import PintIndex from .utils import ( assert_equal, assert_identical, @@ -159,7 +160,17 @@ def test_dimension_coordinate_array_already_quantified(self): arr.pint.quantify({"x": "s"}) def test_dimension_coordinate_array_already_quantified_same_units(self): - ds = xr.Dataset(coords={"x": ("x", [10], {"units": unit_registry.Unit("m")})}) + x = unit_registry.Quantity([10], "m") + coords = xr.Coordinates( + {"x": x}, + indexes={ + "x": PintIndex.from_variables( + {"x": xr.Variable("x", x.magnitude)}, + options={"units": x.units}, + ), + }, + ) + ds = xr.Dataset(coords=coords) arr = ds.x quantified = arr.pint.quantify({"x": "m"}) From c31e6b08a01f934f3e19b0459fd5d8d7ccff94ec Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sun, 23 Jun 2024 22:35:44 +0200 Subject: [PATCH 35/56] quantify the test data This allows us to test with quantified indexes. --- pint_xarray/tests/test_accessors.py | 114 +++++++++++----------------- 1 file changed, 43 insertions(+), 71 deletions(-) diff --git a/pint_xarray/tests/test_accessors.py b/pint_xarray/tests/test_accessors.py index fb4c95bd..8e955bb6 100644 --- a/pint_xarray/tests/test_accessors.py +++ b/pint_xarray/tests/test_accessors.py @@ -23,7 +23,8 @@ # make sure scalars are converted to 0d arrays so quantities can # always be treated like ndarrays -unit_registry = UnitRegistry(force_ndarray=True) +from pint_xarray import unit_registry + Quantity = unit_registry.Quantity nan = np.nan @@ -1412,73 +1413,56 @@ def test_reindex_like(obj, other, expected, error): @requires_scipy @pytest.mark.parametrize( - ["obj", "indexers", "expected", "error"], + ["obj", "units", "indexers", "expected", "expected_units", "error"], ( pytest.param( - xr.Dataset( - { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - } - ), + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, {"x": Quantity([10, 30, 50], "dm"), "y": Quantity([0, 120, 240], "s")}, - xr.Dataset( - { - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 120, 240], {"units": unit_registry.Unit("s")}), - } - ), + xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), + {"x": "dm", "y": "s"}, None, id="Dataset-identical units", ), pytest.param( - xr.Dataset( - { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - } - ), + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, {"x": Quantity([0, 1, 3, 5], "m"), "y": Quantity([0, 2, 4], "min")}, - xr.Dataset( - { - "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2, 4], {"units": unit_registry.Unit("min")}), - } - ), + xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), + {"x": "m", "y": "min"}, None, id="Dataset-compatible units", ), pytest.param( - xr.Dataset( - { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - } - ), + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, {"x": Quantity([1, 3], "s"), "y": Quantity([1], "m")}, None, + {}, ValueError, id="Dataset-incompatible units", ), pytest.param( xr.Dataset( { - "a": (("x", "y"), Quantity([[0, 1], [2, 3], [4, 5]], "kg")), + "a": (("x", "y"), np.array([[0, 1], [2, 3], [4, 5]])), "x": [10, 20, 30], "y": [60, 120], } ), + {"a": "kg"}, { "x": [15, 25], "y": [75, 105], }, xr.Dataset( { - "a": (("x", "y"), Quantity([[1.25, 1.75], [3.25, 3.75]], "kg")), + "a": (("x", "y"), np.array([[1.25, 1.75], [3.25, 3.75]])), "x": [15, 25], "y": [75, 105], } ), + {"a": "kg"}, None, id="Dataset-data units", ), @@ -1486,20 +1470,16 @@ def test_reindex_like(obj, other, expected, error): xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), + {"x": "dm", "y": "s"}, {"x": Quantity([10, 30, 50], "dm"), "y": Quantity([0, 240], "s")}, xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), - coords={ - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 240], {"units": unit_registry.Unit("s")}), - }, + coords={"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}, ), + {"x": "dm", "y": "s"}, None, id="DataArray-identical units", ), @@ -1507,20 +1487,16 @@ def test_reindex_like(obj, other, expected, error): xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), + {"x": "dm", "y": "s"}, {"x": Quantity([1, 3, 5], "m"), "y": Quantity([0, 2], "min")}, xr.DataArray( [[np.nan, 1], [np.nan, 5], [np.nan, np.nan]], dims=("x", "y"), - coords={ - "x": ("x", [1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2], {"units": unit_registry.Unit("min")}), - }, + coords={"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}, ), + {"x": "m", "y": "min"}, None, id="DataArray-compatible units", ), @@ -1528,50 +1504,46 @@ def test_reindex_like(obj, other, expected, error): xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), + {"x": "dm", "y": "s"}, {"x": Quantity([10, 30], "s"), "y": Quantity([60], "m")}, None, + {}, ValueError, id="DataArray-incompatible units", ), pytest.param( xr.DataArray( - Quantity([[0, 1], [2, 3], [4, 5]], "kg"), + np.array([[0, 1], [2, 3], [4, 5]]), dims=("x", "y"), - coords={ - "x": [10, 20, 30], - "y": [60, 120], - }, + coords={"x": [10, 20, 30], "y": [60, 120]}, ), - { - "x": [15, 25], - "y": [75, 105], - }, + {None: "kg"}, + {"x": [15, 25], "y": [75, 105]}, xr.DataArray( Quantity([[1.25, 1.75], [3.25, 3.75]], "kg"), dims=("x", "y"), - coords={ - "x": [15, 25], - "y": [75, 105], - }, + coords={"x": [15, 25], "y": [75, 105]}, ), + {None: "kg"}, None, id="DataArray-data units", ), ), ) -def test_interp(obj, indexers, expected, error): +def test_interp(obj, units, indexers, expected, expected_units, error): + obj_ = obj.pint.quantify(units) + if expected is not None: + expected_ = expected.pint.quantify(expected_units) + if error is not None: with pytest.raises(error): obj.pint.interp(indexers) else: - actual = obj.pint.interp(indexers) - assert_units_equal(actual, expected) - assert_identical(actual, expected) + actual = obj_.pint.interp(indexers) + assert_units_equal(actual, expected_) + assert_identical(actual, expected_) @requires_scipy From 415d05913f15b57b7757ab1f2c6d6fe85e46a55e Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 25 Jun 2024 22:25:35 +0200 Subject: [PATCH 36/56] explicity quantify the input of the `interp_like` tests --- pint_xarray/tests/test_accessors.py | 176 ++++++++++------------------ 1 file changed, 60 insertions(+), 116 deletions(-) diff --git a/pint_xarray/tests/test_accessors.py b/pint_xarray/tests/test_accessors.py index 8e955bb6..1028493f 100644 --- a/pint_xarray/tests/test_accessors.py +++ b/pint_xarray/tests/test_accessors.py @@ -1548,66 +1548,35 @@ def test_interp(obj, units, indexers, expected, expected_units, error): @requires_scipy @pytest.mark.parametrize( - ["obj", "other", "expected", "error"], + ["obj", "units", "other", "other_units", "expected", "expected_units", "error"], ( pytest.param( - xr.Dataset( - { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - } - ), - xr.Dataset( - { - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 120, 240], {"units": unit_registry.Unit("s")}), - } - ), - xr.Dataset( - { - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 120, 240], {"units": unit_registry.Unit("s")}), - } - ), + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), + {"x": "dm", "y": "s"}, None, id="Dataset-identical units", ), pytest.param( - xr.Dataset( - { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - } - ), - xr.Dataset( - { - "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2, 4], {"units": unit_registry.Unit("min")}), - } - ), - xr.Dataset( - { - "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2, 4], {"units": unit_registry.Unit("min")}), - } - ), + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), + {"x": "m", "y": "min"}, + xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), + {"x": "m", "y": "min"}, None, id="Dataset-compatible units", ), pytest.param( - xr.Dataset( - { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - } - ), - xr.Dataset( - { - "x": ("x", [1, 3], {"units": unit_registry.Unit("s")}), - "y": ("y", [1], {"units": unit_registry.Unit("m")}), - } - ), + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [1, 3]), "y": ("y", [1])}), + {"x": "s", "y": "m"}, None, + {}, ValueError, id="Dataset-incompatible units", ), @@ -1615,49 +1584,39 @@ def test_interp(obj, units, indexers, expected, expected_units, error): xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, - ), - xr.Dataset( - { - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 240], {"units": unit_registry.Unit("s")}), - } + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}), + {"x": "dm", "y": "s"}, xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), - coords={ - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 240], {"units": unit_registry.Unit("s")}), - }, + coords={"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}, ), + {"x": "dm", "y": "s"}, None, id="DataArray-identical units", ), pytest.param( xr.Dataset( { - "a": (("x", "y"), Quantity([[0, 1], [2, 3], [4, 5]], "kg")), + "a": (("x", "y"), [[0, 1], [2, 3], [4, 5]]), "x": [10, 20, 30], "y": [60, 120], } ), + {"a": "kg"}, + xr.Dataset({"x": [15, 25], "y": [75, 105]}), + {}, xr.Dataset( { + "a": (("x", "y"), [[1.25, 1.75], [3.25, 3.75]]), "x": [15, 25], "y": [75, 105], } ), - xr.Dataset( - { - "a": (("x", "y"), Quantity([[1.25, 1.75], [3.25, 3.75]], "kg")), - "x": [15, 25], - "y": [75, 105], - } - ), + {"a": "kg"}, None, id="Dataset-data units", ), @@ -1665,25 +1624,17 @@ def test_interp(obj, units, indexers, expected, expected_units, error): xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, - ), - xr.Dataset( - { - "x": ("x", [1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2], {"units": unit_registry.Unit("min")}), - } + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}), + {"x": "m", "y": "min"}, xr.DataArray( [[np.nan, 1], [np.nan, 5], [np.nan, np.nan]], dims=("x", "y"), - coords={ - "x": ("x", [1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2], {"units": unit_registry.Unit("min")}), - }, + coords={"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}, ), + {"x": "m", "y": "min"}, None, id="DataArray-compatible units", ), @@ -1691,57 +1642,50 @@ def test_interp(obj, units, indexers, expected, expected_units, error): xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, - ), - xr.Dataset( - { - "x": ("x", [10, 30], {"units": unit_registry.Unit("s")}), - "y": ("y", [60], {"units": unit_registry.Unit("m")}), - } + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [10, 30]), "y": ("y", [60])}), + {"x": "s", "y": "m"}, None, + {}, ValueError, id="DataArray-incompatible units", ), pytest.param( xr.DataArray( - Quantity([[0, 1], [2, 3], [4, 5]], "kg"), + [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": [10, 20, 30], - "y": [60, 120], - }, - ), - xr.Dataset( - { - "x": [15, 25], - "y": [75, 105], - } + coords={"x": [10, 20, 30], "y": [60, 120]}, ), + {"a": "kg"}, + xr.Dataset({"x": [15, 25], "y": [75, 105]}), + {}, xr.DataArray( - Quantity([[1.25, 1.75], [3.25, 3.75]], "kg"), + [[1.25, 1.75], [3.25, 3.75]], dims=("x", "y"), - coords={ - "x": [15, 25], - "y": [75, 105], - }, + coords={"x": [15, 25], "y": [75, 105]}, ), + {"a": "kg"}, None, id="DataArray-data units", ), ), ) -def test_interp_like(obj, other, expected, error): +def test_interp_like(obj, units, other, other_units, expected, expected_units, error): + obj_ = obj.pint.quantify(units) + other_ = other.pint.quantify(other_units) + + if expected is not None: + expected_ = expected.pint.quantify(expected_units) + if error is not None: with pytest.raises(error): - obj.pint.interp_like(other) + obj_.pint.interp_like(other_) else: - actual = obj.pint.interp_like(other) - assert_units_equal(actual, expected) - assert_identical(actual, expected) + actual = obj_.pint.interp_like(other_) + assert_units_equal(actual, expected_) + assert_identical(actual, expected_) @requires_bottleneck From 1939b2dfee11678ca328b0bb70af50187c213e4f Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 13:27:55 +0200 Subject: [PATCH 37/56] also strip the units of `other` --- pint_xarray/accessors.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pint_xarray/accessors.py b/pint_xarray/accessors.py index bef0cee8..82573451 100644 --- a/pint_xarray/accessors.py +++ b/pint_xarray/accessors.py @@ -756,8 +756,9 @@ def interp_like(self, other, method="linear", assume_sorted=False, kwargs=None): converted = conversion.convert_units(self.da, indexer_units) units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) + stripped_other = conversion.strip_units(other) interpolated = stripped.interp_like( - other, + stripped_other, method=method, assume_sorted=assume_sorted, kwargs=kwargs, @@ -1518,8 +1519,9 @@ def interp_like(self, other, method="linear", assume_sorted=False, kwargs=None): converted = conversion.convert_units(self.ds, indexer_units) units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) + stripped_other = conversion.strip_units(other) interpolated = stripped.interp_like( - other, + stripped_other, method=method, assume_sorted=assume_sorted, kwargs=kwargs, From 25381048c38a61606eefbab3d51b16033f257b4e Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 14:39:12 +0200 Subject: [PATCH 38/56] change expectations in the conversion tests --- pint_xarray/tests/test_conversion.py | 76 +++++++++++++++++++--------- 1 file changed, 51 insertions(+), 25 deletions(-) diff --git a/pint_xarray/tests/test_conversion.py b/pint_xarray/tests/test_conversion.py index 35b8082d..335cdc35 100644 --- a/pint_xarray/tests/test_conversion.py +++ b/pint_xarray/tests/test_conversion.py @@ -1,4 +1,5 @@ import numpy as np +import pandas as pd import pint import pytest from xarray import Coordinates, DataArray, Dataset, Variable @@ -385,15 +386,19 @@ def test_convert_units(self, type, variant, units, error, match): q_u = to_quantity(u, original_units.get("u")) q_x = to_quantity(x, original_units.get("x")) + x_index = PandasIndex(pd.Index(x), "x") + if original_units.get("x") is not None: + x_index = PintIndex(index=x_index, units={"x": original_units.get("x")}) + obj = Dataset( { "a": ("x", q_a), "b": ("x", q_b), }, - coords={ - "u": ("x", q_u), - "x": ("x", x, {"units": original_units.get("x")}), - }, + coords=Coordinates( + {"u": ("x", q_u), "x": ("x", q_x)}, + indexes={"x": x_index}, + ), ) if type == "DataArray": obj = obj["a"] @@ -407,20 +412,22 @@ def test_convert_units(self, type, variant, units, error, match): expected_a = convert_quantity(q_a, units.get("a", original_units.get("a"))) expected_b = convert_quantity(q_b, units.get("b", original_units.get("b"))) expected_u = convert_quantity(q_u, units.get("u", original_units.get("u"))) - expected_x = strip_quantity(convert_quantity(q_x, units.get("x"))) + expected_x = convert_quantity(q_x, units.get("x")) + expected_index = PandasIndex(pd.Index(strip_quantity(expected_x)), "x") + if units.get("x") is not None: + expected_index = PintIndex( + index=expected_index, units={"x": units.get("x")} + ) + expected = Dataset( { "a": ("x", expected_a), "b": ("x", expected_b), }, - coords={ - "u": ("x", expected_u), - "x": ( - "x", - expected_x, - {"units": units.get("x", original_units.get("x"))}, - ), - }, + coords=Coordinates( + {"u": ("x", expected_u), "x": ("x", expected_x)}, + indexes={"x": expected_index}, + ), ) if type == "DataArray": @@ -429,7 +436,7 @@ def test_convert_units(self, type, variant, units, error, match): actual = conversion.convert_units(obj, units) assert conversion.extract_units(actual) == conversion.extract_units(expected) - assert_identical(expected, actual) + assert_identical(actual, expected) @pytest.mark.parametrize( "units", @@ -449,15 +456,22 @@ def test_extract_units(self, type, units): u = np.linspace(0, 100, 2) x = np.arange(2) + index = PandasIndex(x, "x") + if units.get("x") is not None: + index = PintIndex(index=index, units={"x": units.get("x")}) + obj = Dataset( { "a": ("x", to_quantity(a, units.get("a"))), "b": ("x", to_quantity(b, units.get("b"))), }, - coords={ - "u": ("x", to_quantity(u, units.get("u"))), - "x": ("x", x, {"units": units.get("x")}), - }, + coords=Coordinates( + { + "u": ("x", to_quantity(u, units.get("u"))), + "x": ("x", to_quantity(x, units.get("x"))), + }, + indexes={"x": index}, + ), ) if type == "DataArray": obj = obj["a"] @@ -512,21 +526,33 @@ def test_extract_unit_attributes(self, obj, expected): pytest.param( DataArray( dims="x", - data=[0, 4, 3] * unit_registry.m, - coords={"u": ("x", [2, 3, 4] * unit_registry.s)}, + data=Quantity([0, 4, 3], "kg"), + coords=Coordinates( + { + "u": ("x", Quantity([2, 3, 4], "s")), + "x": Quantity([0, 1, 2], "m"), + }, + indexes={}, + ), ), - {None: None, "u": None}, + {None: None, "u": None, "x": None}, id="DataArray", ), pytest.param( Dataset( data_vars={ - "a": ("x", [3, 2, 5] * unit_registry.Pa), - "b": ("x", [0, 2, -1] * unit_registry.kg), + "a": ("x", Quantity([3, 2, 5], "Pa")), + "b": ("x", Quantity([0, 2, -1], "kg")), }, - coords={"u": ("x", [2, 3, 4] * unit_registry.s)}, + coords=Coordinates( + { + "u": ("x", Quantity([2, 3, 4], "s")), + "x": Quantity([0, 1, 2], "m"), + }, + indexes={}, + ), ), - {"a": None, "b": None, "u": None}, + {"a": None, "b": None, "u": None, "x": None}, id="Dataset", ), ), From eb2c405057cc82ebf1178139e6117c338b2f551c Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 14:41:42 +0200 Subject: [PATCH 39/56] refactor `attach_units_dataset` --- pint_xarray/conversion.py | 48 ++++++++++++++++++++++++--------------- 1 file changed, 30 insertions(+), 18 deletions(-) diff --git a/pint_xarray/conversion.py b/pint_xarray/conversion.py index 397165c9..75ff8fda 100644 --- a/pint_xarray/conversion.py +++ b/pint_xarray/conversion.py @@ -122,11 +122,27 @@ def attach_units_variable(variable, units): return new_obj -def dataset_from_variables(variables, coords, attrs): - data_vars = {name: var for name, var in variables.items() if name not in coords} - coords = {name: var for name, var in variables.items() if name in coords} +def dataset_from_variables(variables, coordinate_names, indexes, attrs): + data_vars = { + name: var for name, var in variables.items() if name not in coordinate_names + } + coords = {name: var for name, var in variables.items() if name in coordinate_names} + + new_coords = Coordinates._construct_direct(coords, indexes) + return Dataset(data_vars=data_vars, coords=new_coords, attrs=attrs) + + +def attach_units_index(index, index_vars, units): + if all(unit is None for unit in units.values()): + # skip non-quantity indexed variables + return index + + if isinstance(index, PintIndex) and index.units != units: + raise ValueError( + f"cannot attach units to quantified index: {index.units} != {units}" + ) - return Dataset(data_vars=data_vars, coords=coords, attrs=attrs) + return PintIndex(index=index, units=units) def attach_units_dataset(obj, units): @@ -145,26 +161,22 @@ def attach_units_dataset(obj, units): except ValueError as e: rejected_vars[name] = (unit, e) - ds_xindexes = obj.xindexes - new_indexes, new_index_vars = ds_xindexes.copy_indexes() - - for idx, idx_vars in ds_xindexes.group_by_index(): + indexes, index_vars = obj.xindexes.copy_indexes() + for idx, idx_vars in obj.xindexes.group_by_index(): idx_units = {name: units.get(name) for name in idx_vars.keys()} - if all(unit is None for unit in idx_units.values()): - # skip non-quantity indexed variables - continue - new_idx = PintIndex(index=idx, units=idx_units) - new_indexes.update({k: new_idx for k in idx_vars}) - new_index_vars.update(new_idx.create_variables(idx_vars)) + try: + attached_idx = attach_units_index(idx, idx_vars, idx_units) + indexes.update({k: attached_idx for k in idx_vars}) + index_vars.update(attached_idx.create_variables(idx_vars)) + except ValueError as e: + rejected_vars[name] = (units, e) - new_coords = Coordinates._construct_direct(new_index_vars, new_indexes) + attached.update(index_vars) if rejected_vars: raise ValueError(rejected_vars) - return dataset_from_variables(attached, obj._coord_names, obj.attrs).assign_coords( - new_coords - ) + return dataset_from_variables(attached, obj._coord_names, indexes, obj.attrs) def attach_units(obj, units): From 0d46b669438d49136d866767bdb29f16364248c5 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 14:42:46 +0200 Subject: [PATCH 40/56] get `convert_units` to accept indexes --- pint_xarray/conversion.py | 45 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/pint_xarray/conversion.py b/pint_xarray/conversion.py index 75ff8fda..fe86f578 100644 --- a/pint_xarray/conversion.py +++ b/pint_xarray/conversion.py @@ -249,20 +249,63 @@ def convert_units_variable(variable, units): return new_obj +def convert_units_index(index, index_vars, units): + if not isinstance(index, PintIndex): + raise ValueError("cannot convert non-quantified index") + + converted_vars = {} + failed = {} + for name, var in index_vars.items(): + unit = units.get(name) + try: + converted = convert_units_variable(var, unit) + converted_vars[name] = strip_units_variable(converted) + except (ValueError, pint.errors.PintTypeError) as e: + failed[name] = e + + if failed: + # raise exception group + raise ValueError("failed to convert index variables:", failed) + + # TODO: figure out how to pull out `options` + converted_index = index.index.from_variables(converted_vars, options={}) + return PintIndex(index=converted_index, units=units) + + def convert_units_dataset(obj, units): converted = {} failed = {} + indexed_variables = obj.xindexes.variables for name, var in obj.variables.items(): + if name in indexed_variables: + continue + unit = units.get(name) try: converted[name] = convert_units_variable(var, unit) except (ValueError, pint.errors.PintTypeError) as e: failed[name] = e + indexes, index_vars = obj.xindexes.copy_indexes() + for idx, idx_vars in obj.xindexes.group_by_index(): + idx_units = {name: units.get(name) for name in idx_vars.keys()} + if all(unit is None for unit in idx_units.values()): + continue + + try: + converted_index = convert_units_index(idx, idx_vars, idx_units) + indexes.update({k: converted_index for k in idx_vars}) + index_vars.update(converted_index.create_variables()) + except (ValueError, pint.errors.PintTypeError) as e: + names = tuple(idx_vars) + failed[names] = e + + converted.update(index_vars) + if failed: raise ValueError(failed) - return dataset_from_variables(converted, obj._coord_names, obj.attrs) + return dataset_from_variables(converted, obj._coord_names, indexes, obj.attrs) def convert_units(obj, units): From caf4668b7081d404fd16efa17c2cd8ea05e0dea4 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 14:44:11 +0200 Subject: [PATCH 41/56] strip indexes as well --- pint_xarray/conversion.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/pint_xarray/conversion.py b/pint_xarray/conversion.py index fe86f578..913128ca 100644 --- a/pint_xarray/conversion.py +++ b/pint_xarray/conversion.py @@ -385,7 +385,12 @@ def strip_units_variable(var): def strip_units_dataset(obj): variables = {name: strip_units_variable(var) for name, var in obj.variables.items()} - return dataset_from_variables(variables, obj._coord_names, obj.attrs) + indexes = { + name: (index.index if isinstance(index, PintIndex) else index) + for name, index in obj.xindexes.items() + } + + return dataset_from_variables(variables, obj._coord_names, indexes, obj.attrs) def strip_units(obj): From 7303960977a8632fca980c88eddc08312c80a158 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 14:52:07 +0200 Subject: [PATCH 42/56] change the `.pint.to` tests to not include indexes These are already covered by the conversion tests, so no need to repeat. --- pint_xarray/tests/test_accessors.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/pint_xarray/tests/test_accessors.py b/pint_xarray/tests/test_accessors.py index 1028493f..5742874b 100644 --- a/pint_xarray/tests/test_accessors.py +++ b/pint_xarray/tests/test_accessors.py @@ -494,18 +494,20 @@ def test_roundtrip_data(self, example_unitless_ds): id="Dataset-incompatible units-data", ), pytest.param( - xr.Dataset(coords={"x": ("x", [2, 4], {"units": Unit("s")})}), + xr.Dataset(coords=xr.Coordinates({"x": Quantity([2, 4], "s")}, indexes={})), {"x": "ms"}, - xr.Dataset(coords={"x": ("x", [2000, 4000], {"units": Unit("ms")})}), + xr.Dataset( + coords=xr.Coordinates({"x": Quantity([2000, 4000], "ms")}, indexes={}) + ), None, - id="Dataset-compatible units-dims", + id="Dataset-compatible units-dims-no index", ), pytest.param( - xr.Dataset(coords={"x": ("x", [2, 4], {"units": Unit("s")})}), + xr.Dataset(coords=xr.Coordinates({"x": Quantity([2, 4], "s")}, indexes={})), {"x": "mm"}, None, ValueError, - id="Dataset-incompatible units-dims", + id="Dataset-incompatible units-dims-no index", ), pytest.param( xr.DataArray(Quantity([0, 1], "m"), dims="x"), @@ -537,25 +539,29 @@ def test_roundtrip_data(self, example_unitless_ds): ), pytest.param( xr.DataArray( - [0, 1], dims="x", coords={"x": ("x", [2, 4], {"units": Unit("s")})} + [0, 1], + dims="x", + coords=xr.Coordinates({"x": Quantity([2, 4], "s")}, indexes={}), ), {"x": "ms"}, xr.DataArray( [0, 1], dims="x", - coords={"x": ("x", [2000, 4000], {"units": Unit("ms")})}, + coords=xr.Coordinates({"x": Quantity([2000, 4000], "ms")}, indexes={}), ), None, - id="DataArray-compatible units-dims", + id="DataArray-compatible units-dims-no index", ), pytest.param( xr.DataArray( - [0, 1], dims="x", coords={"x": ("x", [2, 4], {"units": Unit("s")})} + [0, 1], + dims="x", + coords=xr.Coordinates({"x": Quantity([2, 4], "s")}, indexes={}), ), {"x": "mm"}, None, ValueError, - id="DataArray-incompatible units-dims", + id="DataArray-incompatible units-dims-no index", ), ), ) From e88738d02b5118e65b6944fdadaa423b580bdcfd Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 15:02:49 +0200 Subject: [PATCH 43/56] extract the units of `other` in `.pint.interp_like` --- pint_xarray/accessors.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pint_xarray/accessors.py b/pint_xarray/accessors.py index 82573451..882c59d8 100644 --- a/pint_xarray/accessors.py +++ b/pint_xarray/accessors.py @@ -751,7 +751,7 @@ def interp_like(self, other, method="linear", assume_sorted=False, kwargs=None): xarray.DataArray.pint.interp xarray.DataArray.interp_like """ - indexer_units = conversion.extract_unit_attributes(other) + indexer_units = conversion.extract_units(other) converted = conversion.convert_units(self.da, indexer_units) units = conversion.extract_units(converted) @@ -1514,7 +1514,7 @@ def interp_like(self, other, method="linear", assume_sorted=False, kwargs=None): xarray.Dataset.pint.interp xarray.Dataset.interp_like """ - indexer_units = conversion.extract_unit_attributes(other) + indexer_units = conversion.extract_units(other) converted = conversion.convert_units(self.ds, indexer_units) units = conversion.extract_units(converted) From 0b400d0a14db409498aa9ef27a5a39fdc59abed4 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 15:06:31 +0200 Subject: [PATCH 44/56] quantify the input and expected data in the `reindex` tests --- pint_xarray/tests/test_accessors.py | 125 ++++++++++++++++------------ 1 file changed, 71 insertions(+), 54 deletions(-) diff --git a/pint_xarray/tests/test_accessors.py b/pint_xarray/tests/test_accessors.py index 5742874b..d5d9647e 100644 --- a/pint_xarray/tests/test_accessors.py +++ b/pint_xarray/tests/test_accessors.py @@ -1154,72 +1154,73 @@ def test_chunk(obj): @pytest.mark.parametrize( - ["obj", "indexers", "expected", "error"], + ["obj", "units", "indexers", "expected", "expected_units", "error"], ( pytest.param( - xr.Dataset( - { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - } - ), + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, {"x": Quantity([10, 30, 50], "dm"), "y": Quantity([0, 120, 240], "s")}, - xr.Dataset( - { - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 120, 240], {"units": unit_registry.Unit("s")}), - } - ), + xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), + {"x": "dm", "y": "s"}, None, id="Dataset-identical units", ), pytest.param( - xr.Dataset( - { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - } - ), + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, {"x": Quantity([0, 1, 3, 5], "m"), "y": Quantity([0, 2, 4], "min")}, - xr.Dataset( - { - "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2, 4], {"units": unit_registry.Unit("min")}), - } - ), + xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), + {"x": "m", "y": "min"}, None, id="Dataset-compatible units", ), + pytest.param( + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, + {"x": Quantity([1, 3], "s"), "y": Quantity([1], "m")}, + None, + {}, + ValueError, + id="Dataset-incompatible units", + ), pytest.param( xr.Dataset( { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), + "a": (("x", "y"), np.array([[0, 1], [2, 3], [4, 5]])), + "x": [10, 20, 30], + "y": [60, 120], } ), - {"x": Quantity([1, 3], "s"), "y": Quantity([1], "m")}, + {"a": "kg"}, + { + "x": [15, 25], + "y": [75, 105], + }, + xr.Dataset( + { + "a": (("x", "y"), np.array([[np.nan, np.nan], [np.nan, np.nan]])), + "x": [15, 25], + "y": [75, 105], + } + ), + {"a": "kg"}, None, - ValueError, - id="Dataset-incompatible units", + id="Dataset-data units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), + {"x": "dm", "y": "s"}, {"x": Quantity([10, 30, 50], "dm"), "y": Quantity([0, 240], "s")}, xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), - coords={ - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 240], {"units": unit_registry.Unit("s")}), - }, + coords={"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}, ), + {"x": "dm", "y": "s"}, None, id="DataArray-identical units", ), @@ -1227,20 +1228,16 @@ def test_chunk(obj): xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), + {"x": "dm", "y": "s"}, {"x": Quantity([1, 3, 5], "m"), "y": Quantity([0, 2], "min")}, xr.DataArray( [[np.nan, 1], [np.nan, 5], [np.nan, np.nan]], dims=("x", "y"), - coords={ - "x": ("x", [1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2], {"units": unit_registry.Unit("min")}), - }, + coords={"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}, ), + {"x": "m", "y": "min"}, None, id="DataArray-compatible units", ), @@ -1248,26 +1245,46 @@ def test_chunk(obj): xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), + {"x": "dm", "y": "s"}, {"x": Quantity([10, 30], "s"), "y": Quantity([60], "m")}, None, + {}, ValueError, id="DataArray-incompatible units", ), + pytest.param( + xr.DataArray( + np.array([[0, 1], [2, 3], [4, 5]]), + dims=("x", "y"), + coords={"x": [10, 20, 30], "y": [60, 120]}, + ), + {None: "kg"}, + {"x": [15, 25], "y": [75, 105]}, + xr.DataArray( + [[np.nan, np.nan], [np.nan, np.nan]], + dims=("x", "y"), + coords={"x": [15, 25], "y": [75, 105]}, + ), + {None: "kg"}, + None, + id="DataArray-data units", + ), ), ) -def test_reindex(obj, indexers, expected, error): +def test_reindex(obj, units, indexers, expected, expected_units, error): + obj_ = obj.pint.quantify(units) + if expected is not None: + expected_ = expected.pint.quantify(expected_units) + if error is not None: with pytest.raises(error): obj.pint.reindex(indexers) else: - actual = obj.pint.reindex(indexers) - assert_units_equal(actual, expected) - assert_identical(actual, expected) + actual = obj_.pint.reindex(indexers) + assert_units_equal(actual, expected_) + assert_identical(actual, expected_) @pytest.mark.parametrize( From 5bd3ec7304016dd5be8a9045ac9bbe171a91c88a Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 15:16:15 +0200 Subject: [PATCH 45/56] remove the left-over explicit quantification in the `interp` tests --- pint_xarray/tests/test_accessors.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pint_xarray/tests/test_accessors.py b/pint_xarray/tests/test_accessors.py index d5d9647e..fd8f7896 100644 --- a/pint_xarray/tests/test_accessors.py +++ b/pint_xarray/tests/test_accessors.py @@ -1545,7 +1545,7 @@ def test_reindex_like(obj, other, expected, error): {None: "kg"}, {"x": [15, 25], "y": [75, 105]}, xr.DataArray( - Quantity([[1.25, 1.75], [3.25, 3.75]], "kg"), + [[1.25, 1.75], [3.25, 3.75]], dims=("x", "y"), coords={"x": [15, 25], "y": [75, 105]}, ), From 77eef6df1a14eaca5d6256a3e301daa363fb1f73 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 15:20:31 +0200 Subject: [PATCH 46/56] get `.pint.reindex` to work by explicitly converting, stripping, and then reattaching --- pint_xarray/accessors.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/pint_xarray/accessors.py b/pint_xarray/accessors.py index 882c59d8..cb8fb3cd 100644 --- a/pint_xarray/accessors.py +++ b/pint_xarray/accessors.py @@ -647,17 +647,19 @@ def reindex( # convert the indexes to the indexer's units converted = conversion.convert_units(self.da, indexer_units) + converted_units = conversion.extract_units(converted) + stripped = conversion.strip_units(converted) # index stripped_indexers = conversion.strip_indexer_units(indexers) - indexed = converted.reindex( + indexed = stripped.reindex( stripped_indexers, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) - return indexed + return conversion.attach_units(indexed, converted_units) def reindex_like( self, other, method=None, tolerance=None, copy=True, fill_value=NA @@ -1410,17 +1412,19 @@ def reindex( # convert the indexes to the indexer's units converted = conversion.convert_units(self.ds, indexer_units) + converted_units = conversion.extract_units(converted) + stripped = conversion.strip_units(converted) # index stripped_indexers = conversion.strip_indexer_units(indexers) - indexed = converted.reindex( + indexed = stripped.reindex( stripped_indexers, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) - return indexed + return conversion.attach_units(indexed, converted_units) def reindex_like( self, other, method=None, tolerance=None, copy=True, fill_value=NA From c38eb5a535f9828fb98172890783142085a82ad5 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 15:28:46 +0200 Subject: [PATCH 47/56] quantify the input and expected objects in the `reindex_like` tests --- pint_xarray/tests/test_accessors.py | 176 ++++++++++++++-------------- 1 file changed, 85 insertions(+), 91 deletions(-) diff --git a/pint_xarray/tests/test_accessors.py b/pint_xarray/tests/test_accessors.py index fd8f7896..da6e59e1 100644 --- a/pint_xarray/tests/test_accessors.py +++ b/pint_xarray/tests/test_accessors.py @@ -1288,66 +1288,35 @@ def test_reindex(obj, units, indexers, expected, expected_units, error): @pytest.mark.parametrize( - ["obj", "other", "expected", "error"], + ["obj", "units", "other", "other_units", "expected", "expected_units", "error"], ( pytest.param( - xr.Dataset( - { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - } - ), - xr.Dataset( - { - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 120, 240], {"units": unit_registry.Unit("s")}), - } - ), - xr.Dataset( - { - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 120, 240], {"units": unit_registry.Unit("s")}), - } - ), + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), + {"x": "dm", "y": "s"}, None, id="Dataset-identical units", ), pytest.param( - xr.Dataset( - { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - } - ), - xr.Dataset( - { - "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2, 4], {"units": unit_registry.Unit("min")}), - } - ), - xr.Dataset( - { - "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2, 4], {"units": unit_registry.Unit("min")}), - } - ), + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), + {"x": "m", "y": "min"}, + xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), + {"x": "m", "y": "min"}, None, id="Dataset-compatible units", ), pytest.param( - xr.Dataset( - { - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - } - ), - xr.Dataset( - { - "x": ("x", [1, 3], {"units": unit_registry.Unit("s")}), - "y": ("y", [1], {"units": unit_registry.Unit("m")}), - } - ), + xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [1, 3]), "y": ("y", [1])}), + {"x": "s", "y": "m"}, None, + {}, ValueError, id="Dataset-incompatible units", ), @@ -1355,51 +1324,57 @@ def test_reindex(obj, units, indexers, expected, expected_units, error): xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, - ), - xr.Dataset( - { - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 240], {"units": unit_registry.Unit("s")}), - } + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}), + {"x": "dm", "y": "s"}, xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), - coords={ - "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), - "y": ("y", [0, 240], {"units": unit_registry.Unit("s")}), - }, + coords={"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}, ), + {"x": "dm", "y": "s"}, None, id="DataArray-identical units", ), pytest.param( - xr.DataArray( - [[0, 1], [2, 3], [4, 5]], - dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, + xr.Dataset( + { + "a": (("x", "y"), [[0, 1], [2, 3], [4, 5]]), + "x": [10, 20, 30], + "y": [60, 120], + } ), + {"a": "kg"}, + xr.Dataset({"x": [15, 25], "y": [75, 105]}), + {}, xr.Dataset( { - "x": ("x", [1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2], {"units": unit_registry.Unit("min")}), + "a": (("x", "y"), [[np.nan, np.nan], [np.nan, np.nan]]), + "x": [15, 25], + "y": [75, 105], } ), + {"a": "kg"}, + None, + id="Dataset-data units", + ), + pytest.param( + xr.DataArray( + [[0, 1], [2, 3], [4, 5]], + dims=("x", "y"), + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, + ), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}), + {"x": "m", "y": "min"}, xr.DataArray( [[np.nan, 1], [np.nan, 5], [np.nan, np.nan]], dims=("x", "y"), - coords={ - "x": ("x", [1, 3, 5], {"units": unit_registry.Unit("m")}), - "y": ("y", [0, 2], {"units": unit_registry.Unit("min")}), - }, + coords={"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}, ), + {"x": "m", "y": "min"}, None, id="DataArray-compatible units", ), @@ -1407,31 +1382,50 @@ def test_reindex(obj, units, indexers, expected, expected_units, error): xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), - coords={ - "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), - "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), - }, - ), - xr.Dataset( - { - "x": ("x", [10, 30], {"units": unit_registry.Unit("s")}), - "y": ("y", [60], {"units": unit_registry.Unit("m")}), - } + coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), + {"x": "dm", "y": "s"}, + xr.Dataset({"x": ("x", [10, 30]), "y": ("y", [60])}), + {"x": "s", "y": "m"}, None, + {}, ValueError, id="DataArray-incompatible units", ), + pytest.param( + xr.DataArray( + [[0, 1], [2, 3], [4, 5]], + dims=("x", "y"), + coords={"x": [10, 20, 30], "y": [60, 120]}, + ), + {"a": "kg"}, + xr.Dataset({"x": [15, 25], "y": [75, 105]}), + {}, + xr.DataArray( + [[np.nan, np.nan], [np.nan, np.nan]], + dims=("x", "y"), + coords={"x": [15, 25], "y": [75, 105]}, + ), + {"a": "kg"}, + None, + id="DataArray-data units", + ), ), ) -def test_reindex_like(obj, other, expected, error): +def test_reindex_like(obj, units, other, other_units, expected, expected_units, error): + obj_ = obj.pint.quantify(units) + other_ = other.pint.quantify(other_units) + + if expected is not None: + expected_ = expected.pint.quantify(expected_units) + if error is not None: with pytest.raises(error): - obj.pint.reindex_like(other) + obj_.pint.reindex_like(other_) else: - actual = obj.pint.reindex_like(other) - assert_units_equal(actual, expected) - assert_identical(actual, expected) + actual = obj_.pint.reindex_like(other_) + assert_units_equal(actual, expected_) + assert_identical(actual, expected_) @requires_scipy From 7277eb52fca9cbfc23f3077458aa995fe388b091 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 15:29:37 +0200 Subject: [PATCH 48/56] get `reindex_like` to work with indexes --- pint_xarray/accessors.py | 26 ++++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/pint_xarray/accessors.py b/pint_xarray/accessors.py index cb8fb3cd..560ffccf 100644 --- a/pint_xarray/accessors.py +++ b/pint_xarray/accessors.py @@ -679,19 +679,24 @@ def reindex_like( xarray.DataArray.pint.reindex xarray.DataArray.reindex_like """ - indexer_units = conversion.extract_unit_attributes(other) + indexer_units = conversion.extract_units(other) + + converted = conversion.convert_units(self.da, indexer_units) + units = conversion.extract_units(converted) + stripped = conversion.strip_units(converted) + stripped_other = conversion.strip_units(other) # TODO: handle tolerance # TODO: handle fill_value - converted = conversion.convert_units(self.da, indexer_units) - return converted.reindex_like( - other, + reindexed = stripped.reindex_like( + stripped_other, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) + return conversion.attach_units(reindexed, units) def interp( self, @@ -1444,19 +1449,24 @@ def reindex_like( xarray.Dataset.pint.reindex xarray.Dataset.reindex_like """ - indexer_units = conversion.extract_unit_attributes(other) + indexer_units = conversion.extract_units(other) + + converted = conversion.convert_units(self.ds, indexer_units) + units = conversion.extract_units(converted) + stripped = conversion.strip_units(converted) + stripped_other = conversion.strip_units(other) # TODO: handle tolerance # TODO: handle fill_value - converted = conversion.convert_units(self.ds, indexer_units) - return converted.reindex_like( - other, + reindexed = stripped.reindex_like( + stripped_other, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) + return conversion.attach_units(reindexed, units) def interp( self, From c7cf3406e8d07f1e261ace478a68263da2a12450 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 16:09:39 +0200 Subject: [PATCH 49/56] quantify expected only if we expect to make use of it --- pint_xarray/tests/test_accessors.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/pint_xarray/tests/test_accessors.py b/pint_xarray/tests/test_accessors.py index da6e59e1..59a6b9d8 100644 --- a/pint_xarray/tests/test_accessors.py +++ b/pint_xarray/tests/test_accessors.py @@ -1275,13 +1275,13 @@ def test_chunk(obj): ) def test_reindex(obj, units, indexers, expected, expected_units, error): obj_ = obj.pint.quantify(units) - if expected is not None: - expected_ = expected.pint.quantify(expected_units) if error is not None: with pytest.raises(error): obj.pint.reindex(indexers) else: + expected_ = expected.pint.quantify(expected_units) + actual = obj_.pint.reindex(indexers) assert_units_equal(actual, expected_) assert_identical(actual, expected_) @@ -1416,13 +1416,12 @@ def test_reindex_like(obj, units, other, other_units, expected, expected_units, obj_ = obj.pint.quantify(units) other_ = other.pint.quantify(other_units) - if expected is not None: - expected_ = expected.pint.quantify(expected_units) - if error is not None: with pytest.raises(error): obj_.pint.reindex_like(other_) else: + expected_ = expected.pint.quantify(expected_units) + actual = obj_.pint.reindex_like(other_) assert_units_equal(actual, expected_) assert_identical(actual, expected_) @@ -1551,13 +1550,13 @@ def test_reindex_like(obj, units, other, other_units, expected, expected_units, ) def test_interp(obj, units, indexers, expected, expected_units, error): obj_ = obj.pint.quantify(units) - if expected is not None: - expected_ = expected.pint.quantify(expected_units) if error is not None: with pytest.raises(error): obj.pint.interp(indexers) else: + expected_ = expected.pint.quantify(expected_units) + actual = obj_.pint.interp(indexers) assert_units_equal(actual, expected_) assert_identical(actual, expected_) @@ -1693,13 +1692,12 @@ def test_interp_like(obj, units, other, other_units, expected, expected_units, e obj_ = obj.pint.quantify(units) other_ = other.pint.quantify(other_units) - if expected is not None: - expected_ = expected.pint.quantify(expected_units) - if error is not None: with pytest.raises(error): obj_.pint.interp_like(other_) else: + expected_ = expected.pint.quantify(expected_units) + actual = obj_.pint.interp_like(other_) assert_units_equal(actual, expected_) assert_identical(actual, expected_) From 948d20f614772d65d3946303faebd3b14c5f9a3c Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 16:10:30 +0200 Subject: [PATCH 50/56] quantify input and expected objects in the `sel` and `loc` tests --- pint_xarray/tests/test_accessors.py | 24 ++++++++++++++++-------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/pint_xarray/tests/test_accessors.py b/pint_xarray/tests/test_accessors.py index 59a6b9d8..b33c7866 100644 --- a/pint_xarray/tests/test_accessors.py +++ b/pint_xarray/tests/test_accessors.py @@ -684,13 +684,17 @@ def test_to(obj, units, expected, error): ), ) def test_sel(obj, indexers, expected, error): + obj_ = obj.pint.quantify() + if error is not None: with pytest.raises(error): - obj.pint.sel(indexers) + obj_.pint.sel(indexers) else: - actual = obj.pint.sel(indexers) - assert_units_equal(actual, expected) - assert_identical(actual, expected) + expected_ = expected.pint.quantify() + + actual = obj_.pint.sel(indexers) + assert_units_equal(actual, expected_) + assert_identical(actual, expected_) @pytest.mark.parametrize( @@ -801,13 +805,17 @@ def test_sel(obj, indexers, expected, error): ), ) def test_loc(obj, indexers, expected, error): + obj_ = obj.pint.quantify() + if error is not None: with pytest.raises(error): - obj.pint.loc[indexers] + obj_.pint.loc[indexers] else: - actual = obj.pint.loc[indexers] - assert_units_equal(actual, expected) - assert_identical(actual, expected) + expected_ = expected.pint.quantify() + + actual = obj_.pint.loc[indexers] + assert_units_equal(actual, expected_) + assert_identical(actual, expected_) @pytest.mark.parametrize( From 8c76cbc5fc8fe9cabe2f0683b67d3e80981607be Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 16:11:05 +0200 Subject: [PATCH 51/56] get `.pint.sel` and `.pint.loc` to work with the indexes In theory, we could also use `sel` and `loc` directly, but that would not allow us to change the units of the result (or at least, as far as I can tell). --- pint_xarray/accessors.py | 28 ++++++++++++++++++++++------ 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/pint_xarray/accessors.py b/pint_xarray/accessors.py index 560ffccf..caf96f45 100644 --- a/pint_xarray/accessors.py +++ b/pint_xarray/accessors.py @@ -165,7 +165,12 @@ def __getitem__(self, indexers): # index stripped_indexers = conversion.strip_indexer_units(indexers) - return converted.loc[stripped_indexers] + + stripped = conversion.strip_units(converted) + converted_units = conversion.extract_units(converted) + indexed = stripped.loc[stripped_indexers] + + return conversion.attach_units(indexed, converted_units) class DataArrayLocIndexer: @@ -192,7 +197,12 @@ def __getitem__(self, indexers): # index stripped_indexers = conversion.strip_indexer_units(indexers) - return converted.loc[stripped_indexers] + + stripped = conversion.strip_units(converted) + converted_units = conversion.extract_units(converted) + indexed = stripped.loc[stripped_indexers] + + return conversion.attach_units(indexed, converted_units) def __setitem__(self, indexers, values): if not is_dict_like(indexers): @@ -808,14 +818,17 @@ def sel( # index stripped_indexers = conversion.strip_indexer_units(indexers) - indexed = converted.sel( + + stripped = conversion.strip_units(converted) + converted_units = conversion.extract_units(converted) + indexed = stripped.sel( stripped_indexers, method=method, tolerance=tolerance, drop=drop, ) - return indexed + return conversion.attach_units(indexed, converted_units) @property def loc(self): @@ -1578,14 +1591,17 @@ def sel( # index stripped_indexers = conversion.strip_indexer_units(indexers) - indexed = converted.sel( + + stripped = conversion.strip_units(converted) + converted_units = conversion.extract_units(converted) + indexed = stripped.sel( stripped_indexers, method=method, tolerance=tolerance, drop=drop, ) - return indexed + return conversion.attach_units(indexed, converted_units) @property def loc(self): From f9cb15c04bed44758d84ef89d4088fb358077b25 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 16:21:05 +0200 Subject: [PATCH 52/56] remove the warning about indexed coordinates --- pint_xarray/accessors.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/pint_xarray/accessors.py b/pint_xarray/accessors.py index caf96f45..2104021b 100644 --- a/pint_xarray/accessors.py +++ b/pint_xarray/accessors.py @@ -251,12 +251,6 @@ def quantify(self, units=_default, unit_registry=None, **unit_kwargs): the data into memory. To avoid that, consider converting to ``dask`` first (e.g. using ``chunk``). - .. warning:: - - As units in dimension coordinates are not supported until - ``xarray`` changes the way it implements indexes, these - units will be set as attributes. - .. note:: Also note that datetime units (i.e. ones that match ``{units} since {date}``) in unit attributes will be @@ -979,12 +973,6 @@ def quantify(self, units=_default, unit_registry=None, **unit_kwargs): the data into memory. To avoid that, consider converting to ``dask`` first (e.g. using ``chunk``). - .. warning:: - - As units in dimension coordinates are not supported until - ``xarray`` changes the way it implements indexes, these - units will be set as attributes. - .. note:: Also note that datetime units (i.e. ones that match ``{units} since {date}``) in unit attributes will be From 49942bfa76f1b441a5ff09f3062a39279c48ba74 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Sat, 6 Jul 2024 16:23:51 +0200 Subject: [PATCH 53/56] preserve the order of the variables --- pint_xarray/conversion.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pint_xarray/conversion.py b/pint_xarray/conversion.py index 913128ca..a015633c 100644 --- a/pint_xarray/conversion.py +++ b/pint_xarray/conversion.py @@ -176,7 +176,8 @@ def attach_units_dataset(obj, units): if rejected_vars: raise ValueError(rejected_vars) - return dataset_from_variables(attached, obj._coord_names, indexes, obj.attrs) + reordered = {name: attached[name] for name in obj.variables.keys()} + return dataset_from_variables(reordered, obj._coord_names, indexes, obj.attrs) def attach_units(obj, units): @@ -305,7 +306,8 @@ def convert_units_dataset(obj, units): if failed: raise ValueError(failed) - return dataset_from_variables(converted, obj._coord_names, indexes, obj.attrs) + reordered = {name: converted[name] for name in obj.variables.keys()} + return dataset_from_variables(reordered, obj._coord_names, indexes, obj.attrs) def convert_units(obj, units): From 20dd15cdba7aaea15e65d1512ad30f2f9d9d93f6 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Mon, 8 Jul 2024 23:57:23 +0200 Subject: [PATCH 54/56] remove the remaining uses of `Coordinates._construct_direct` --- pint_xarray/conversion.py | 2 +- pint_xarray/tests/test_conversion.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pint_xarray/conversion.py b/pint_xarray/conversion.py index a015633c..5b801fb4 100644 --- a/pint_xarray/conversion.py +++ b/pint_xarray/conversion.py @@ -128,7 +128,7 @@ def dataset_from_variables(variables, coordinate_names, indexes, attrs): } coords = {name: var for name, var in variables.items() if name in coordinate_names} - new_coords = Coordinates._construct_direct(coords, indexes) + new_coords = Coordinates(coords, indexes=indexes) return Dataset(data_vars=data_vars, coords=new_coords, attrs=attrs) diff --git a/pint_xarray/tests/test_conversion.py b/pint_xarray/tests/test_conversion.py index 335cdc35..5101eac1 100644 --- a/pint_xarray/tests/test_conversion.py +++ b/pint_xarray/tests/test_conversion.py @@ -256,7 +256,7 @@ def test_attach_units(self, type, units): index = PintIndex(index=index, units=units.get("x")) obj = Dataset({"a": ("x", a), "b": ("x", b)}, coords={"u": ("x", u), "x": x}) - coords = Coordinates._construct_direct( + coords = Coordinates( coords={"u": Variable("x", q_u), "x": Variable("x", q_x)}, indexes={"x": index}, ) From 5efb318f58efefea5104e43c3b9b41ceacb5037d Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 9 Jul 2024 10:36:54 +0200 Subject: [PATCH 55/56] whats-new entry --- docs/whats-new.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/whats-new.rst b/docs/whats-new.rst index a90a8423..babbf532 100644 --- a/docs/whats-new.rst +++ b/docs/whats-new.rst @@ -6,6 +6,8 @@ What's new ------------------ - drop support for python 3.9 (:pull:`266`) By `Justus Magin `_. +- create a `PintIndex` to allow units on indexed coordinates (:pull:`163`, :issue:`162`) + By `Justus Magin `_ and `Benoit Bovy `_. 0.4 (23 Jun 2024) ----------------- From f53539ae5c686c0595dee53e3369eb003452a0b2 Mon Sep 17 00:00:00 2001 From: Justus Magin Date: Tue, 9 Jul 2024 10:47:16 +0200 Subject: [PATCH 56/56] expose the index --- pint_xarray/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pint_xarray/__init__.py b/pint_xarray/__init__.py index 3ce42d86..9991351b 100644 --- a/pint_xarray/__init__.py +++ b/pint_xarray/__init__.py @@ -5,6 +5,7 @@ from . import accessors, formatting, testing # noqa: F401 from .accessors import default_registry as unit_registry from .accessors import setup_registry +from .index import PintIndex try: __version__ = version("pint-xarray") @@ -21,4 +22,5 @@ "testing", "unit_registry", "setup_registry", + "PintIndex", ]