diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 9ac671d5858..090b6022b05 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -68,6 +68,8 @@ Enhancements - :py:meth:`pandas.Series.dropna` is now supported for a :py:class:`pandas.Series` indexed by a :py:class:`~xarray.CFTimeIndex` (:issue:`2688`). By `Spencer Clark `_. +- :py:meth:`~xarray.cftime_range` now supports QuarterBegin and QuarterEnd offsets (:issue:`2663`). + By `Jwen Fai Low `_ - :py:meth:`~xarray.open_dataset` now accepts a ``use_cftime`` argument, which can be used to require that ``cftime.datetime`` objects are always used, or never used when decoding dates encoded with a standard calendar. This can be @@ -76,7 +78,7 @@ Enhancements serialization warnings raised if dates from a standard calendar are found to be outside the :py:class:`pandas.Timestamp`-valid range (:issue:`2754`). By `Spencer Clark `_. - + Bug fixes ~~~~~~~~~ diff --git a/xarray/coding/cftime_offsets.py b/xarray/coding/cftime_offsets.py index 4b5770ac90a..a74c735224b 100644 --- a/xarray/coding/cftime_offsets.py +++ b/xarray/coding/cftime_offsets.py @@ -75,6 +75,7 @@ def get_date_type(calendar): class BaseCFTimeOffset(object): _freq = None # type: ClassVar[str] + _day_option = None def __init__(self, n=1): if not isinstance(n, int): @@ -151,6 +152,41 @@ def __str__(self): def __repr__(self): return str(self) + def _get_offset_day(self, other): + # subclass must implement `_day_option`; calling from the base class + # will raise NotImplementedError. + return _get_day_of_month(other, self._day_option) + + +def _get_day_of_month(other, day_option): + """Find the day in `other`'s month that satisfies a BaseCFTimeOffset's + onOffset policy, as described by the `day_option` argument. + + Parameters + ---------- + other : cftime.datetime + day_option : 'start', 'end' + 'start': returns 1 + 'end': returns last day of the month + + Returns + ------- + day_of_month : int + + """ + + if day_option == 'start': + return 1 + elif day_option == 'end': + days_in_month = _days_in_month(other) + return days_in_month + elif day_option is None: + # Note: unlike `_shift_month`, _get_day_of_month does not + # allow day_option = None + raise NotImplementedError + else: + raise ValueError(day_option) + def _days_in_month(date): """The number of days in the month of the given date""" @@ -186,7 +222,7 @@ def _adjust_n_years(other, n, month, reference_day): return n -def _shift_months(date, months, day_option='start'): +def _shift_month(date, months, day_option='start'): """Shift the date to a month start or end a given number of months away. """ delta_year = (date.month + months) // 12 @@ -211,12 +247,69 @@ def _shift_months(date, months, day_option='start'): return date.replace(year=year, month=month, day=day, dayofwk=-1) +def roll_qtrday(other, n, month, day_option, modby=3): + """Possibly increment or decrement the number of periods to shift + based on rollforward/rollbackward conventions. + + Parameters + ---------- + other : cftime.datetime + n : number of periods to increment, before adjusting for rolling + month : int reference month giving the first month of the year + day_option : 'start', 'end' + The convention to use in finding the day in a given month against + which to compare for rollforward/rollbackward decisions. + modby : int 3 for quarters, 12 for years + + Returns + ------- + n : int number of periods to increment + + See Also + -------- + _get_day_of_month : Find the day in a month provided an offset. + """ + + months_since = other.month % modby - month % modby + + if n > 0: + if months_since < 0 or ( + months_since == 0 and + other.day < _get_day_of_month(other, day_option)): + # pretend to roll back if on same month but + # before compare_day + n -= 1 + else: + if months_since > 0 or ( + months_since == 0 and + other.day > _get_day_of_month(other, day_option)): + # make sure to roll forward, so negate + n += 1 + return n + + +def _validate_month(month, default_month): + if month is None: + result_month = default_month + else: + result_month = month + if not isinstance(result_month, int): + raise TypeError("'self.month' must be an integer value between 1 " + "and 12. Instead, it was set to a value of " + "{!r}".format(result_month)) + elif not (1 <= result_month <= 12): + raise ValueError("'self.month' must be an integer value between 1 " + "and 12. Instead, it was set to a value of " + "{!r}".format(result_month)) + return result_month + + class MonthBegin(BaseCFTimeOffset): _freq = 'MS' def __apply__(self, other): n = _adjust_n_months(other.day, self.n, 1) - return _shift_months(other, n, 'start') + return _shift_month(other, n, 'start') def onOffset(self, date): """Check if the given date is in the set of possible dates created @@ -229,7 +322,7 @@ class MonthEnd(BaseCFTimeOffset): def __apply__(self, other): n = _adjust_n_months(other.day, self.n, _days_in_month(other)) - return _shift_months(other, n, 'end') + return _shift_month(other, n, 'end') def onOffset(self, date): """Check if the given date is in the set of possible dates created @@ -253,6 +346,105 @@ def onOffset(self, date): } +class QuarterOffset(BaseCFTimeOffset): + """Quarter representation copied off of pandas/tseries/offsets.py + """ + _freq = None # type: ClassVar[str] + _default_month = None # type: ClassVar[int] + + def __init__(self, n=1, month=None): + BaseCFTimeOffset.__init__(self, n) + self.month = _validate_month(month, self._default_month) + + def __apply__(self, other): + # months_since: find the calendar quarter containing other.month, + # e.g. if other.month == 8, the calendar quarter is [Jul, Aug, Sep]. + # Then find the month in that quarter containing an onOffset date for + # self. `months_since` is the number of months to shift other.month + # to get to this on-offset month. + months_since = other.month % 3 - self.month % 3 + qtrs = roll_qtrday(other, self.n, self.month, + day_option=self._day_option, modby=3) + months = qtrs * 3 - months_since + return _shift_month(other, months, self._day_option) + + def onOffset(self, date): + """Check if the given date is in the set of possible dates created + using a length-one version of this offset class.""" + mod_month = (date.month - self.month) % 3 + return mod_month == 0 and date.day == self._get_offset_day(date) + + def __sub__(self, other): + import cftime + + if isinstance(other, cftime.datetime): + raise TypeError('Cannot subtract cftime.datetime from offset.') + elif type(other) == type(self) and other.month == self.month: + return type(self)(self.n - other.n, month=self.month) + else: + return NotImplemented + + def __mul__(self, other): + return type(self)(n=other * self.n, month=self.month) + + def rule_code(self): + return '{}-{}'.format(self._freq, _MONTH_ABBREVIATIONS[self.month]) + + def __str__(self): + return '<{}: n={}, month={}>'.format( + type(self).__name__, self.n, self.month) + + +class QuarterBegin(QuarterOffset): + # When converting a string to an offset, pandas converts + # 'QS' to a QuarterBegin offset starting in the month of + # January. When creating a QuarterBegin offset directly + # from the constructor, however, the default month is March. + # We follow that behavior here. + _default_month = 3 + _freq = 'QS' + _day_option = 'start' + + def rollforward(self, date): + """Roll date forward to nearest start of quarter""" + if self.onOffset(date): + return date + else: + return date + QuarterBegin(month=self.month) + + def rollback(self, date): + """Roll date backward to nearest start of quarter""" + if self.onOffset(date): + return date + else: + return date - QuarterBegin(month=self.month) + + +class QuarterEnd(QuarterOffset): + # When converting a string to an offset, pandas converts + # 'Q' to a QuarterEnd offset starting in the month of + # December. When creating a QuarterEnd offset directly + # from the constructor, however, the default month is March. + # We follow that behavior here. + _default_month = 3 + _freq = 'Q' + _day_option = 'end' + + def rollforward(self, date): + """Roll date forward to nearest end of quarter""" + if self.onOffset(date): + return date + else: + return date + QuarterEnd(month=self.month) + + def rollback(self, date): + """Roll date backward to nearest end of quarter""" + if self.onOffset(date): + return date + else: + return date - QuarterEnd(month=self.month) + + class YearOffset(BaseCFTimeOffset): _freq = None # type: ClassVar[str] _day_option = None # type: ClassVar[str] @@ -260,29 +452,13 @@ class YearOffset(BaseCFTimeOffset): def __init__(self, n=1, month=None): BaseCFTimeOffset.__init__(self, n) - if month is None: - self.month = self._default_month - else: - self.month = month - if not isinstance(self.month, int): - raise TypeError("'self.month' must be an integer value between 1 " - "and 12. Instead, it was set to a value of " - "{!r}".format(self.month)) - elif not (1 <= self.month <= 12): - raise ValueError("'self.month' must be an integer value between 1 " - "and 12. Instead, it was set to a value of " - "{!r}".format(self.month)) + self.month = _validate_month(month, self._default_month) def __apply__(self, other): - if self._day_option == 'start': - reference_day = 1 - elif self._day_option == 'end': - reference_day = _days_in_month(other) - else: - raise ValueError(self._day_option) + reference_day = _get_day_of_month(other, self._day_option) years = _adjust_n_years(other, self.n, self.month, reference_day) months = years * 12 + (self.month - other.month) - return _shift_months(other, months, self._day_option) + return _shift_month(other, months, self._day_option) def __sub__(self, other): import cftime @@ -400,6 +576,8 @@ def __apply__(self, other): 'AS': YearBegin, 'Y': YearEnd, 'YS': YearBegin, + 'Q': partial(QuarterEnd, month=12), + 'QS': partial(QuarterBegin, month=1), 'M': MonthEnd, 'MS': MonthBegin, 'D': Day, @@ -430,7 +608,31 @@ def __apply__(self, other): 'A-SEP': partial(YearEnd, month=9), 'A-OCT': partial(YearEnd, month=10), 'A-NOV': partial(YearEnd, month=11), - 'A-DEC': partial(YearEnd, month=12) + 'A-DEC': partial(YearEnd, month=12), + 'QS-JAN': partial(QuarterBegin, month=1), + 'QS-FEB': partial(QuarterBegin, month=2), + 'QS-MAR': partial(QuarterBegin, month=3), + 'QS-APR': partial(QuarterBegin, month=4), + 'QS-MAY': partial(QuarterBegin, month=5), + 'QS-JUN': partial(QuarterBegin, month=6), + 'QS-JUL': partial(QuarterBegin, month=7), + 'QS-AUG': partial(QuarterBegin, month=8), + 'QS-SEP': partial(QuarterBegin, month=9), + 'QS-OCT': partial(QuarterBegin, month=10), + 'QS-NOV': partial(QuarterBegin, month=11), + 'QS-DEC': partial(QuarterBegin, month=12), + 'Q-JAN': partial(QuarterEnd, month=1), + 'Q-FEB': partial(QuarterEnd, month=2), + 'Q-MAR': partial(QuarterEnd, month=3), + 'Q-APR': partial(QuarterEnd, month=4), + 'Q-MAY': partial(QuarterEnd, month=5), + 'Q-JUN': partial(QuarterEnd, month=6), + 'Q-JUL': partial(QuarterEnd, month=7), + 'Q-AUG': partial(QuarterEnd, month=8), + 'Q-SEP': partial(QuarterEnd, month=9), + 'Q-OCT': partial(QuarterEnd, month=10), + 'Q-NOV': partial(QuarterEnd, month=11), + 'Q-DEC': partial(QuarterEnd, month=12) } @@ -624,55 +826,84 @@ def cftime_range(start=None, end=None, periods=None, freq='D', Valid simple frequency strings for use with ``cftime``-calendars include any multiples of the following. - +--------+-----------------------+ - | Alias | Description | - +========+=======================+ - | A, Y | Year-end frequency | - +--------+-----------------------+ - | AS, YS | Year-start frequency | - +--------+-----------------------+ - | M | Month-end frequency | - +--------+-----------------------+ - | MS | Month-start frequency | - +--------+-----------------------+ - | D | Day frequency | - +--------+-----------------------+ - | H | Hour frequency | - +--------+-----------------------+ - | T, min | Minute frequency | - +--------+-----------------------+ - | S | Second frequency | - +--------+-----------------------+ + +--------+--------------------------+ + | Alias | Description | + +========+==========================+ + | A, Y | Year-end frequency | + +--------+--------------------------+ + | AS, YS | Year-start frequency | + +--------+--------------------------+ + | Q | Quarter-end frequency | + +--------+--------------------------+ + | QS | Quarter-start frequency | + +--------+--------------------------+ + | M | Month-end frequency | + +--------+--------------------------+ + | MS | Month-start frequency | + +--------+--------------------------+ + | D | Day frequency | + +--------+--------------------------+ + | H | Hour frequency | + +--------+--------------------------+ + | T, min | Minute frequency | + +--------+--------------------------+ + | S | Second frequency | + +--------+--------------------------+ Any multiples of the following anchored offsets are also supported. - +----------+-------------------------------------------------------------------+ - | Alias | Description | - +==========+===================================================================+ - | A(S)-JAN | Annual frequency, anchored at the end (or beginning) of January | - +----------+-------------------------------------------------------------------+ - | A(S)-FEB | Annual frequency, anchored at the end (or beginning) of February | - +----------+-------------------------------------------------------------------+ - | A(S)-MAR | Annual frequency, anchored at the end (or beginning) of March | - +----------+-------------------------------------------------------------------+ - | A(S)-APR | Annual frequency, anchored at the end (or beginning) of April | - +----------+-------------------------------------------------------------------+ - | A(S)-MAY | Annual frequency, anchored at the end (or beginning) of May | - +----------+-------------------------------------------------------------------+ - | A(S)-JUN | Annual frequency, anchored at the end (or beginning) of June | - +----------+-------------------------------------------------------------------+ - | A(S)-JUL | Annual frequency, anchored at the end (or beginning) of July | - +----------+-------------------------------------------------------------------+ - | A(S)-AUG | Annual frequency, anchored at the end (or beginning) of August | - +----------+-------------------------------------------------------------------+ - | A(S)-SEP | Annual frequency, anchored at the end (or beginning) of September | - +----------+-------------------------------------------------------------------+ - | A(S)-OCT | Annual frequency, anchored at the end (or beginning) of October | - +----------+-------------------------------------------------------------------+ - | A(S)-NOV | Annual frequency, anchored at the end (or beginning) of November | - +----------+-------------------------------------------------------------------+ - | A(S)-DEC | Annual frequency, anchored at the end (or beginning) of December | - +----------+-------------------------------------------------------------------+ + +----------+--------------------------------------------------------------------+ + | Alias | Description | + +==========+====================================================================+ + | A(S)-JAN | Annual frequency, anchored at the end (or beginning) of January | + +----------+--------------------------------------------------------------------+ + | A(S)-FEB | Annual frequency, anchored at the end (or beginning) of February | + +----------+--------------------------------------------------------------------+ + | A(S)-MAR | Annual frequency, anchored at the end (or beginning) of March | + +----------+--------------------------------------------------------------------+ + | A(S)-APR | Annual frequency, anchored at the end (or beginning) of April | + +----------+--------------------------------------------------------------------+ + | A(S)-MAY | Annual frequency, anchored at the end (or beginning) of May | + +----------+--------------------------------------------------------------------+ + | A(S)-JUN | Annual frequency, anchored at the end (or beginning) of June | + +----------+--------------------------------------------------------------------+ + | A(S)-JUL | Annual frequency, anchored at the end (or beginning) of July | + +----------+--------------------------------------------------------------------+ + | A(S)-AUG | Annual frequency, anchored at the end (or beginning) of August | + +----------+--------------------------------------------------------------------+ + | A(S)-SEP | Annual frequency, anchored at the end (or beginning) of September | + +----------+--------------------------------------------------------------------+ + | A(S)-OCT | Annual frequency, anchored at the end (or beginning) of October | + +----------+--------------------------------------------------------------------+ + | A(S)-NOV | Annual frequency, anchored at the end (or beginning) of November | + +----------+--------------------------------------------------------------------+ + | A(S)-DEC | Annual frequency, anchored at the end (or beginning) of December | + +----------+--------------------------------------------------------------------+ + | Q(S)-JAN | Quarter frequency, anchored at the end (or beginning) of January | + +----------+--------------------------------------------------------------------+ + | Q(S)-FEB | Quarter frequency, anchored at the end (or beginning) of February | + +----------+--------------------------------------------------------------------+ + | Q(S)-MAR | Quarter frequency, anchored at the end (or beginning) of March | + +----------+--------------------------------------------------------------------+ + | Q(S)-APR | Quarter frequency, anchored at the end (or beginning) of April | + +----------+--------------------------------------------------------------------+ + | Q(S)-MAY | Quarter frequency, anchored at the end (or beginning) of May | + +----------+--------------------------------------------------------------------+ + | Q(S)-JUN | Quarter frequency, anchored at the end (or beginning) of June | + +----------+--------------------------------------------------------------------+ + | Q(S)-JUL | Quarter frequency, anchored at the end (or beginning) of July | + +----------+--------------------------------------------------------------------+ + | Q(S)-AUG | Quarter frequency, anchored at the end (or beginning) of August | + +----------+--------------------------------------------------------------------+ + | Q(S)-SEP | Quarter frequency, anchored at the end (or beginning) of September | + +----------+--------------------------------------------------------------------+ + | Q(S)-OCT | Quarter frequency, anchored at the end (or beginning) of October | + +----------+--------------------------------------------------------------------+ + | Q(S)-NOV | Quarter frequency, anchored at the end (or beginning) of November | + +----------+--------------------------------------------------------------------+ + | Q(S)-DEC | Quarter frequency, anchored at the end (or beginning) of December | + +----------+--------------------------------------------------------------------+ + Finally, the following calendar aliases are supported. diff --git a/xarray/core/resample_cftime.py b/xarray/core/resample_cftime.py index 6b6d214768e..161945f118d 100644 --- a/xarray/core/resample_cftime.py +++ b/xarray/core/resample_cftime.py @@ -38,7 +38,7 @@ from ..coding.cftimeindex import CFTimeIndex from ..coding.cftime_offsets import (cftime_range, normalize_date, - Day, MonthEnd, YearEnd, + Day, MonthEnd, QuarterEnd, YearEnd, CFTIME_TICKS, to_offset) import datetime import numpy as np @@ -50,14 +50,14 @@ class CFTimeGrouper(object): single method, the only one required for resampling in xarray. It cannot be used in a call to groupby like a pandas.Grouper object can.""" - def __init__(self, freq, closed, label, base, loffset): + def __init__(self, freq, closed=None, label=None, base=0, loffset=None): self.freq = to_offset(freq) self.closed = closed self.label = label self.base = base self.loffset = loffset - if isinstance(self.freq, (MonthEnd, YearEnd)): + if isinstance(self.freq, (MonthEnd, QuarterEnd, YearEnd)): if self.closed is None: self.closed = 'right' if self.label is None: @@ -199,7 +199,7 @@ def _adjust_bin_edges(datetime_bins, offset, closed, index, labels): This is also required for daily frequencies longer than one day and year-end frequencies. """ - is_super_daily = (isinstance(offset, (MonthEnd, YearEnd)) or + is_super_daily = (isinstance(offset, (MonthEnd, QuarterEnd, YearEnd)) or (isinstance(offset, Day) and offset.n > 1)) if is_super_daily: if closed == 'right': diff --git a/xarray/tests/test_cftime_offsets.py b/xarray/tests/test_cftime_offsets.py index 29caa88cc53..1cf257c96eb 100644 --- a/xarray/tests/test_cftime_offsets.py +++ b/xarray/tests/test_cftime_offsets.py @@ -6,9 +6,9 @@ from xarray import CFTimeIndex from xarray.coding.cftime_offsets import ( - _MONTH_ABBREVIATIONS, BaseCFTimeOffset, Day, Hour, Minute, MonthBegin, - MonthEnd, Second, YearBegin, YearEnd, _days_in_month, cftime_range, - get_date_type, to_cftime_datetime, to_offset) + _MONTH_ABBREVIATIONS, BaseCFTimeOffset, Day, Hour, Minute, Second, + MonthBegin, MonthEnd, YearBegin, YearEnd, QuarterBegin, QuarterEnd, + _days_in_month, cftime_range, get_date_type, to_cftime_datetime, to_offset) cftime = pytest.importorskip('cftime') @@ -32,9 +32,13 @@ def calendar(request): [(BaseCFTimeOffset(), 1), (YearBegin(), 1), (YearEnd(), 1), + (QuarterBegin(), 1), + (QuarterEnd(), 1), (BaseCFTimeOffset(n=2), 2), (YearBegin(n=2), 2), - (YearEnd(n=2), 2)], + (YearEnd(n=2), 2), + (QuarterBegin(n=2), 2), + (QuarterEnd(n=2), 2)], ids=_id_func ) def test_cftime_offset_constructor_valid_n(offset, expected_n): @@ -45,7 +49,9 @@ def test_cftime_offset_constructor_valid_n(offset, expected_n): ('offset', 'invalid_n'), [(BaseCFTimeOffset, 1.5), (YearBegin, 1.5), - (YearEnd, 1.5)], + (YearEnd, 1.5), + (QuarterBegin, 1.5), + (QuarterEnd, 1.5)], ids=_id_func ) def test_cftime_offset_constructor_invalid_n(offset, invalid_n): @@ -58,7 +64,11 @@ def test_cftime_offset_constructor_invalid_n(offset, invalid_n): [(YearBegin(), 1), (YearEnd(), 12), (YearBegin(month=5), 5), - (YearEnd(month=5), 5)], + (YearEnd(month=5), 5), + (QuarterBegin(), 3), + (QuarterEnd(), 3), + (QuarterBegin(month=5), 5), + (QuarterEnd(month=5), 5)], ids=_id_func ) def test_year_offset_constructor_valid_month(offset, expected_month): @@ -72,7 +82,13 @@ def test_year_offset_constructor_valid_month(offset, expected_month): (YearBegin, 13, ValueError,), (YearEnd, 13, ValueError), (YearBegin, 1.5, TypeError), - (YearEnd, 1.5, TypeError)], + (YearEnd, 1.5, TypeError), + (QuarterBegin, 0, ValueError), + (QuarterEnd, 0, ValueError), + (QuarterBegin, 1.5, TypeError), + (QuarterEnd, 1.5, TypeError), + (QuarterBegin, 13, ValueError), + (QuarterEnd, 13, ValueError)], ids=_id_func ) def test_year_offset_constructor_invalid_month( @@ -85,7 +101,8 @@ def test_year_offset_constructor_invalid_month( ('offset', 'expected'), [(BaseCFTimeOffset(), None), (MonthBegin(), 'MS'), - (YearBegin(), 'AS-JAN')], + (YearBegin(), 'AS-JAN'), + (QuarterBegin(), 'QS-MAR')], ids=_id_func ) def test_rule_code(offset, expected): @@ -95,7 +112,8 @@ def test_rule_code(offset, expected): @pytest.mark.parametrize( ('offset', 'expected'), [(BaseCFTimeOffset(), ''), - (YearBegin(), '')], + (YearBegin(), ''), + (QuarterBegin(), '')], ids=_id_func ) def test_str_and_repr(offset, expected): @@ -105,7 +123,7 @@ def test_str_and_repr(offset, expected): @pytest.mark.parametrize( 'offset', - [BaseCFTimeOffset(), MonthBegin(), YearBegin()], + [BaseCFTimeOffset(), MonthBegin(), QuarterBegin(), YearBegin()], ids=_id_func ) def test_to_offset_offset_input(offset): @@ -164,7 +182,47 @@ def test_to_offset_annual(month_label, month_int, multiple, offset_str): assert result == expected -@pytest.mark.parametrize('freq', ['Z', '7min2', 'AM', 'M-', 'AS-', '1H1min']) +_QUARTER_OFFSET_TYPES = { + 'Q': QuarterEnd, + 'QS': QuarterBegin +} + + +@pytest.mark.parametrize(('month_int', 'month_label'), + list(_MONTH_ABBREVIATIONS.items()) + [(0, '')]) +@pytest.mark.parametrize('multiple', [None, 2]) +@pytest.mark.parametrize('offset_str', ['QS', 'Q']) +def test_to_offset_quarter(month_label, month_int, multiple, offset_str): + freq = offset_str + offset_type = _QUARTER_OFFSET_TYPES[offset_str] + if month_label: + freq = '-'.join([freq, month_label]) + if multiple: + freq = '{}'.format(multiple) + freq + result = to_offset(freq) + + if multiple and month_int: + expected = offset_type(n=multiple, month=month_int) + elif multiple: + if month_int: + expected = offset_type(n=multiple) + else: + if offset_type == QuarterBegin: + expected = offset_type(n=multiple, month=1) + elif offset_type == QuarterEnd: + expected = offset_type(n=multiple, month=12) + elif month_int: + expected = offset_type(month=month_int) + else: + if offset_type == QuarterBegin: + expected = offset_type(month=1) + elif offset_type == QuarterEnd: + expected = offset_type(month=12) + assert result == expected + + +@pytest.mark.parametrize('freq', ['Z', '7min2', 'AM', 'M-', 'AS-', 'QS-', + '1H1min']) def test_invalid_to_offset_str(freq): with pytest.raises(ValueError): to_offset(freq) @@ -197,13 +255,16 @@ def test_to_cftime_datetime_error_type_error(): _EQ_TESTS_A = [ BaseCFTimeOffset(), YearBegin(), YearEnd(), YearBegin(month=2), - YearEnd(month=2), MonthBegin(), MonthEnd(), Day(), Hour(), Minute(), + YearEnd(month=2), QuarterBegin(), QuarterEnd(), QuarterBegin(month=2), + QuarterEnd(month=2), MonthBegin(), MonthEnd(), Day(), Hour(), Minute(), Second() ] _EQ_TESTS_B = [ BaseCFTimeOffset(n=2), YearBegin(n=2), YearEnd(n=2), - YearBegin(n=2, month=2), YearEnd(n=2, month=2), MonthBegin(n=2), - MonthEnd(n=2), Day(n=2), Hour(n=2), Minute(n=2), Second(n=2) + YearBegin(n=2, month=2), YearEnd(n=2, month=2), QuarterBegin(n=2), + QuarterEnd(n=2), QuarterBegin(n=2, month=2), QuarterEnd(n=2, month=2), + MonthBegin(n=2), MonthEnd(n=2), Day(n=2), Hour(n=2), Minute(n=2), + Second(n=2) ] @@ -216,8 +277,10 @@ def test_neq(a, b): _EQ_TESTS_B_COPY = [ BaseCFTimeOffset(n=2), YearBegin(n=2), YearEnd(n=2), - YearBegin(n=2, month=2), YearEnd(n=2, month=2), MonthBegin(n=2), - MonthEnd(n=2), Day(n=2), Hour(n=2), Minute(n=2), Second(n=2) + YearBegin(n=2, month=2), YearEnd(n=2, month=2), QuarterBegin(n=2), + QuarterEnd(n=2), QuarterBegin(n=2, month=2), QuarterEnd(n=2, month=2), + MonthBegin(n=2), MonthEnd(n=2), Day(n=2), Hour(n=2), Minute(n=2), + Second(n=2) ] @@ -232,6 +295,8 @@ def test_eq(a, b): (BaseCFTimeOffset(), BaseCFTimeOffset(n=3)), (YearEnd(), YearEnd(n=3)), (YearBegin(), YearBegin(n=3)), + (QuarterEnd(), QuarterEnd(n=3)), + (QuarterBegin(), QuarterBegin(n=3)), (MonthEnd(), MonthEnd(n=3)), (MonthBegin(), MonthBegin(n=3)), (Day(), Day(n=3)), @@ -256,6 +321,8 @@ def test_rmul(offset, expected): [(BaseCFTimeOffset(), BaseCFTimeOffset(n=-1)), (YearEnd(), YearEnd(n=-1)), (YearBegin(), YearBegin(n=-1)), + (QuarterEnd(), QuarterEnd(n=-1)), + (QuarterBegin(), QuarterBegin(n=-1)), (MonthEnd(), MonthEnd(n=-1)), (MonthBegin(), MonthBegin(n=-1)), (Day(), Day(n=-1)), @@ -536,6 +603,89 @@ def test_add_year_end_onOffset( assert result == expected +@pytest.mark.parametrize( + ('initial_date_args', 'offset', 'expected_date_args'), + [((1, 1, 1), QuarterBegin(), (1, 3, 1)), + ((1, 1, 1), QuarterBegin(n=2), (1, 6, 1)), + ((1, 1, 1), QuarterBegin(month=2), (1, 2, 1)), + ((1, 1, 7), QuarterBegin(n=2), (1, 6, 1)), + ((2, 2, 1), QuarterBegin(n=-1), (1, 12, 1)), + ((1, 3, 2), QuarterBegin(n=-1), (1, 3, 1)), + ((1, 1, 1, 5, 5, 5, 5), QuarterBegin(), (1, 3, 1, 5, 5, 5, 5)), + ((2, 1, 1, 5, 5, 5, 5), QuarterBegin(n=-1), (1, 12, 1, 5, 5, 5, 5))], + ids=_id_func +) +def test_add_quarter_begin(calendar, initial_date_args, offset, + expected_date_args): + date_type = get_date_type(calendar) + initial = date_type(*initial_date_args) + result = initial + offset + expected = date_type(*expected_date_args) + assert result == expected + + +@pytest.mark.parametrize( + ('initial_date_args', 'offset', 'expected_year_month', + 'expected_sub_day'), + [((1, 1, 1), QuarterEnd(), (1, 3), ()), + ((1, 1, 1), QuarterEnd(n=2), (1, 6), ()), + ((1, 1, 1), QuarterEnd(month=1), (1, 1), ()), + ((2, 3, 1), QuarterEnd(n=-1), (1, 12), ()), + ((1, 3, 1), QuarterEnd(n=-1, month=2), (1, 2), ()), + ((1, 1, 1, 5, 5, 5, 5), QuarterEnd(), (1, 3), (5, 5, 5, 5)), + ((1, 1, 1, 5, 5, 5, 5), QuarterEnd(n=2), (1, 6), (5, 5, 5, 5))], + ids=_id_func +) +def test_add_quarter_end( + calendar, initial_date_args, offset, expected_year_month, + expected_sub_day +): + date_type = get_date_type(calendar) + initial = date_type(*initial_date_args) + result = initial + offset + reference_args = expected_year_month + (1,) + reference = date_type(*reference_args) + + # Here the days at the end of each month varies based on the calendar used + expected_date_args = (expected_year_month + + (_days_in_month(reference),) + expected_sub_day) + expected = date_type(*expected_date_args) + assert result == expected + + +@pytest.mark.parametrize( + ('initial_year_month', 'initial_sub_day', 'offset', 'expected_year_month', + 'expected_sub_day'), + [((1, 12), (), QuarterEnd(), (2, 3), ()), + ((1, 12), (), QuarterEnd(n=2), (2, 6), ()), + ((1, 12), (), QuarterEnd(n=-1), (1, 9), ()), + ((1, 12), (), QuarterEnd(n=-2), (1, 6), ()), + ((1, 1), (), QuarterEnd(month=2), (1, 2), ()), + ((1, 12), (5, 5, 5, 5), QuarterEnd(), (2, 3), (5, 5, 5, 5)), + ((1, 12), (5, 5, 5, 5), QuarterEnd(n=-1), (1, 9), (5, 5, 5, 5))], + ids=_id_func +) +def test_add_quarter_end_onOffset( + calendar, initial_year_month, initial_sub_day, offset, expected_year_month, + expected_sub_day +): + date_type = get_date_type(calendar) + reference_args = initial_year_month + (1,) + reference = date_type(*reference_args) + initial_date_args = (initial_year_month + (_days_in_month(reference),) + + initial_sub_day) + initial = date_type(*initial_date_args) + result = initial + offset + reference_args = expected_year_month + (1,) + reference = date_type(*reference_args) + + # Here the days at the end of each month varies based on the calendar used + expected_date_args = (expected_year_month + + (_days_in_month(reference),) + expected_sub_day) + expected = date_type(*expected_date_args) + assert result == expected + + # Note for all sub-monthly offsets, pandas always returns True for onOffset @pytest.mark.parametrize( ('date_args', 'offset', 'expected'), @@ -543,6 +693,10 @@ def test_add_year_end_onOffset( ((1, 1, 1, 1), MonthBegin(), True), ((1, 1, 5), MonthBegin(), False), ((1, 1, 5), MonthEnd(), False), + ((1, 3, 1), QuarterBegin(), True), + ((1, 3, 1, 1), QuarterBegin(), True), + ((1, 3, 5), QuarterBegin(), False), + ((1, 12, 1), QuarterEnd(), False), ((1, 1, 1), YearBegin(), True), ((1, 1, 1, 1), YearBegin(), True), ((1, 1, 5), YearBegin(), False), @@ -565,16 +719,19 @@ def test_onOffset(calendar, date_args, offset, expected): ('year_month_args', 'sub_day_args', 'offset'), [((1, 1), (), MonthEnd()), ((1, 1), (1,), MonthEnd()), + ((1, 12), (), QuarterEnd()), + ((1, 1), (), QuarterEnd(month=1)), ((1, 12), (), YearEnd()), ((1, 1), (), YearEnd(month=1))], ids=_id_func ) -def test_onOffset_month_or_year_end( +def test_onOffset_month_or_quarter_or_year_end( calendar, year_month_args, sub_day_args, offset): date_type = get_date_type(calendar) reference_args = year_month_args + (1,) reference = date_type(*reference_args) - date_args = year_month_args + (_days_in_month(reference),) + sub_day_args + date_args = (year_month_args + (_days_in_month(reference),) + + sub_day_args) date = date_type(*date_args) result = offset.onOffset(date) assert result @@ -590,6 +747,14 @@ def test_onOffset_month_or_year_end( (YearEnd(n=2), (1, 3, 1), (1, 12)), (YearEnd(n=2, month=2), (1, 3, 1), (2, 2)), (YearEnd(n=2, month=4), (1, 4, 30), (1, 4)), + (QuarterBegin(), (1, 3, 2), (1, 6)), + (QuarterBegin(), (1, 4, 1), (1, 6)), + (QuarterBegin(n=2), (1, 4, 1), (1, 6)), + (QuarterBegin(n=2, month=2), (1, 4, 1), (1, 5)), + (QuarterEnd(), (1, 3, 1), (1, 3)), + (QuarterEnd(n=2), (1, 3, 1), (1, 3)), + (QuarterEnd(n=2, month=2), (1, 3, 1), (1, 5)), + (QuarterEnd(n=2, month=4), (1, 4, 30), (1, 4)), (MonthBegin(), (1, 3, 2), (1, 4)), (MonthBegin(), (1, 3, 1), (1, 3)), (MonthBegin(n=2), (1, 3, 2), (1, 4)), @@ -606,9 +771,9 @@ def test_rollforward(calendar, offset, initial_date_args, partial_expected_date_args): date_type = get_date_type(calendar) initial = date_type(*initial_date_args) - if isinstance(offset, (MonthBegin, YearBegin)): + if isinstance(offset, (MonthBegin, QuarterBegin, YearBegin)): expected_date_args = partial_expected_date_args + (1,) - elif isinstance(offset, (MonthEnd, YearEnd)): + elif isinstance(offset, (MonthEnd, QuarterEnd, YearEnd)): reference_args = partial_expected_date_args + (1,) reference = date_type(*reference_args) expected_date_args = (partial_expected_date_args + @@ -631,6 +796,14 @@ def test_rollforward(calendar, offset, initial_date_args, (YearEnd(n=2), (2, 3, 1), (1, 12)), (YearEnd(n=2, month=2), (2, 3, 1), (2, 2)), (YearEnd(month=4), (1, 4, 30), (1, 4)), + (QuarterBegin(), (1, 3, 2), (1, 3)), + (QuarterBegin(), (1, 4, 1), (1, 3)), + (QuarterBegin(n=2), (1, 4, 1), (1, 3)), + (QuarterBegin(n=2, month=2), (1, 4, 1), (1, 2)), + (QuarterEnd(), (2, 3, 1), (1, 12)), + (QuarterEnd(n=2), (2, 3, 1), (1, 12)), + (QuarterEnd(n=2, month=2), (2, 3, 1), (2, 2)), + (QuarterEnd(n=2, month=4), (1, 4, 30), (1, 4)), (MonthBegin(), (1, 3, 2), (1, 3)), (MonthBegin(n=2), (1, 3, 2), (1, 3)), (MonthBegin(), (1, 3, 1), (1, 3)), @@ -647,9 +820,9 @@ def test_rollback(calendar, offset, initial_date_args, partial_expected_date_args): date_type = get_date_type(calendar) initial = date_type(*initial_date_args) - if isinstance(offset, (MonthBegin, YearBegin)): + if isinstance(offset, (MonthBegin, QuarterBegin, YearBegin)): expected_date_args = partial_expected_date_args + (1,) - elif isinstance(offset, (MonthEnd, YearEnd)): + elif isinstance(offset, (MonthEnd, QuarterEnd, YearEnd)): reference_args = partial_expected_date_args + (1,) reference = date_type(*reference_args) expected_date_args = (partial_expected_date_args + @@ -687,7 +860,9 @@ def test_rollback(calendar, offset, initial_date_args, ('0010', None, 4, YearBegin(n=-2), None, False, [(10, 1, 1), (8, 1, 1), (6, 1, 1), (4, 1, 1)]), ('0001-01-01', '0001-01-04', 4, None, None, False, - [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 1, 4)]) + [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 1, 4)]), + ('0001-06-01', None, 4, '3QS-JUN', None, False, + [(1, 6, 1), (2, 3, 1), (2, 12, 1), (3, 9, 1)]) ] diff --git a/xarray/tests/test_cftimeindex_resample.py b/xarray/tests/test_cftimeindex_resample.py index 0b56f1d1fc6..636f9ef7b0e 100644 --- a/xarray/tests/test_cftimeindex_resample.py +++ b/xarray/tests/test_cftimeindex_resample.py @@ -4,6 +4,7 @@ import numpy as np import pandas as pd import xarray as xr +from xarray.core.resample_cftime import CFTimeGrouper pytest.importorskip('cftime') pytest.importorskip('pandas', minversion='0.24') @@ -13,10 +14,10 @@ params=[ dict(start='2004-01-01T12:07:01', periods=91, freq='3D'), dict(start='1892-01-03T12:07:01', periods=15, freq='41987T'), - dict(start='2004-01-01T12:07:01', periods=31, freq='2MS'), + dict(start='2004-01-01T12:07:01', periods=7, freq='3Q-AUG'), dict(start='1892-01-03T12:07:01', periods=10, freq='3AS-JUN') ], - ids=['3D', '41987T', '2MS', '3AS_JUN'] + ids=['3D', '41987T', '3Q_AUG', '3AS_JUN'] ) def time_range_kwargs(request): return request.param @@ -40,15 +41,18 @@ def da(index): @pytest.mark.parametrize('freq', [ '700T', '8001T', '12H', '8001H', - '3D', '8D', '8001D', - '2MS', '2M', '3MS', '3M', '4MS', '4M', - '3AS', '3A', '4AS', '4A']) -@pytest.mark.parametrize('closed', [None, 'left', 'right']) -@pytest.mark.parametrize('label', [None, 'left', 'right']) -@pytest.mark.parametrize('base', [17, 24]) + '8D', '8001D', + '2MS', '3MS', + '2QS-AUG', '3QS-SEP', + '3AS-MAR', '4A-MAY']) +@pytest.mark.parametrize('closed', [None, 'right']) +@pytest.mark.parametrize('label', [None, 'right']) +@pytest.mark.parametrize('base', [24, 31]) def test_resampler(freq, closed, label, base, datetime_index, cftime_index): # Fairly extensive testing for standard/proleptic Gregorian calendar + # For any frequencies which are not greater-than-day and anchored + # at the end, the default values for closed and label are 'left'. loffset = '12H' try: da_datetime = da(datetime_index).resample( @@ -67,11 +71,51 @@ def test_resampler(freq, closed, label, base, xr.testing.assert_identical(da_cftime, da_datetime) +@pytest.mark.parametrize('freq', [ + '2M', '3M', + '2Q-JUN', '3Q-JUL', + '3A-FEB', '4A-APR']) +@pytest.mark.parametrize('closed', ['left', None]) +@pytest.mark.parametrize('label', ['left', None]) +@pytest.mark.parametrize('base', [17, 24]) +def test_resampler_end_super_day(freq, closed, label, base, + datetime_index, cftime_index): + # Fairly extensive testing for standard/proleptic Gregorian calendar. + # For greater-than-day frequencies anchored at the end, the default values + # for closed and label are 'right'. + loffset = '12H' + try: + da_datetime = da(datetime_index).resample( + time=freq, closed=closed, label=label, base=base, + loffset=loffset).mean() + except ValueError: + with pytest.raises(ValueError): + da(cftime_index).resample( + time=freq, closed=closed, label=label, base=base, + loffset=loffset).mean() + else: + da_cftime = da(cftime_index).resample(time=freq, closed=closed, + label=label, base=base, + loffset=loffset).mean() + da_cftime['time'] = da_cftime.indexes['time'].to_datetimeindex() + xr.testing.assert_identical(da_cftime, da_datetime) + + +@pytest.mark.parametrize( + ('freq', 'expected'), + [('S', 'left'), ('T', 'left'), ('H', 'left'), ('D', 'left'), + ('M', 'right'), ('MS', 'left'), ('Q', 'right'), ('QS', 'left'), + ('A', 'right'), ('AS', 'left')]) +def test_closed_label_defaults(freq, expected): + assert CFTimeGrouper(freq=freq).closed == expected + assert CFTimeGrouper(freq=freq).label == expected + + @pytest.mark.parametrize('calendar', ['gregorian', 'noleap', 'all_leap', '360_day', 'julian']) def test_calendars(calendar): # Limited testing for non-standard calendars - freq, closed, label, base = '81T', None, None, 17 + freq, closed, label, base = '8001T', None, None, 17 loffset = datetime.timedelta(hours=12) xr_index = xr.cftime_range(start='2004-01-01T12:07:01', periods=7, freq='3D', calendar=calendar)