Skip to content

Commit de4a4a2

Browse files
keewiskeewis
authored
interp and interp_like (#72)
* 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
1 parent 1052105 commit de4a4a2

File tree

5 files changed

+521
-7
lines changed

5 files changed

+521
-7
lines changed

‎ci/requirements.txt

+1
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
pint
22
numpy
3+
scipy
34
xarray
45
isort
56
black

‎docs/api.rst

+4
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ Dataset
1515

1616
xarray.Dataset.pint.quantify
1717
xarray.Dataset.pint.dequantify
18+
xarray.Dataset.pint.interp
19+
xarray.Dataset.pint.interp_like
1820
xarray.Dataset.pint.reindex
1921
xarray.Dataset.pint.reindex_like
2022
xarray.Dataset.pint.sel
@@ -37,6 +39,8 @@ DataArray
3739

3840
xarray.DataArray.pint.quantify
3941
xarray.DataArray.pint.dequantify
42+
xarray.DataArray.pint.interp
43+
xarray.DataArray.pint.interp_like
4044
xarray.DataArray.pint.reindex
4145
xarray.DataArray.pint.reindex_like
4246
xarray.DataArray.pint.sel

‎docs/whats-new.rst

+3
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,9 @@ What's new
2828
- implement :py:meth:`Dataset.pint.reindex`, :py:meth:`Dataset.pint.reindex_like`,
2929
:py:meth:`DataArray.pint.reindex` and :py:meth:`DataArray.pint.reindex_like` (:pull:`69`).
3030
By `Justus Magin <https://github.com/keewis>`_.
31+
- implement :py:meth:`Dataset.pint.interp`, :py:meth:`Dataset.pint.interp_like`,
32+
:py:meth:`DataArray.pint.interp` and :py:meth:`DataArray.pint.interp_like` (:pull:`72`).
33+
By `Justus Magin <https://github.com/keewis>`_.
3134

3235
v0.1 (October 26 2020)
3336
----------------------

‎pint_xarray/accessors.py

+249-7
Original file line numberDiff line numberDiff line change
@@ -422,7 +422,7 @@ def reindex(
422422
):
423423
"""unit-aware version of reindex
424424
425-
Just like :py:meth:`xarray.DataArray.reindex`, except the dataset's indexes are converted
425+
Just like :py:meth:`xarray.DataArray.reindex`, except the object's indexes are converted
426426
to the units of the indexers first.
427427
428428
.. note::
@@ -490,7 +490,7 @@ def reindex_like(
490490
):
491491
"""unit-aware version of reindex_like
492492
493-
Just like :py:meth:`xarray.DataArray.reindex_like`, except the dataset's indexes are converted
493+
Just like :py:meth:`xarray.DataArray.reindex_like`, except the object's indexes are converted
494494
to the units of the indexers first.
495495
496496
.. note::
@@ -540,12 +540,133 @@ def reindex_like(
540540
fill_value=fill_value,
541541
)
542542

543+
def interp(
544+
self,
545+
coords=None,
546+
method="linear",
547+
assume_sorted=False,
548+
kwargs=None,
549+
**coords_kwargs,
550+
):
551+
"""unit-aware version of interp
552+
553+
Just like :py:meth:`xarray.DataArray.interp`, except the object's indexes are converted
554+
to the units of the indexers first.
555+
556+
.. note::
557+
``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to
558+
``DataArray.interp`` unmodified.
559+
560+
See Also
561+
--------
562+
xarray.Dataset.pint.interp
563+
xarray.DataArray.pint.interp_like
564+
xarray.DataArray.interp
565+
"""
566+
indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp")
567+
568+
indexer_units = {
569+
name: conversion.extract_indexer_units(indexer)
570+
for name, indexer in indexers.items()
571+
}
572+
573+
# make sure we only have compatible units
574+
dims = self.da.dims
575+
unit_attrs = conversion.extract_unit_attributes(self.da)
576+
index_units = {
577+
name: units for name, units in unit_attrs.items() if name in dims
578+
}
579+
580+
registry = get_registry(None, index_units, indexer_units)
581+
582+
units = zip_mappings(indexer_units, index_units)
583+
incompatible_units = [
584+
key
585+
for key, (indexer_unit, index_unit) in units.items()
586+
if (
587+
None not in (indexer_unit, index_unit)
588+
and not registry.is_compatible_with(indexer_unit, index_unit)
589+
)
590+
]
591+
if incompatible_units:
592+
units1 = {key: indexer_units[key] for key in incompatible_units}
593+
units2 = {key: index_units[key] for key in incompatible_units}
594+
raise DimensionalityError(units1, units2)
595+
596+
# convert the indexes to the indexer's units
597+
converted = conversion.convert_units(self.da, indexer_units)
598+
stripped = conversion.strip_units(converted)
599+
600+
# index
601+
stripped_indexers = {
602+
name: conversion.strip_indexer_units(indexer)
603+
for name, indexer in indexers.items()
604+
}
605+
interpolated = stripped.interp(
606+
stripped_indexers,
607+
method=method,
608+
assume_sorted=False,
609+
kwargs=None,
610+
)
611+
return conversion.attach_units(interpolated, indexer_units)
612+
613+
def interp_like(self, other, method="linear", assume_sorted=False, kwargs=None):
614+
"""unit-aware version of interp_like
615+
616+
Just like :py:meth:`xarray.DataArray.interp_like`, except the object's indexes are converted
617+
to the units of the indexers first.
618+
619+
.. note::
620+
``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to
621+
``DataArray.interp_like`` unmodified.
622+
623+
See Also
624+
--------
625+
xarray.Dataset.pint.interp_like
626+
xarray.DataArray.pint.interp
627+
xarray.DataArray.interp_like
628+
"""
629+
indexer_units = conversion.extract_unit_attributes(other)
630+
631+
# make sure we only have compatible units
632+
dims = self.da.dims
633+
unit_attrs = conversion.extract_unit_attributes(self.da)
634+
index_units = {
635+
name: units for name, units in unit_attrs.items() if name in dims
636+
}
637+
638+
registry = get_registry(None, index_units, indexer_units)
639+
640+
units = zip_mappings(indexer_units, index_units)
641+
incompatible_units = [
642+
key
643+
for key, (indexer_unit, index_unit) in units.items()
644+
if (
645+
None not in (indexer_unit, index_unit)
646+
and not registry.is_compatible_with(indexer_unit, index_unit)
647+
)
648+
]
649+
if incompatible_units:
650+
units1 = {key: indexer_units[key] for key in incompatible_units}
651+
units2 = {key: index_units[key] for key in incompatible_units}
652+
raise DimensionalityError(units1, units2)
653+
654+
converted = conversion.convert_units(self.da, indexer_units)
655+
stripped = conversion.strip_units(converted)
656+
interpolated = stripped.interp_like(
657+
other,
658+
method=method,
659+
assume_sorted=assume_sorted,
660+
kwargs=kwargs,
661+
)
662+
return conversion.attach_units(interpolated, indexer_units)
663+
543664
def sel(
544665
self, indexers=None, method=None, tolerance=None, drop=False, **indexers_kwargs
545666
):
546667
"""unit-aware version of sel
547668
548-
Just like :py:meth:`xarray.DataArray.sel`, except the dataset's indexes are converted
669+
Just like :py:meth:`xarray.DataArray.sel`, except the object's indexes are converted
549670
to the units of the indexers first.
550671
551672
.. note::
@@ -743,7 +864,7 @@ def to(self, units=None, **unit_kwargs):
743864
----------
744865
units : unit-like or mapping of hashable to unit-like, optional
745866
The units to convert to. If a unit name or ``pint.Unit``
746-
object, convert all the Dataset's data variables. If a dict-like, it
867+
object, convert all the object's data variables. If a dict-like, it
747868
maps variable names to unit names or ``pint.Unit``
748869
objects.
749870
**unit_kwargs
@@ -891,7 +1012,7 @@ def reindex(
8911012
):
8921013
"""unit-aware version of reindex
8931014
894-
Just like :py:meth:`xarray.Dataset.reindex`, except the dataset's indexes are converted
1015+
Just like :py:meth:`xarray.Dataset.reindex`, except the object's indexes are converted
8951016
to the units of the indexers first.
8961017
8971018
.. note::
@@ -959,7 +1080,7 @@ def reindex_like(
9591080
):
9601081
"""unit-aware version of reindex_like
9611082
962-
Just like :py:meth:`xarray.Dataset.reindex_like`, except the dataset's indexes are converted
1083+
Just like :py:meth:`xarray.Dataset.reindex_like`, except the object's indexes are converted
9631084
to the units of the indexers first.
9641085
9651086
.. note::
@@ -1009,12 +1130,133 @@ def reindex_like(
10091130
fill_value=fill_value,
10101131
)
10111132

1133+
def interp(
1134+
self,
1135+
coords=None,
1136+
method="linear",
1137+
assume_sorted=False,
1138+
kwargs=None,
1139+
**coords_kwargs,
1140+
):
1141+
"""unit-aware version of interp
1142+
1143+
Just like :py:meth:`xarray.Dataset.interp`, except the object's indexes are converted
1144+
to the units of the indexers first.
1145+
1146+
.. note::
1147+
``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to
1148+
``Dataset.interp`` unmodified.
1149+
1150+
See Also
1151+
--------
1152+
xarray.DataArray.pint.interp
1153+
xarray.Dataset.pint.interp_like
1154+
xarray.Dataset.interp
1155+
"""
1156+
indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp")
1157+
1158+
indexer_units = {
1159+
name: conversion.extract_indexer_units(indexer)
1160+
for name, indexer in indexers.items()
1161+
}
1162+
1163+
# make sure we only have compatible units
1164+
dims = self.ds.dims
1165+
unit_attrs = conversion.extract_unit_attributes(self.ds)
1166+
index_units = {
1167+
name: units for name, units in unit_attrs.items() if name in dims
1168+
}
1169+
1170+
registry = get_registry(None, index_units, indexer_units)
1171+
1172+
units = zip_mappings(indexer_units, index_units)
1173+
incompatible_units = [
1174+
key
1175+
for key, (indexer_unit, index_unit) in units.items()
1176+
if (
1177+
None not in (indexer_unit, index_unit)
1178+
and not registry.is_compatible_with(indexer_unit, index_unit)
1179+
)
1180+
]
1181+
if incompatible_units:
1182+
units1 = {key: indexer_units[key] for key in incompatible_units}
1183+
units2 = {key: index_units[key] for key in incompatible_units}
1184+
raise DimensionalityError(units1, units2)
1185+
1186+
# convert the indexes to the indexer's units
1187+
converted = conversion.convert_units(self.ds, indexer_units)
1188+
stripped = conversion.strip_units(converted)
1189+
1190+
# index
1191+
stripped_indexers = {
1192+
name: conversion.strip_indexer_units(indexer)
1193+
for name, indexer in indexers.items()
1194+
}
1195+
interpolated = stripped.interp(
1196+
stripped_indexers,
1197+
method=method,
1198+
assume_sorted=False,
1199+
kwargs=None,
1200+
)
1201+
return conversion.attach_units(interpolated, indexer_units)
1202+
1203+
def interp_like(self, other, method="linear", assume_sorted=False, kwargs=None):
1204+
"""unit-aware version of interp_like
1205+
1206+
Just like :py:meth:`xarray.Dataset.interp_like`, except the object's indexes are converted
1207+
to the units of the indexers first.
1208+
1209+
.. note::
1210+
``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to
1211+
``Dataset.interp_like`` unmodified.
1212+
1213+
See Also
1214+
--------
1215+
xarray.DataArray.pint.interp_like
1216+
xarray.Dataset.pint.interp
1217+
xarray.Dataset.interp_like
1218+
"""
1219+
indexer_units = conversion.extract_unit_attributes(other)
1220+
1221+
# make sure we only have compatible units
1222+
dims = self.ds.dims
1223+
unit_attrs = conversion.extract_unit_attributes(self.ds)
1224+
index_units = {
1225+
name: units for name, units in unit_attrs.items() if name in dims
1226+
}
1227+
1228+
registry = get_registry(None, index_units, indexer_units)
1229+
1230+
units = zip_mappings(indexer_units, index_units)
1231+
incompatible_units = [
1232+
key
1233+
for key, (indexer_unit, index_unit) in units.items()
1234+
if (
1235+
None not in (indexer_unit, index_unit)
1236+
and not registry.is_compatible_with(indexer_unit, index_unit)
1237+
)
1238+
]
1239+
if incompatible_units:
1240+
units1 = {key: indexer_units[key] for key in incompatible_units}
1241+
units2 = {key: index_units[key] for key in incompatible_units}
1242+
raise DimensionalityError(units1, units2)
1243+
1244+
converted = conversion.convert_units(self.ds, indexer_units)
1245+
stripped = conversion.strip_units(converted)
1246+
interpolated = stripped.interp_like(
1247+
other,
1248+
method=method,
1249+
assume_sorted=assume_sorted,
1250+
kwargs=kwargs,
1251+
)
1252+
return conversion.attach_units(interpolated, indexer_units)
1253+
10121254
def sel(
10131255
self, indexers=None, method=None, tolerance=None, drop=False, **indexers_kwargs
10141256
):
10151257
"""unit-aware version of sel
10161258
1017-
Just like :py:meth:`xarray.Dataset.sel`, except the dataset's indexes are converted to the units
1259+
Just like :py:meth:`xarray.Dataset.sel`, except the object's indexes are converted to the units
10181260
of the indexers first.
10191261
10201262
.. note::

0 commit comments

Comments
 (0)