interp and interp_like

ci/requirements.txt

@@ -1,5 +1,6 @@
 pint
 numpy
+scipy
 xarray
 isort
 black

docs/api.rst

@@ -15,6 +15,8 @@ Dataset
 
    xarray.Dataset.pint.quantify
    xarray.Dataset.pint.dequantify
+   xarray.Dataset.pint.interp
+   xarray.Dataset.pint.interp_like
    xarray.Dataset.pint.reindex
    xarray.Dataset.pint.reindex_like
    xarray.Dataset.pint.sel
@@ -37,6 +39,8 @@ DataArray
 
    xarray.DataArray.pint.quantify
    xarray.DataArray.pint.dequantify
+   xarray.DataArray.pint.interp
+   xarray.DataArray.pint.interp_like
    xarray.DataArray.pint.reindex
    xarray.DataArray.pint.reindex_like
    xarray.DataArray.pint.sel

docs/whats-new.rst

@@ -28,6 +28,9 @@ What's new
 - implement :py:meth:`Dataset.pint.reindex`, :py:meth:`Dataset.pint.reindex_like`,
   :py:meth:`DataArray.pint.reindex` and :py:meth:`DataArray.pint.reindex_like` (:pull:`69`).
   By `Justus Magin <https://github.com/keewis>`_.
+- implement :py:meth:`Dataset.pint.interp`, :py:meth:`Dataset.pint.interp_like`,
+  :py:meth:`DataArray.pint.interp` and :py:meth:`DataArray.pint.interp_like` (:pull:`72`).
+  By `Justus Magin <https://github.com/keewis>`_.
 
 v0.1 (October 26 2020)
 ----------------------

pint_xarray/accessors.py

@@ -422,7 +422,7 @@ def reindex(
     ):
         """unit-aware version of reindex
 
-        Just like :py:meth:`xarray.DataArray.reindex`, except the dataset's indexes are converted
+        Just like :py:meth:`xarray.DataArray.reindex`, except the object's indexes are converted
         to the units of the indexers first.
 
         .. note::
@@ -490,7 +490,7 @@ def reindex_like(
     ):
         """unit-aware version of reindex_like
 
-        Just like :py:meth:`xarray.DataArray.reindex_like`, except the dataset's indexes are converted
+        Just like :py:meth:`xarray.DataArray.reindex_like`, except the object's indexes are converted
         to the units of the indexers first.
 
         .. note::
@@ -540,12 +540,133 @@ def reindex_like(
             fill_value=fill_value,
         )
 
+    def interp(
+        self,
+        coords=None,
+        method="linear",
+        assume_sorted=False,
+        kwargs=None,
+        **coords_kwargs,
+    ):
+        """unit-aware version of interp
+
+        Just like :py:meth:`xarray.DataArray.interp`, except the object's indexes are converted
+        to the units of the indexers first.
+
+        .. note::
+            ``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to
+            ``DataArray.interp`` unmodified.
+
+        See Also
+        --------
+        xarray.Dataset.pint.interp
+        xarray.DataArray.pint.interp_like
+        xarray.DataArray.interp
+        """
+        indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp")
+
+        indexer_units = {
+            name: conversion.extract_indexer_units(indexer)
+            for name, indexer in indexers.items()
+        }
+
+        # make sure we only have compatible units
+        dims = self.da.dims
+        unit_attrs = conversion.extract_unit_attributes(self.da)
+        index_units = {
+            name: units for name, units in unit_attrs.items() if name in dims
+        }
+
+        registry = get_registry(None, index_units, indexer_units)
+
+        units = zip_mappings(indexer_units, index_units)
+        incompatible_units = [
+            key
+            for key, (indexer_unit, index_unit) in units.items()
+            if (
+                None not in (indexer_unit, index_unit)
+                and not registry.is_compatible_with(indexer_unit, index_unit)
+            )
+        ]
+        if incompatible_units:
+            units1 = {key: indexer_units[key] for key in incompatible_units}
+            units2 = {key: index_units[key] for key in incompatible_units}
+            raise DimensionalityError(units1, units2)
+
+        # convert the indexes to the indexer's units
+        converted = conversion.convert_units(self.da, indexer_units)
+        stripped = conversion.strip_units(converted)
+
+        # index
+        stripped_indexers = {
+            name: conversion.strip_indexer_units(indexer)
+            for name, indexer in indexers.items()
+        }
+        interpolated = stripped.interp(
+            stripped_indexers,
+            method=method,
+            assume_sorted=False,
+            kwargs=None,
+        )
+        return conversion.attach_units(interpolated, indexer_units)
+
+    def interp_like(self, other, method="linear", assume_sorted=False, kwargs=None):
+        """unit-aware version of interp_like
+
+        Just like :py:meth:`xarray.DataArray.interp_like`, except the object's indexes are converted
+        to the units of the indexers first.
+
+        .. note::
+            ``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to
+            ``DataArray.interp_like`` unmodified.
+
+        See Also
+        --------
+        xarray.Dataset.pint.interp_like
+        xarray.DataArray.pint.interp
+        xarray.DataArray.interp_like
+        """
+        indexer_units = conversion.extract_unit_attributes(other)
+
+        # make sure we only have compatible units
+        dims = self.da.dims
+        unit_attrs = conversion.extract_unit_attributes(self.da)
+        index_units = {
+            name: units for name, units in unit_attrs.items() if name in dims
+        }
+
+        registry = get_registry(None, index_units, indexer_units)
+
+        units = zip_mappings(indexer_units, index_units)
+        incompatible_units = [
+            key
+            for key, (indexer_unit, index_unit) in units.items()
+            if (
+                None not in (indexer_unit, index_unit)
+                and not registry.is_compatible_with(indexer_unit, index_unit)
+            )
+        ]
+        if incompatible_units:
+            units1 = {key: indexer_units[key] for key in incompatible_units}
+            units2 = {key: index_units[key] for key in incompatible_units}
+            raise DimensionalityError(units1, units2)
+
+        converted = conversion.convert_units(self.da, indexer_units)
+        stripped = conversion.strip_units(converted)
+        interpolated = stripped.interp_like(
+            other,
+            method=method,
+            assume_sorted=assume_sorted,
+            kwargs=kwargs,
+        )
+        return conversion.attach_units(interpolated, indexer_units)
+
     def sel(
         self, indexers=None, method=None, tolerance=None, drop=False, **indexers_kwargs
     ):
         """unit-aware version of sel
 
-        Just like :py:meth:`xarray.DataArray.sel`, except the dataset's indexes are converted
+        Just like :py:meth:`xarray.DataArray.sel`, except the object's indexes are converted
         to the units of the indexers first.
 
         .. note::
@@ -743,7 +864,7 @@ def to(self, units=None, **unit_kwargs):
         ----------
         units : unit-like or mapping of hashable to unit-like, optional
             The units to convert to. If a unit name or ``pint.Unit``
-            object, convert all the Dataset's data variables. If a dict-like, it
+            object, convert all the object's data variables. If a dict-like, it
             maps variable names to unit names or ``pint.Unit``
             objects.
         **unit_kwargs
@@ -891,7 +1012,7 @@ def reindex(
     ):
         """unit-aware version of reindex
 
-        Just like :py:meth:`xarray.Dataset.reindex`, except the dataset's indexes are converted
+        Just like :py:meth:`xarray.Dataset.reindex`, except the object's indexes are converted
         to the units of the indexers first.
 
         .. note::
@@ -959,7 +1080,7 @@ def reindex_like(
     ):
         """unit-aware version of reindex_like
 
-        Just like :py:meth:`xarray.Dataset.reindex_like`, except the dataset's indexes are converted
+        Just like :py:meth:`xarray.Dataset.reindex_like`, except the object's indexes are converted
         to the units of the indexers first.
 
         .. note::
@@ -1009,12 +1130,133 @@ def reindex_like(
             fill_value=fill_value,
         )
 
+    def interp(
+        self,
+        coords=None,
+        method="linear",
+        assume_sorted=False,
+        kwargs=None,
+        **coords_kwargs,
+    ):
+        """unit-aware version of interp
+
+        Just like :py:meth:`xarray.Dataset.interp`, except the object's indexes are converted
+        to the units of the indexers first.
+
+        .. note::
+            ``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to
+            ``Dataset.interp`` unmodified.
+
+        See Also
+        --------
+        xarray.DataArray.pint.interp
+        xarray.Dataset.pint.interp_like
+        xarray.Dataset.interp
+        """
+        indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp")
+
+        indexer_units = {
+            name: conversion.extract_indexer_units(indexer)
+            for name, indexer in indexers.items()
+        }
+
+        # make sure we only have compatible units
+        dims = self.ds.dims
+        unit_attrs = conversion.extract_unit_attributes(self.ds)
+        index_units = {
+            name: units for name, units in unit_attrs.items() if name in dims
+        }
+
+        registry = get_registry(None, index_units, indexer_units)
+
+        units = zip_mappings(indexer_units, index_units)
+        incompatible_units = [
+            key
+            for key, (indexer_unit, index_unit) in units.items()
+            if (
+                None not in (indexer_unit, index_unit)
+                and not registry.is_compatible_with(indexer_unit, index_unit)
+            )
+        ]
+        if incompatible_units:
+            units1 = {key: indexer_units[key] for key in incompatible_units}
+            units2 = {key: index_units[key] for key in incompatible_units}
+            raise DimensionalityError(units1, units2)
+
+        # convert the indexes to the indexer's units
+        converted = conversion.convert_units(self.ds, indexer_units)
+        stripped = conversion.strip_units(converted)
+
+        # index
+        stripped_indexers = {
+            name: conversion.strip_indexer_units(indexer)
+            for name, indexer in indexers.items()
+        }
+        interpolated = stripped.interp(
+            stripped_indexers,
+            method=method,
+            assume_sorted=False,
+            kwargs=None,
+        )
+        return conversion.attach_units(interpolated, indexer_units)
+
+    def interp_like(self, other, method="linear", assume_sorted=False, kwargs=None):
+        """unit-aware version of interp_like
+
+        Just like :py:meth:`xarray.Dataset.interp_like`, except the object's indexes are converted
+        to the units of the indexers first.
+
+        .. note::
+            ``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to
+            ``Dataset.interp_like`` unmodified.
+
+        See Also
+        --------
+        xarray.DataArray.pint.interp_like
+        xarray.Dataset.pint.interp
+        xarray.Dataset.interp_like
+        """
+        indexer_units = conversion.extract_unit_attributes(other)
+
+        # make sure we only have compatible units
+        dims = self.ds.dims
+        unit_attrs = conversion.extract_unit_attributes(self.ds)
+        index_units = {
+            name: units for name, units in unit_attrs.items() if name in dims
+        }
+
+        registry = get_registry(None, index_units, indexer_units)
+
+        units = zip_mappings(indexer_units, index_units)
+        incompatible_units = [
+            key
+            for key, (indexer_unit, index_unit) in units.items()
+            if (
+                None not in (indexer_unit, index_unit)
+                and not registry.is_compatible_with(indexer_unit, index_unit)
+            )
+        ]
+        if incompatible_units:
+            units1 = {key: indexer_units[key] for key in incompatible_units}
+            units2 = {key: index_units[key] for key in incompatible_units}
+            raise DimensionalityError(units1, units2)
+
+        converted = conversion.convert_units(self.ds, indexer_units)
+        stripped = conversion.strip_units(converted)
+        interpolated = stripped.interp_like(
+            other,
+            method=method,
+            assume_sorted=assume_sorted,
+            kwargs=kwargs,
+        )
+        return conversion.attach_units(interpolated, indexer_units)
+
     def sel(
         self, indexers=None, method=None, tolerance=None, drop=False, **indexers_kwargs
     ):
         """unit-aware version of sel
 
-        Just like :py:meth:`xarray.Dataset.sel`, except the dataset's indexes are converted to the units
+        Just like :py:meth:`xarray.Dataset.sel`, except the object's indexes are converted to the units
         of the indexers first.
 
         .. note::