Skip to content

Commit 0f891ea

Browse files
keewiskeewis
authored
chunk (#83)
* fix the tests of bfill and ffill * add tests for chunk * implement chunk * fix the chunk tests * add chunk to api.rst * add dask to the environment * update whats-new.rst * install dask[array] * fix the ?fill tests
1 parent 87157a3 commit 0f891ea

File tree

5 files changed

+123
-6
lines changed

5 files changed

+123
-6
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+
dask[array]
45
bottleneck
56
xarray
67
isort

‎docs/api.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ Dataset
2828
xarray.Dataset.pint.drop_sel
2929
xarray.Dataset.pint.sel
3030
xarray.Dataset.pint.to
31+
xarray.Dataset.pint.chunk
3132
xarray.Dataset.pint.ffill
3233
xarray.Dataset.pint.bfill
3334
xarray.Dataset.pint.interpolate_na
@@ -58,6 +59,7 @@ DataArray
5859
xarray.DataArray.pint.drop_sel
5960
xarray.DataArray.pint.sel
6061
xarray.DataArray.pint.to
62+
xarray.DataArray.pint.chunk
6163
xarray.DataArray.pint.ffill
6264
xarray.DataArray.pint.bfill
6365
xarray.DataArray.pint.interpolate_na

‎docs/whats-new.rst

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@ What's new
2929
By `Justus Magin <https://github.com/keewis>`_.
3030
- implement :py:meth:`Dataset.pint.drop_sel` and :py:meth:`DataArray.pint.drop_sel` (:pull:`73`).
3131
By `Justus Magin <https://github.com/keewis>`_.
32+
- implement :py:meth:`Dataset.pint.chunk` and :py:meth:`DataArray.pint.chunk` (:pull:`83`).
33+
By `Justus Magin <https://github.com/keewis>`_.
3234
- implement :py:meth:`Dataset.pint.reindex`, :py:meth:`Dataset.pint.reindex_like`,
3335
:py:meth:`DataArray.pint.reindex` and :py:meth:`DataArray.pint.reindex_like` (:pull:`69`).
3436
By `Justus Magin <https://github.com/keewis>`_.

‎pint_xarray/accessors.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -569,6 +569,29 @@ def to(self, units=None, **unit_kwargs):
569569

570570
return conversion.convert_units(self.da, units)
571571

572+
def chunk(self, chunks, name_prefix="xarray-", token=None, lock=False):
573+
"""unit-aware version of chunk
574+
575+
Like :py:meth:`xarray.DataArray.chunk`, but chunking a quantity will change the
576+
wrapped type to ``dask``.
577+
578+
.. note::
579+
It is recommended to only use this when chunking in-memory arrays. To
580+
rechunk please use :py:meth:`xarray.DataArray.chunk`.
581+
582+
See Also
583+
--------
584+
xarray.DataArray.chunk
585+
xarray.Dataset.pint.chunk
586+
"""
587+
units = conversion.extract_units(self.da)
588+
stripped = conversion.strip_units(self.da)
589+
590+
chunked = stripped.chunk(
591+
chunks, name_prefix=name_prefix, token=token, lock=lock
592+
)
593+
return conversion.attach_units(chunked, units)
594+
572595
def reindex(
573596
self,
574597
indexers=None,
@@ -1305,6 +1328,29 @@ def to(self, units=None, **unit_kwargs):
13051328

13061329
return conversion.convert_units(self.ds, units)
13071330

1331+
def chunk(self, chunks, name_prefix="xarray-", token=None, lock=False):
1332+
"""unit-aware version of chunk
1333+
1334+
Like :py:meth:`xarray.Dataset.chunk`, but chunking a quantity will change the
1335+
wrapped type to ``dask``.
1336+
1337+
.. note::
1338+
It is recommended to only use this when chunking in-memory arrays. To
1339+
rechunk please use :py:meth:`xarray.Dataset.chunk`.
1340+
1341+
See Also
1342+
--------
1343+
xarray.Dataset.chunk
1344+
xarray.DataArray.pint.chunk
1345+
"""
1346+
units = conversion.extract_units(self.ds)
1347+
stripped = conversion.strip_units(self.ds)
1348+
1349+
chunked = stripped.chunk(
1350+
chunks, name_prefix=name_prefix, token=token, lock=lock
1351+
)
1352+
return conversion.attach_units(chunked, units)
1353+
13081354
def reindex(
13091355
self,
13101356
indexers=None,

‎pint_xarray/tests/test_accessors.py

Lines changed: 72 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -814,6 +814,72 @@ def test_drop_sel(obj, indexers, expected, error):
814814
assert_identical(actual, expected)
815815

816816

817+
@pytest.mark.parametrize(
818+
"obj",
819+
(
820+
pytest.param(
821+
xr.Dataset(
822+
{"a": ("x", np.linspace(0, 1, 11))},
823+
coords={"u": ("x", np.arange(11))},
824+
),
825+
id="Dataset-no units",
826+
),
827+
pytest.param(
828+
xr.Dataset(
829+
{
830+
"a": (
831+
"x",
832+
Quantity(np.linspace(0, 1, 11), "m"),
833+
)
834+
},
835+
coords={
836+
"u": (
837+
"x",
838+
Quantity(np.arange(11), "m"),
839+
)
840+
},
841+
),
842+
id="Dataset-units",
843+
),
844+
pytest.param(
845+
xr.DataArray(
846+
np.linspace(0, 1, 11),
847+
coords={
848+
"u": (
849+
"x",
850+
np.arange(11),
851+
)
852+
},
853+
dims="x",
854+
),
855+
id="DataArray-no units",
856+
),
857+
pytest.param(
858+
xr.DataArray(
859+
Quantity(np.linspace(0, 1, 11), "m"),
860+
coords={
861+
"u": (
862+
"x",
863+
Quantity(np.arange(11), "m"),
864+
)
865+
},
866+
dims="x",
867+
),
868+
id="DataArray-units",
869+
),
870+
),
871+
)
872+
def test_chunk(obj):
873+
actual = obj.pint.chunk({"x": 2})
874+
875+
expected = (
876+
obj.pint.dequantify().chunk({"x": 2}).pint.quantify(unit_registry=unit_registry)
877+
)
878+
879+
assert_units_equal(actual, expected)
880+
assert_identical(actual, expected)
881+
882+
817883
@pytest.mark.parametrize(
818884
["obj", "indexers", "expected", "error"],
819885
(
@@ -1484,7 +1550,7 @@ def test_interp_like(obj, other, expected, error):
14841550
coords={
14851551
"u": (
14861552
"x",
1487-
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1553+
[nan, 0, nan, 1, nan, nan, 2, nan],
14881554
)
14891555
},
14901556
dims="x",
@@ -1494,12 +1560,12 @@ def test_interp_like(obj, other, expected, error):
14941560
coords={
14951561
"u": (
14961562
"x",
1497-
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1563+
[nan, 0, nan, 1, nan, nan, 2, nan],
14981564
)
14991565
},
15001566
dims="x",
15011567
),
1502-
id="DataArray-units",
1568+
id="DataArray-no units",
15031569
),
15041570
pytest.param(
15051571
xr.DataArray(
@@ -1578,7 +1644,7 @@ def test_ffill(obj, expected):
15781644
coords={
15791645
"u": (
15801646
"x",
1581-
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1647+
[nan, 0, nan, 1, nan, nan, 2, nan],
15821648
)
15831649
},
15841650
dims="x",
@@ -1588,12 +1654,12 @@ def test_ffill(obj, expected):
15881654
coords={
15891655
"u": (
15901656
"x",
1591-
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
1657+
[nan, 0, nan, 1, nan, nan, 2, nan],
15921658
)
15931659
},
15941660
dims="x",
15951661
),
1596-
id="DataArray-units",
1662+
id="DataArray-no units",
15971663
),
15981664
pytest.param(
15991665
xr.DataArray(

0 commit comments

Comments
 (0)