from __future__ import annotations import datetime from typing import TypedDict import numpy as np import pandas as pd import pytest from packaging.version import Version import xarray as xr from xarray.coding.cftime_offsets import _new_to_legacy_freq from xarray.coding.cftimeindex import CFTimeIndex from xarray.core.resample_cftime import CFTimeGrouper cftime = pytest.importorskip("cftime") # Create a list of pairs of similar-length initial and resample frequencies # that cover: # - Resampling from shorter to longer frequencies # - Resampling from longer to shorter frequencies # - Resampling from one initial frequency to another. # These are used to test the cftime version of resample against pandas # with a standard calendar. FREQS = [ ("8003D", "4001D"), ("8003D", "16006D"), ("8003D", "21YS"), ("6h", "3h"), ("6h", "12h"), ("6h", "400min"), ("3D", "D"), ("3D", "6D"), ("11D", "MS"), ("3MS", "MS"), ("3MS", "6MS"), ("3MS", "85D"), ("7ME", "3ME"), ("7ME", "14ME"), ("7ME", "2QS-APR"), ("43QS-AUG", "21QS-AUG"), ("43QS-AUG", "86QS-AUG"), ("43QS-AUG", "11YE-JUN"), ("11QE-JUN", "5QE-JUN"), ("11QE-JUN", "22QE-JUN"), ("11QE-JUN", "51MS"), ("3YS-MAR", "YS-MAR"), ("3YS-MAR", "6YS-MAR"), ("3YS-MAR", "14QE-FEB"), ("7YE-MAY", "3YE-MAY"), ("7YE-MAY", "14YE-MAY"), ("7YE-MAY", "85ME"), ] def compare_against_pandas( da_datetimeindex, da_cftimeindex, freq, closed=None, label=None, offset=None, origin=None, ) -> None: if isinstance(origin, tuple): origin_pandas = pd.Timestamp(datetime.datetime(*origin)) origin_cftime = cftime.DatetimeGregorian(*origin) else: origin_pandas = origin origin_cftime = origin try: result_datetimeindex = da_datetimeindex.resample( time=freq, closed=closed, label=label, offset=offset, origin=origin_pandas, ).mean() except ValueError: with pytest.raises(ValueError): da_cftimeindex.resample( time=freq, closed=closed, label=label, origin=origin_cftime, offset=offset, ).mean() else: result_cftimeindex = da_cftimeindex.resample( time=freq, closed=closed, label=label, origin=origin_cftime, offset=offset, ).mean() # TODO (benbovy - flexible indexes): update when CFTimeIndex is a xarray Index subclass result_cftimeindex["time"] = ( result_cftimeindex.xindexes["time"] .to_pandas_index() .to_datetimeindex(time_unit="ns") ) xr.testing.assert_identical(result_cftimeindex, result_datetimeindex) def da(index) -> xr.DataArray: return xr.DataArray( np.arange(100.0, 100.0 + index.size), coords=[index], dims=["time"] ) @pytest.mark.parametrize("freqs", FREQS, ids=lambda x: "{}->{}".format(*x)) @pytest.mark.parametrize("closed", [None, "left", "right"]) @pytest.mark.parametrize("label", [None, "left", "right"]) @pytest.mark.parametrize("offset", [None, "5s"], ids=lambda x: f"{x}") def test_resample(freqs, closed, label, offset) -> None: initial_freq, resample_freq = freqs if ( resample_freq == "4001D" and closed == "right" and Version(pd.__version__) < Version("2.2") ): pytest.skip( "Pandas fixed a bug in this test case in version 2.2, which we " "ported to xarray, so this test no longer produces the same " "result as pandas for earlier pandas versions." ) start = "2000-01-01T12:07:01" origin = "start" datetime_index = pd.date_range( start=start, periods=5, freq=_new_to_legacy_freq(initial_freq) ) cftime_index = xr.date_range( start=start, periods=5, freq=initial_freq, use_cftime=True ) da_datetimeindex = da(datetime_index) da_cftimeindex = da(cftime_index) compare_against_pandas( da_datetimeindex, da_cftimeindex, resample_freq, closed=closed, label=label, offset=offset, origin=origin, ) @pytest.mark.parametrize( ("freq", "expected"), [ ("s", "left"), ("min", "left"), ("h", "left"), ("D", "left"), ("ME", "right"), ("MS", "left"), ("QE", "right"), ("QS", "left"), ("YE", "right"), ("YS", "left"), ], ) def test_closed_label_defaults(freq, expected) -> None: assert CFTimeGrouper(freq=freq).closed == expected assert CFTimeGrouper(freq=freq).label == expected @pytest.mark.filterwarnings("ignore:Converting a CFTimeIndex") @pytest.mark.parametrize( "calendar", ["gregorian", "noleap", "all_leap", "360_day", "julian"] ) def test_calendars(calendar: str) -> None: # Limited testing for non-standard calendars freq, closed, label = "8001min", None, None xr_index = xr.date_range( start="2004-01-01T12:07:01", periods=7, freq="3D", calendar=calendar, use_cftime=True, ) pd_index = pd.date_range(start="2004-01-01T12:07:01", periods=7, freq="3D") da_cftime = da(xr_index).resample(time=freq, closed=closed, label=label).mean() da_datetime = da(pd_index).resample(time=freq, closed=closed, label=label).mean() # TODO (benbovy - flexible indexes): update when CFTimeIndex is a xarray Index subclass new_pd_index = da_cftime.xindexes["time"].to_pandas_index() assert isinstance(new_pd_index, CFTimeIndex) # shouldn't that be a pd.Index? da_cftime["time"] = new_pd_index.to_datetimeindex(time_unit="ns") xr.testing.assert_identical(da_cftime, da_datetime) class DateRangeKwargs(TypedDict): start: str periods: int freq: str @pytest.mark.parametrize("closed", ["left", "right"]) @pytest.mark.parametrize( "origin", ["start_day", "start", "end", "end_day", "epoch", (1970, 1, 1, 3, 2)], ids=lambda x: f"{x}", ) def test_origin(closed, origin) -> None: initial_freq, resample_freq = ("3h", "9h") start = "1969-12-31T12:07:01" index_kwargs: DateRangeKwargs = dict(start=start, periods=12, freq=initial_freq) datetime_index = pd.date_range(**index_kwargs) cftime_index = xr.date_range(**index_kwargs, use_cftime=True) da_datetimeindex = da(datetime_index) da_cftimeindex = da(cftime_index) compare_against_pandas( da_datetimeindex, da_cftimeindex, resample_freq, closed=closed, origin=origin, ) @pytest.mark.parametrize("offset", ["foo", "5MS", 10]) def test_invalid_offset_error(offset: str | int) -> None: cftime_index = xr.date_range("2000", periods=5, use_cftime=True) da_cftime = da(cftime_index) with pytest.raises(ValueError, match="offset must be"): da_cftime.resample(time="2D", offset=offset) # type: ignore[arg-type] def test_timedelta_offset() -> None: timedelta = datetime.timedelta(seconds=5) string = "5s" cftime_index = xr.date_range("2000", periods=5, use_cftime=True) da_cftime = da(cftime_index) timedelta_result = da_cftime.resample(time="2D", offset=timedelta).mean() string_result = da_cftime.resample(time="2D", offset=string).mean() xr.testing.assert_identical(timedelta_result, string_result)