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.

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

_season_from_months can now handle np.nan #5876

Merged
merged 12 commits into from
Jan 11, 2022
4 changes: 4 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,10 @@ Bug fixes
- Fix applying function with non-xarray arguments using :py:func:`xr.map_blocks`.
By `Cindy Chiao <https://github.com/tcchiao>`_.

- `dt.season <https://xarray.pydata.org/en/stable/generated/xarray.DataArray.dt.season.html>`_ can now handle NaN and NaT. (:pull:`5876`).
By `Pierre Loicq <https://github.com/pierreloicq>`_ and `Spencer Clark <https://github.com/spencerkclark>`_.
pierreloicq marked this conversation as resolved.
Show resolved Hide resolved


Documentation
~~~~~~~~~~~~~

Expand Down
15 changes: 13 additions & 2 deletions xarray/core/accessor_dt.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,20 @@
def _season_from_months(months):
"""Compute season (DJF, MAM, JJA, SON) from month ordinal"""
# TODO: Move "season" accessor upstream into pandas
seasons = np.array(["DJF", "MAM", "JJA", "SON"])
seasons = np.array(["DJF", "MAM", "JJA", "SON", "nan"])
months = np.asarray(months)
return seasons[(months // 3) % 4]

with warnings.catch_warnings():
warnings.filterwarnings(
"ignore", message="invalid value encountered in floor_divide"
)
warnings.filterwarnings(
"ignore", message="invalid value encountered in remainder"
)
idx = (months // 3) % 4

idx[np.isnan(idx)] = 4
return seasons[idx.astype(int)]


def _access_through_cftimeindex(values, name):
Expand Down
2 changes: 2 additions & 0 deletions xarray/tests/test_accessor_dt.py
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ def test_dask_accessor_method(self, method, parameters) -> None:

def test_seasons(self) -> None:
dates = pd.date_range(start="2000/01/01", freq="M", periods=12)
dates = dates.append(pd.Index([np.datetime64("NaT")]))
dates = xr.DataArray(dates)
seasons = xr.DataArray(
[
Expand All @@ -237,6 +238,7 @@ def test_seasons(self) -> None:
"SON",
"SON",
"DJF",
"nan",
]
)

Expand Down