Skip to content

Commit 901fa8b

Browse files
keewiskeewis
authored
bfill and ffill (#78)
* 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]
1 parent 6cdab46 commit 901fa8b

File tree

5 files changed

+266
-0
lines changed

5 files changed

+266
-0
lines changed

‎ci/requirements.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
pint
22
numpy
33
scipy
4+
bottleneck
45
xarray
56
isort
67
black

‎docs/api.rst

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ Dataset
2222
xarray.Dataset.pint.drop_sel
2323
xarray.Dataset.pint.sel
2424
xarray.Dataset.pint.to
25+
xarray.Dataset.pint.ffill
26+
xarray.Dataset.pint.bfill
2527

2628
DataArray
2729
---------
@@ -47,6 +49,8 @@ DataArray
4749
xarray.DataArray.pint.drop_sel
4850
xarray.DataArray.pint.sel
4951
xarray.DataArray.pint.to
52+
xarray.DataArray.pint.ffill
53+
xarray.DataArray.pint.bfill
5054

5155
Testing
5256
-------

‎docs/whats-new.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@ What's new
3333
- implement :py:meth:`Dataset.pint.interp`, :py:meth:`Dataset.pint.interp_like`,
3434
:py:meth:`DataArray.pint.interp` and :py:meth:`DataArray.pint.interp_like` (:pull:`72`, :pull:`76`).
3535
By `Justus Magin <https://github.com/keewis>`_.
36+
- implement :py:meth:`Dataset.pint.ffill`, :py:meth:`Dataset.pint.bfill`,
37+
:py:meth:`DataArray.pint.ffill` and :py:meth:`DataArray.pint.bfill` (:pull:`78`).
38+
By `Justus Magin <https://github.com/keewis>`_.
3639

3740
v0.1 (October 26 2020)
3841
----------------------

‎pint_xarray/accessors.py

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -792,6 +792,40 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs):
792792

793793
return indexed
794794

795+
def ffill(self, dim, limit=None):
796+
"""unit-aware version of ffill
797+
798+
Like :py:meth:`DataArray.ffill` but without stripping the data units.
799+
800+
See Also
801+
--------
802+
xarray.DataArray.ffill
803+
xarray.DataArray.pint.bfill
804+
"""
805+
units = conversion.extract_units(self.da)
806+
stripped = conversion.strip_units(self.da)
807+
808+
filled = stripped.ffill(dim=dim, limit=limit)
809+
810+
return conversion.attach_units(filled, units)
811+
812+
def bfill(self, dim, limit=None):
813+
"""unit-aware version of bfill
814+
815+
Like :py:meth:`DataArray.bfill` but without stripping the data units.
816+
817+
See Also
818+
--------
819+
xarray.DataArray.bfill
820+
xarray.DataArray.pint.ffill
821+
"""
822+
units = conversion.extract_units(self.da)
823+
stripped = conversion.strip_units(self.da)
824+
825+
filled = stripped.bfill(dim=dim, limit=limit)
826+
827+
return conversion.attach_units(filled, units)
828+
795829

796830
@register_dataset_accessor("pint")
797831
class PintDatasetAccessor:
@@ -1441,3 +1475,37 @@ def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs):
14411475
)
14421476

14431477
return indexed
1478+
1479+
def ffill(self, dim, limit=None):
1480+
"""unit-aware version of ffill
1481+
1482+
Like :py:meth:`Dataset.ffill` but without stripping the data units.
1483+
1484+
See Also
1485+
--------
1486+
xarray.Dataset.ffill
1487+
xarray.Dataset.pint.bfill
1488+
"""
1489+
units = conversion.extract_units(self.ds)
1490+
stripped = conversion.strip_units(self.ds)
1491+
1492+
filled = stripped.ffill(dim=dim, limit=limit)
1493+
1494+
return conversion.attach_units(filled, units)
1495+
1496+
def bfill(self, dim, limit=None):
1497+
"""unit-aware version of bfill
1498+
1499+
Like :py:meth:`Dataset.bfill` but without stripping the data units.
1500+
1501+
See Also
1502+
--------
1503+
xarray.Dataset.bfill
1504+
xarray.Dataset.pint.ffill
1505+
"""
1506+
units = conversion.extract_units(self.ds)
1507+
stripped = conversion.strip_units(self.ds)
1508+
1509+
filled = stripped.bfill(dim=dim, limit=limit)
1510+
1511+
return conversion.attach_units(filled, units)

‎pint_xarray/tests/test_accessors.py

Lines changed: 190 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
unit_registry = UnitRegistry(force_ndarray=True)
1919
Quantity = unit_registry.Quantity
2020

21+
nan = np.nan
22+
2123

2224
def assert_all_str_or_none(mapping):
2325
__tracebackhide__ = True
@@ -1183,3 +1185,191 @@ def test_interp_like(obj, other, expected, error):
11831185
actual = obj.pint.interp_like(other)
11841186
assert_units_equal(actual, expected)
11851187
assert_identical(actual, expected)
1188+
1189+
1190+
@pytest.mark.parametrize(
1191+
["obj", "expected"],
1192+
(
1193+
pytest.param(
1194+
xr.Dataset(
1195+
{"a": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
1196+
coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
1197+
),
1198+
xr.Dataset(
1199+
{"a": ("x", [nan, 0, 0, 1, 1, 1, 2, 2])},
1200+
coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
1201+
),
1202+
id="Dataset-no units",
1203+
),
1204+
pytest.param(
1205+
xr.Dataset(
1206+
{
1207+
"a": (
1208+
"x",
1209+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1210+
)
1211+
},
1212+
coords={
1213+
"u": (
1214+
"x",
1215+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1216+
)
1217+
},
1218+
),
1219+
xr.Dataset(
1220+
{"a": ("x", Quantity([nan, 0, 0, 1, 1, 1, 2, 2], "m"))},
1221+
coords={
1222+
"u": (
1223+
"x",
1224+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1225+
)
1226+
},
1227+
),
1228+
id="Dataset-units",
1229+
),
1230+
pytest.param(
1231+
xr.DataArray(
1232+
[nan, 0, nan, 1, nan, nan, 2, nan],
1233+
coords={
1234+
"u": (
1235+
"x",
1236+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1237+
)
1238+
},
1239+
dims="x",
1240+
),
1241+
xr.DataArray(
1242+
[nan, 0, 0, 1, 1, 1, 2, 2],
1243+
coords={
1244+
"u": (
1245+
"x",
1246+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1247+
)
1248+
},
1249+
dims="x",
1250+
),
1251+
id="DataArray-units",
1252+
),
1253+
pytest.param(
1254+
xr.DataArray(
1255+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1256+
coords={
1257+
"u": (
1258+
"x",
1259+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1260+
)
1261+
},
1262+
dims="x",
1263+
),
1264+
xr.DataArray(
1265+
Quantity([nan, 0, 0, 1, 1, 1, 2, 2], "m"),
1266+
coords={
1267+
"u": (
1268+
"x",
1269+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1270+
)
1271+
},
1272+
dims="x",
1273+
),
1274+
id="DataArray-units",
1275+
),
1276+
),
1277+
)
1278+
def test_ffill(obj, expected):
1279+
actual = obj.pint.ffill(dim="x")
1280+
assert_identical(actual, expected)
1281+
assert_units_equal(actual, expected)
1282+
1283+
1284+
@pytest.mark.parametrize(
1285+
["obj", "expected"],
1286+
(
1287+
pytest.param(
1288+
xr.Dataset(
1289+
{"a": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
1290+
coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
1291+
),
1292+
xr.Dataset(
1293+
{"a": ("x", [0, 0, 1, 1, 2, 2, 2, nan])},
1294+
coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])},
1295+
),
1296+
id="Dataset-no units",
1297+
),
1298+
pytest.param(
1299+
xr.Dataset(
1300+
{
1301+
"a": (
1302+
"x",
1303+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1304+
)
1305+
},
1306+
coords={
1307+
"u": (
1308+
"x",
1309+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1310+
)
1311+
},
1312+
),
1313+
xr.Dataset(
1314+
{"a": ("x", Quantity([0, 0, 1, 1, 2, 2, 2, nan], "m"))},
1315+
coords={
1316+
"u": (
1317+
"x",
1318+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1319+
)
1320+
},
1321+
),
1322+
id="Dataset-units",
1323+
),
1324+
pytest.param(
1325+
xr.DataArray(
1326+
[nan, 0, nan, 1, nan, nan, 2, nan],
1327+
coords={
1328+
"u": (
1329+
"x",
1330+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1331+
)
1332+
},
1333+
dims="x",
1334+
),
1335+
xr.DataArray(
1336+
[0, 0, 1, 1, 2, 2, 2, nan],
1337+
coords={
1338+
"u": (
1339+
"x",
1340+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1341+
)
1342+
},
1343+
dims="x",
1344+
),
1345+
id="DataArray-units",
1346+
),
1347+
pytest.param(
1348+
xr.DataArray(
1349+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1350+
coords={
1351+
"u": (
1352+
"x",
1353+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1354+
)
1355+
},
1356+
dims="x",
1357+
),
1358+
xr.DataArray(
1359+
Quantity([0, 0, 1, 1, 2, 2, 2, nan], "m"),
1360+
coords={
1361+
"u": (
1362+
"x",
1363+
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1364+
)
1365+
},
1366+
dims="x",
1367+
),
1368+
id="DataArray-units",
1369+
),
1370+
),
1371+
)
1372+
def test_bfill(obj, expected):
1373+
actual = obj.pint.bfill(dim="x")
1374+
assert_identical(actual, expected)
1375+
assert_units_equal(actual, expected)

0 commit comments

Comments
 (0)