Skip to content

Commit 87157a3

Browse files
keewiskeewis
authored
interpolate_na (#82)
* add tests for interpolate_na * implement interpolate_na * add interpolate_na to the api documentation * update whats_new.rst
1 parent 038ac60 commit 87157a3

File tree

4 files changed

+174
-0
lines changed

4 files changed

+174
-0
lines changed

‎docs/api.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ Dataset
3030
xarray.Dataset.pint.to
3131
xarray.Dataset.pint.ffill
3232
xarray.Dataset.pint.bfill
33+
xarray.Dataset.pint.interpolate_na
3334

3435
DataArray
3536
---------
@@ -59,6 +60,7 @@ DataArray
5960
xarray.DataArray.pint.to
6061
xarray.DataArray.pint.ffill
6162
xarray.DataArray.pint.bfill
63+
xarray.DataArray.pint.interpolate_na
6264

6365
Testing
6466
-------

‎docs/whats-new.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ What's new
3838
- implement :py:meth:`Dataset.pint.ffill`, :py:meth:`Dataset.pint.bfill`,
3939
:py:meth:`DataArray.pint.ffill` and :py:meth:`DataArray.pint.bfill` (:pull:`78`).
4040
By `Justus Magin <https://github.com/keewis>`_.
41+
- implement :py:meth:`Dataset.pint.interpolate_na` and :py:meth:`DataArray.pint.interpolate_na` (:pull:`82`).
42+
By `Justus Magin <https://github.com/keewis>`_.
4143

4244
v0.1 (October 26 2020)
4345
----------------------

‎pint_xarray/accessors.py

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -998,6 +998,44 @@ def bfill(self, dim, limit=None):
998998

999999
return conversion.attach_units(filled, units)
10001000

1001+
def interpolate_na(
1002+
self,
1003+
dim=None,
1004+
method="linear",
1005+
limit=None,
1006+
use_coordinate=True,
1007+
max_gap=None,
1008+
keep_attrs=None,
1009+
**kwargs,
1010+
):
1011+
"""unit-aware version of interpolate_na
1012+
1013+
Like :py:meth:`DataArray.interpolate_na` but without stripping the units on data or coordinates.
1014+
1015+
.. note::
1016+
``max_gap`` is not supported, yet, and will be passed through to
1017+
``DataArray.interpolate_na`` unmodified.
1018+
1019+
See Also
1020+
--------
1021+
xarray.Dataset.pint.interpolate_na
1022+
xarray.DataArray.interpolate_na
1023+
"""
1024+
units = conversion.extract_units(self.da)
1025+
stripped = conversion.strip_units(self.da)
1026+
1027+
interpolated = stripped.interpolate_na(
1028+
dim=dim,
1029+
method=method,
1030+
limit=limit,
1031+
use_coordinate=use_coordinate,
1032+
max_gap=max_gap,
1033+
keep_attrs=keep_attrs,
1034+
**kwargs,
1035+
)
1036+
1037+
return conversion.attach_units(interpolated, units)
1038+
10011039

10021040
@register_dataset_accessor("pint")
10031041
class PintDatasetAccessor:
@@ -1697,3 +1735,41 @@ def bfill(self, dim, limit=None):
16971735
filled = stripped.bfill(dim=dim, limit=limit)
16981736

16991737
return conversion.attach_units(filled, units)
1738+
1739+
def interpolate_na(
1740+
self,
1741+
dim=None,
1742+
method="linear",
1743+
limit=None,
1744+
use_coordinate=True,
1745+
max_gap=None,
1746+
keep_attrs=None,
1747+
**kwargs,
1748+
):
1749+
"""unit-aware version of interpolate_na
1750+
1751+
Like :py:meth:`Dataset.interpolate_na` but without stripping the units on data or coordinates.
1752+
1753+
.. note::
1754+
``max_gap`` is not supported, yet, and will be passed through to
1755+
``Dataset.interpolate_na`` unmodified.
1756+
1757+
See Also
1758+
--------
1759+
xarray.DataArray.pint.interpolate_na
1760+
xarray.Dataset.interpolate_na
1761+
"""
1762+
units = conversion.extract_units(self.ds)
1763+
stripped = conversion.strip_units(self.ds)
1764+
1765+
interpolated = stripped.interpolate_na(
1766+
dim=dim,
1767+
method=method,
1768+
limit=limit,
1769+
use_coordinate=use_coordinate,
1770+
max_gap=max_gap,
1771+
keep_attrs=keep_attrs,
1772+
**kwargs,
1773+
)
1774+
1775+
return conversion.attach_units(interpolated, units)

‎pint_xarray/tests/test_accessors.py

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1624,3 +1624,97 @@ def test_bfill(obj, expected):
16241624
actual = obj.pint.bfill(dim="x")
16251625
assert_identical(actual, expected)
16261626
assert_units_equal(actual, expected)
1627+
1628+
1629+
@pytest.mark.parametrize(
1630+
["obj", "expected"],
1631+
(
1632+
pytest.param(
1633+
xr.Dataset(
1634+
{"a": ("x", [nan, 0, nan, 1, nan, nan, nan, 2, nan])},
1635+
coords={"u": ("x", [nan, 0, nan, 1, nan, nan, nan, 2, nan])},
1636+
),
1637+
xr.Dataset(
1638+
{"a": ("x", [nan, 0, 0.5, 1, 1.25, 1.5, 1.75, 2, nan])},
1639+
coords={"u": ("x", [nan, 0, nan, 1, nan, nan, nan, 2, nan])},
1640+
),
1641+
id="Dataset-no units",
1642+
),
1643+
pytest.param(
1644+
xr.Dataset(
1645+
{
1646+
"a": (
1647+
"x",
1648+
Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"),
1649+
)
1650+
},
1651+
coords={
1652+
"u": (
1653+
"x",
1654+
Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"),
1655+
)
1656+
},
1657+
),
1658+
xr.Dataset(
1659+
{"a": ("x", Quantity([nan, 0, 0.5, 1, 1.25, 1.5, 1.75, 2, nan], "m"))},
1660+
coords={
1661+
"u": (
1662+
"x",
1663+
Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"),
1664+
)
1665+
},
1666+
),
1667+
id="Dataset-units",
1668+
),
1669+
pytest.param(
1670+
xr.DataArray(
1671+
[nan, 0, nan, 1, nan, nan, nan, 2, nan],
1672+
coords={
1673+
"u": (
1674+
"x",
1675+
Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"),
1676+
)
1677+
},
1678+
dims="x",
1679+
),
1680+
xr.DataArray(
1681+
[nan, 0, 0.5, 1, 1.25, 1.5, 1.75, 2, nan],
1682+
coords={
1683+
"u": (
1684+
"x",
1685+
Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"),
1686+
)
1687+
},
1688+
dims="x",
1689+
),
1690+
id="DataArray-units",
1691+
),
1692+
pytest.param(
1693+
xr.DataArray(
1694+
Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"),
1695+
coords={
1696+
"u": (
1697+
"x",
1698+
Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"),
1699+
)
1700+
},
1701+
dims="x",
1702+
),
1703+
xr.DataArray(
1704+
Quantity([nan, 0, 0.5, 1, 1.25, 1.5, 1.75, 2, nan], "m"),
1705+
coords={
1706+
"u": (
1707+
"x",
1708+
Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"),
1709+
)
1710+
},
1711+
dims="x",
1712+
),
1713+
id="DataArray-units",
1714+
),
1715+
),
1716+
)
1717+
def test_interpolate_na(obj, expected):
1718+
actual = obj.pint.interpolate_na(dim="x")
1719+
assert_identical(actual, expected)
1720+
assert_units_equal(actual, expected)

0 commit comments

Comments
 (0)