Skip to content

Commit

Permalink
bfill and ffill (#78)
Browse files Browse the repository at this point in the history 
* add tests for ffill and bfill

* implement ffill and bfill

* add ffill and bfill to api.rst

* update whats-new.rst [skip-ci]

* add bottleneck to the environment

* fix whats-new.rst [skip-ci]

* fix the whats-new.rst entry [skip-ci]

* fix [skip-ci]
  • Loading branch information
@keewis
keewis authored Mar 6, 2021
1 parent 6cdab46 commit 901fa8b
Show file tree
Hide file tree
Showing 5 changed files with 266 additions and 0 deletions.
1 change: 1 addition & 0 deletions ci/requirements.txt
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
pint
numpy
scipy
bottleneck
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 @@ -22,6 +22,8 @@ Dataset
xarray.Dataset.pint.drop_sel
xarray.Dataset.pint.sel
xarray.Dataset.pint.to
xarray.Dataset.pint.ffill
xarray.Dataset.pint.bfill

DataArray
---------
Expand All @@ -47,6 +49,8 @@ DataArray
xarray.DataArray.pint.drop_sel
xarray.DataArray.pint.sel
xarray.DataArray.pint.to
xarray.DataArray.pint.ffill
xarray.DataArray.pint.bfill

Testing
-------
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 @@ -33,6 +33,9 @@ What's new
- 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`, :pull:`76`).
By `Justus Magin <https://github.com/keewis>`_.
- implement :py:meth:`Dataset.pint.ffill`, :py:meth:`Dataset.pint.bfill`,
:py:meth:`DataArray.pint.ffill` and :py:meth:`DataArray.pint.bfill` (:pull:`78`).
By `Justus Magin <https://github.com/keewis>`_.

v0.1 (October 26 2020)
----------------------
Expand Down
68 changes: 68 additions & 0 deletions pint_xarray/accessors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -792,6 +792,40 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs):

return indexed

def ffill(self, dim, limit=None):
"""unit-aware version of ffill
Like :py:meth:`DataArray.ffill` but without stripping the data units.
See Also
--------
xarray.DataArray.ffill
xarray.DataArray.pint.bfill
"""
units = conversion.extract_units(self.da)
stripped = conversion.strip_units(self.da)

filled = stripped.ffill(dim=dim, limit=limit)

return conversion.attach_units(filled, units)

def bfill(self, dim, limit=None):
"""unit-aware version of bfill
Like :py:meth:`DataArray.bfill` but without stripping the data units.
See Also
--------
xarray.DataArray.bfill
xarray.DataArray.pint.ffill
"""
units = conversion.extract_units(self.da)
stripped = conversion.strip_units(self.da)

filled = stripped.bfill(dim=dim, limit=limit)

return conversion.attach_units(filled, units)


@register_dataset_accessor("pint")
class PintDatasetAccessor:
Expand Down Expand Up @@ -1441,3 +1475,37 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs):
)

return indexed

def ffill(self, dim, limit=None):
"""unit-aware version of ffill
Like :py:meth:`Dataset.ffill` but without stripping the data units.
See Also
--------
xarray.Dataset.ffill
xarray.Dataset.pint.bfill
"""
units = conversion.extract_units(self.ds)
stripped = conversion.strip_units(self.ds)

filled = stripped.ffill(dim=dim, limit=limit)

return conversion.attach_units(filled, units)

def bfill(self, dim, limit=None):
"""unit-aware version of bfill
Like :py:meth:`Dataset.bfill` but without stripping the data units.
See Also
--------
xarray.Dataset.bfill
xarray.Dataset.pint.ffill
"""
units = conversion.extract_units(self.ds)
stripped = conversion.strip_units(self.ds)

filled = stripped.bfill(dim=dim, limit=limit)

return conversion.attach_units(filled, units)
190 changes: 190 additions & 0 deletions pint_xarray/tests/test_accessors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
unit_registry = UnitRegistry(force_ndarray=True)
Quantity = unit_registry.Quantity

nan = np.nan


def assert_all_str_or_none(mapping):
__tracebackhide__ = True
Expand Down Expand Up @@ -1183,3 +1185,191 @@ def test_interp_like(obj, other, expected, error):
actual = obj.pint.interp_like(other)
assert_units_equal(actual, expected)
assert_identical(actual, expected)


@pytest.mark.parametrize(
["obj", "expected"],
(
pytest.param(
xr.Dataset(
{"a": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
),
xr.Dataset(
{"a": ("x", [nan, 0, 0, 1, 1, 1, 2, 2])},
coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
),
id="Dataset-no units",
),
pytest.param(
xr.Dataset(
{
"a": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
),
xr.Dataset(
{"a": ("x", Quantity([nan, 0, 0, 1, 1, 1, 2, 2], "m"))},
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
),
id="Dataset-units",
),
pytest.param(
xr.DataArray(
[nan, 0, nan, 1, nan, nan, 2, nan],
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
dims="x",
),
xr.DataArray(
[nan, 0, 0, 1, 1, 1, 2, 2],
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
dims="x",
),
id="DataArray-units",
),
pytest.param(
xr.DataArray(
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
dims="x",
),
xr.DataArray(
Quantity([nan, 0, 0, 1, 1, 1, 2, 2], "m"),
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
dims="x",
),
id="DataArray-units",
),
),
)
def test_ffill(obj, expected):
actual = obj.pint.ffill(dim="x")
assert_identical(actual, expected)
assert_units_equal(actual, expected)


@pytest.mark.parametrize(
["obj", "expected"],
(
pytest.param(
xr.Dataset(
{"a": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
),
xr.Dataset(
{"a": ("x", [0, 0, 1, 1, 2, 2, 2, nan])},
coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
),
id="Dataset-no units",
),
pytest.param(
xr.Dataset(
{
"a": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
),
xr.Dataset(
{"a": ("x", Quantity([0, 0, 1, 1, 2, 2, 2, nan], "m"))},
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
),
id="Dataset-units",
),
pytest.param(
xr.DataArray(
[nan, 0, nan, 1, nan, nan, 2, nan],
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
dims="x",
),
xr.DataArray(
[0, 0, 1, 1, 2, 2, 2, nan],
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
dims="x",
),
id="DataArray-units",
),
pytest.param(
xr.DataArray(
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
dims="x",
),
xr.DataArray(
Quantity([0, 0, 1, 1, 2, 2, 2, nan], "m"),
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
)
},
dims="x",
),
id="DataArray-units",
),
),
)
def test_bfill(obj, expected):
actual = obj.pint.bfill(dim="x")
assert_identical(actual, expected)
assert_units_equal(actual, expected)

0 comments on commit 901fa8b

Please sign in to comment.