Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

chunk #83

Merged
merged 9 commits into from
Mar 15, 2021
Merged

chunk #83

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
1e19f2c
fix the tests of bfill and ffill
keewis Mar 15, 2021
50c938e
add tests for chunk
keewis Mar 15, 2021
1354540
implement chunk
keewis Mar 15, 2021
0760cb2
fix the chunk tests
keewis Mar 15, 2021
bdfec09
add chunk to api.rst
keewis Mar 15, 2021
d398686
add dask to the environment
keewis Mar 15, 2021
5249672
update whats-new.rst
keewis Mar 15, 2021
230ed8c
install dask[array]
keewis Mar 15, 2021
3d3b74f
fix the ?fill tests
keewis Mar 15, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions ci/requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
pint
numpy
scipy
dask[array]
bottleneck
xarray
isort
Expand Down
2 changes: 2 additions & 0 deletions docs/api.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ Dataset
xarray.Dataset.pint.drop_sel
xarray.Dataset.pint.sel
xarray.Dataset.pint.to
xarray.Dataset.pint.chunk
xarray.Dataset.pint.ffill
xarray.Dataset.pint.bfill

Expand Down Expand Up @@ -57,6 +58,7 @@ DataArray
xarray.DataArray.pint.drop_sel
xarray.DataArray.pint.sel
xarray.DataArray.pint.to
xarray.DataArray.pint.chunk
xarray.DataArray.pint.ffill
xarray.DataArray.pint.bfill

Expand Down
2 changes: 2 additions & 0 deletions docs/whats-new.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,8 @@ What's new
By `Justus Magin <https://github.com/keewis>`_.
- implement :py:meth:`Dataset.pint.drop_sel` and :py:meth:`DataArray.pint.drop_sel` (:pull:`73`).
By `Justus Magin <https://github.com/keewis>`_.
- implement :py:meth:`Dataset.pint.chunk` and :py:meth:`DataArray.pint.chunk` (:pull:`83`).
By `Justus Magin <https://github.com/keewis>`_.
- implement :py:meth:`Dataset.pint.reindex`, :py:meth:`Dataset.pint.reindex_like`,
:py:meth:`DataArray.pint.reindex` and :py:meth:`DataArray.pint.reindex_like` (:pull:`69`).
By `Justus Magin <https://github.com/keewis>`_.
Expand Down
46 changes: 46 additions & 0 deletions pint_xarray/accessors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -569,6 +569,29 @@ def to(self, units=None, **unit_kwargs):

return conversion.convert_units(self.da, units)

def chunk(self, chunks, name_prefix="xarray-", token=None, lock=False):
"""unit-aware version of chunk

Like :py:meth:`xarray.DataArray.chunk`, but chunking a quantity will change the
wrapped type to ``dask``.

.. note::
It is recommended to only use this when chunking in-memory arrays. To
rechunk please use :py:meth:`xarray.DataArray.chunk`.

See Also
--------
xarray.DataArray.chunk
xarray.Dataset.pint.chunk
"""
units = conversion.extract_units(self.da)
stripped = conversion.strip_units(self.da)

chunked = stripped.chunk(
chunks, name_prefix=name_prefix, token=token, lock=lock
)
return conversion.attach_units(chunked, units)

def reindex(
self,
indexers=None,
Expand Down Expand Up @@ -1267,6 +1290,29 @@ def to(self, units=None, **unit_kwargs):

return conversion.convert_units(self.ds, units)

def chunk(self, chunks, name_prefix="xarray-", token=None, lock=False):
"""unit-aware version of chunk

Like :py:meth:`xarray.Dataset.chunk`, but chunking a quantity will change the
wrapped type to ``dask``.

.. note::
It is recommended to only use this when chunking in-memory arrays. To
rechunk please use :py:meth:`xarray.Dataset.chunk`.

See Also
--------
xarray.Dataset.chunk
xarray.DataArray.pint.chunk
"""
units = conversion.extract_units(self.ds)
stripped = conversion.strip_units(self.ds)

chunked = stripped.chunk(
chunks, name_prefix=name_prefix, token=token, lock=lock
)
return conversion.attach_units(chunked, units)

def reindex(
self,
indexers=None,
Expand Down
78 changes: 72 additions & 6 deletions pint_xarray/tests/test_accessors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -814,6 +814,72 @@ def test_drop_sel(obj, indexers, expected, error):
assert_identical(actual, expected)


@pytest.mark.parametrize(
"obj",
(
pytest.param(
xr.Dataset(
{"a": ("x", np.linspace(0, 1, 11))},
coords={"u": ("x", np.arange(11))},
),
id="Dataset-no units",
),
pytest.param(
xr.Dataset(
{
"a": (
"x",
Quantity(np.linspace(0, 1, 11), "m"),
)
},
coords={
"u": (
"x",
Quantity(np.arange(11), "m"),
)
},
),
id="Dataset-units",
),
pytest.param(
xr.DataArray(
np.linspace(0, 1, 11),
coords={
"u": (
"x",
np.arange(11),
)
},
dims="x",
),
id="DataArray-no units",
),
pytest.param(
xr.DataArray(
Quantity(np.linspace(0, 1, 11), "m"),
coords={
"u": (
"x",
Quantity(np.arange(11), "m"),
)
},
dims="x",
),
id="DataArray-units",
),
),
)
def test_chunk(obj):
actual = obj.pint.chunk({"x": 2})

expected = (
obj.pint.dequantify().chunk({"x": 2}).pint.quantify(unit_registry=unit_registry)
)

assert_units_equal(actual, expected)
assert_identical(actual, expected)


@pytest.mark.parametrize(
["obj", "indexers", "expected", "error"],
(
Expand Down Expand Up @@ -1484,7 +1550,7 @@ def test_interp_like(obj, other, expected, error):
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
[nan, 0, nan, 1, nan, nan, 2, nan],
)
},
dims="x",
Expand All @@ -1494,12 +1560,12 @@ def test_interp_like(obj, other, expected, error):
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
[nan, 0, nan, 1, nan, nan, 2, nan],
)
},
dims="x",
),
id="DataArray-units",
id="DataArray-no units",
),
pytest.param(
xr.DataArray(
Expand Down Expand Up @@ -1578,7 +1644,7 @@ def test_ffill(obj, expected):
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
[nan, 0, nan, 1, nan, nan, 2, nan],
)
},
dims="x",
Expand All @@ -1588,12 +1654,12 @@ def test_ffill(obj, expected):
coords={
"u": (
"x",
Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"),
[nan, 0, nan, 1, nan, nan, 2, nan],
)
},
dims="x",
),
id="DataArray-units",
id="DataArray-no units",
),
pytest.param(
xr.DataArray(
Expand Down