Skip to content

Commit

Permalink
interp and interp_like (#72)
Browse files Browse the repository at this point in the history 
* implement and test interp

* make the docstrings more accurate

* implement and test interp_like

* add both to api.rst

* update whats-new.rst

* add scipy to the CI requirements
  • Loading branch information
@keewis
keewis committed Mar 2, 2021
1 parent 1052105 commit de4a4a2
Show file tree
Hide file tree
Showing 5 changed files with 521 additions and 7 deletions.
1 change: 1 addition & 0 deletions ci/requirements.txt
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
pint
numpy
scipy
xarray
isort
black
Expand Down
4 changes: 4 additions & 0 deletions docs/api.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
3 changes: 3 additions & 0 deletions docs/whats-new.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
----------------------
Expand Down
256 changes: 249 additions & 7 deletions pint_xarray/accessors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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::
Expand Down Expand Up @@ -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::
Expand Down Expand Up @@ -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::
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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::
Expand Down Expand Up @@ -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::
Expand Down Expand Up @@ -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::
Expand Down

0 comments on commit de4a4a2

Please sign in to comment.