-
-
Notifications
You must be signed in to change notification settings - Fork 1.1k
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.
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
Use duck array ops in more places #8267
Changes from 5 commits
2fba050
46c5ace
a9509f2
c83f3d8
92b5430
6f3bf0d
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
|
@@ -7,6 +7,7 @@ | |||||
import pandas as pd | ||||||
|
||||||
from xarray.coding.times import infer_calendar_name | ||||||
from xarray.core import duck_array_ops | ||||||
from xarray.core.common import ( | ||||||
_contains_datetime_like_objects, | ||||||
is_np_datetime_like, | ||||||
|
@@ -50,7 +51,7 @@ def _access_through_cftimeindex(values, name): | |||||
from xarray.coding.cftimeindex import CFTimeIndex | ||||||
|
||||||
if not isinstance(values, CFTimeIndex): | ||||||
values_as_cftimeindex = CFTimeIndex(values.ravel()) | ||||||
values_as_cftimeindex = CFTimeIndex(duck_array_ops.ravel(values)) | ||||||
else: | ||||||
values_as_cftimeindex = values | ||||||
if name == "season": | ||||||
|
@@ -69,7 +70,7 @@ def _access_through_series(values, name): | |||||
"""Coerce an array of datetime-like values to a pandas Series and | ||||||
access requested datetime component | ||||||
""" | ||||||
values_as_series = pd.Series(values.ravel(), copy=False) | ||||||
values_as_series = pd.Series(duck_array_ops.ravel(values), copy=False) | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
if name == "season": | ||||||
months = values_as_series.dt.month.values | ||||||
field_values = _season_from_months(months) | ||||||
|
@@ -148,10 +149,10 @@ def _round_through_series_or_index(values, name, freq): | |||||
from xarray.coding.cftimeindex import CFTimeIndex | ||||||
|
||||||
if is_np_datetime_like(values.dtype): | ||||||
values_as_series = pd.Series(values.ravel(), copy=False) | ||||||
values_as_series = pd.Series(duck_array_ops.ravel(values), copy=False) | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
method = getattr(values_as_series.dt, name) | ||||||
else: | ||||||
values_as_cftimeindex = CFTimeIndex(values.ravel()) | ||||||
values_as_cftimeindex = CFTimeIndex(duck_array_ops.ravel(values)) | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
method = getattr(values_as_cftimeindex, name) | ||||||
|
||||||
field_values = method(freq=freq).values | ||||||
|
@@ -195,7 +196,7 @@ def _strftime_through_cftimeindex(values, date_format: str): | |||||
""" | ||||||
from xarray.coding.cftimeindex import CFTimeIndex | ||||||
|
||||||
values_as_cftimeindex = CFTimeIndex(values.ravel()) | ||||||
values_as_cftimeindex = CFTimeIndex(duck_array_ops.ravel(values)) | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
|
||||||
field_values = values_as_cftimeindex.strftime(date_format) | ||||||
return field_values.values.reshape(values.shape) | ||||||
|
@@ -205,7 +206,7 @@ def _strftime_through_series(values, date_format: str): | |||||
"""Coerce an array of datetime-like values to a pandas Series and | ||||||
apply string formatting | ||||||
""" | ||||||
values_as_series = pd.Series(values.ravel(), copy=False) | ||||||
values_as_series = pd.Series(duck_array_ops.ravel(values), copy=False) | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
strs = values_as_series.dt.strftime(date_format) | ||||||
return strs.values.reshape(values.shape) | ||||||
|
||||||
|
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
|
@@ -2123,7 +2123,8 @@ def _calc_idxminmax( | |||||
chunkmanager = get_chunked_array_type(array.data) | ||||||
chunks = dict(zip(array.dims, array.chunks)) | ||||||
dask_coord = chunkmanager.from_array(array[dim].data, chunks=chunks[dim]) | ||||||
res = indx.copy(data=dask_coord[indx.data.ravel()].reshape(indx.shape)) | ||||||
data = dask_coord[duck_array_ops.ravel(indx.data)] | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
res = indx.copy(data=duck_array_ops.reshape(data, indx.shape)) | ||||||
# we need to attach back the dim name | ||||||
res.name = dim | ||||||
else: | ||||||
|
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
|
@@ -337,6 +337,10 @@ def reshape(array, shape): | |||||
return xp.reshape(array, shape) | ||||||
|
||||||
|
||||||
def ravel(array): | ||||||
return reshape(array, (-1,)) | ||||||
|
||||||
|
||||||
Comment on lines
+340
to
+343
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
@contextlib.contextmanager | ||||||
def _ignore_warnings_if(condition): | ||||||
if condition: | ||||||
|
@@ -363,7 +367,7 @@ def f(values, axis=None, skipna=None, **kwargs): | |||||
values = asarray(values) | ||||||
|
||||||
if coerce_strings and values.dtype.kind in "SU": | ||||||
values = values.astype(object) | ||||||
values = astype(values, object) | ||||||
|
||||||
func = None | ||||||
if skipna or (skipna is None and values.dtype.kind in "cfO"): | ||||||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -6,6 +6,7 @@ | |
|
||
import xarray as xr | ||
from xarray import DataArray, Dataset, set_options | ||
from xarray.core import duck_array_ops | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This effectively changes the test from only testing public xarray API to now relying on some internals too, which is not ideal from a code separation perspective. In practice I think this is probably fine here, and I know you're doing it to allow Cubed arrays to go through the modified tests, but just something to bear in mind. |
||
from xarray.tests import ( | ||
assert_allclose, | ||
assert_equal, | ||
|
@@ -272,21 +273,24 @@ def test_coarsen_construct(self, dask: bool) -> None: | |
expected = xr.Dataset(attrs={"foo": "bar"}) | ||
expected["vart"] = ( | ||
("year", "month"), | ||
ds.vart.data.reshape((-1, 12)), | ||
duck_array_ops.reshape(ds.vart.data, (-1, 12)), | ||
{"a": "b"}, | ||
) | ||
expected["varx"] = ( | ||
("x", "x_reshaped"), | ||
ds.varx.data.reshape((-1, 5)), | ||
duck_array_ops.reshape(ds.varx.data, (-1, 5)), | ||
{"a": "b"}, | ||
) | ||
expected["vartx"] = ( | ||
("x", "x_reshaped", "year", "month"), | ||
ds.vartx.data.reshape(2, 5, 4, 12), | ||
duck_array_ops.reshape(ds.vartx.data, (2, 5, 4, 12)), | ||
{"a": "b"}, | ||
) | ||
expected["vary"] = ds.vary | ||
expected.coords["time"] = (("year", "month"), ds.time.data.reshape((-1, 12))) | ||
expected.coords["time"] = ( | ||
("year", "month"), | ||
duck_array_ops.reshape(ds.time.data, (-1, 12)), | ||
) | ||
|
||
with raise_if_dask_computes(): | ||
actual = ds.coarsen(time=12, x=5).construct( | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can't we just get used to using reshape instead?