diff --git a/docs/api/index.rst b/docs/api/index.rst index 35ac74d..f3987c3 100644 --- a/docs/api/index.rst +++ b/docs/api/index.rst @@ -7,4 +7,5 @@ API Reference sequence climatology + timesteps calendar diff --git a/docs/api/timesteps.rst b/docs/api/timesteps.rst new file mode 100644 index 0000000..a50cbd7 --- /dev/null +++ b/docs/api/timesteps.rst @@ -0,0 +1,5 @@ +earthkit.time.timesteps - Manipulate time steps and step ranges +=============================================================== + +.. automodule:: earthkit.time.timesteps + :members: diff --git a/docs/guide/api.rst b/docs/guide/api.rst index fff667f..9b7a7a4 100644 --- a/docs/guide/api.rst +++ b/docs/guide/api.rst @@ -163,3 +163,246 @@ To combine yearly dates with multiple reference dates taken from a sequence, use 20180731, 20180803, 20180807, 20180810, 20190731, 20190803, 20190807, 20190810, 20200731, 20200803, 20200807, 20200810 >>> print_dates(model_climate_dates(date(2023, 1, 1), RelativeYear(-7), RelativeYear(-4), 5, 5, seq)) 20151229, 20160102, 20160105, 20161229, 20170102, 20170105, 20171229, 20180102, 20180105, 20181229, 20190102, 20190105 + + +Manipulating time steps and ranges +---------------------------------- + +A step range is represented as a ``start-end`` string, or a ``(start, end)`` +tuple. The :func:`~earthkit.time.timesteps.parse_range` function converts the +string representation to a tuple of integers. + + +Simple manipulations +~~~~~~~~~~~~~~~~~~~~ + +Regularly-spaced step ranges of the same width repeating at a given interval can +be created with :func:`~earthkit.time.timesteps.regular_ranges`: + +.. code-block:: pycon + + >>> from earthkit.time import regular_ranges + >>> list(regular_ranges(24, 48, 12, 6)) + [(24, 36), (30, 42), (36, 48)] + +To expand a step range into its steps, given an interval, use +:func:`~earthkit.time.timesteps.expand_range`: + +.. code-block:: pycon + + >>> from earthkit.time import expand_range + >>> list(expand_range("0-24", 6)) + [0, 6, 12, 18, 24] + + +The :func:`~earthkit.time.timesteps.hours_from_delta` function converts +:class:`~datetime.timedelta` objects to time step numbers in hours: + +.. code-block:: pycon + + >>> from datetime import timedelta + >>> from earthkit.time import hours_from_delta + >>> hours_from_delta(timedelta(days=2)) + 48 + + +Daily dates and ranges +~~~~~~~~~~~~~~~~~~~~~~ + +Daily date computations can be done with +:func:`~earthkit.time.timesteps.start_from_day` and +:func:`~earthkit.time.timesteps.day_from_start`. Both functions take a forecast +base datetime. :func:`~earthkit.time.timesteps.start_from_day` optionally takes +a start time for the day. + +.. code-block:: pycon + + >>> from datetime import datetime, time + >>> from earthkit.time import start_from_day, day_from_start + >>> start_from_day(1, datetime(2026, 9, 5)) + datetime.datetime(2026, 9, 5, 0, 0) + >>> start_from_day(2, datetime(2026, 9, 5), 12) + datetime.datetime(2026, 9, 6, 12, 0) + >>> start_from_day(1, datetime(2026, 9, 5, 12), time(6)) + datetime.datetime(2026, 9, 6, 6, 0) + >>> day_from_start(datetime(2026, 9, 5), datetime(2026, 9, 7)) + 3 + >>> day_from_start(datetime(2026, 9, 5, 12), datetime(2026, 9, 6, 6)) + 1 + +Conversion between day numbers and the corresponding step ranges can be done +with :func:`~earthkit.time.timesteps.range_from_day` and +:func:`~earthkit.time.timesteps.day_from_range`. Both functions take a forecast +base time (which can be a datetime or a time) and the starting time of the first +day. + +.. code-block:: pycon + + >>> from datetime import time + >>> from earthkit.time import range_from_day, day_from_range + >>> range_from_day(1) + (0, 24) + >>> range_from_day(2, 12) + (36, 60) + >>> range_from_day(1, time(6), time(12)) + (6, 30) + >>> day_from_range((24, 48)) + 2 + >>> day_from_range((6, 30), time(6), time(12)) + 1 + +:func:`~earthkit.time.timesteps.day_from_range` also supports providing only +either end of the range: + +.. code-block:: pycon + + >>> from datetime import time + >>> from earthkit.time import day_from_range + >>> day_from_range((None, 48)) + 2 + >>> day_from_range((6, None), time(6), time(12)) + 1 + + +Weekly dates and ranges +~~~~~~~~~~~~~~~~~~~~~~~ + +Weekly date computations can be done with +:func:`~earthkit.time.timesteps.startdate_from_week` and +:func:`~earthkit.time.timesteps.week_from_startdate`. Both functions take a +forecast base date (or datetime). +:func:`~earthkit.time.timesteps.start_from_day` optionally takes a starting day +of the week. The first week starts on the first full day matching the given +starting day (or the first full day if not given), at 00:00. + +.. code-block:: pycon + + >>> from datetime import date, datetime + >>> from earthkit.time.calendar import WEDNESDAY + >>> from earthkit.time import startdate_from_week, week_from_startdate + >>> startdate_from_week(2, date(2022, 1, 11)) + datetime.date(2022, 1, 18) + >>> startdate_from_week(1, datetime(2025, 7, 22, 12), WEDNESDAY) # 2025-07-22 is a Tuesday + datetime.date(2025, 7, 23) + >>> week_from_startdate(date(2022, 1, 11), date(2022, 1, 18)) + 2 + >>> week_from_startdate(datetime(2025, 7, 22, 12), datetime(2025, 7, 23)) + 1 + +Conversion between week numbers and the corresponding step ranges can be done +with :func:`~earthkit.time.timesteps.range_from_week` and +:func:`~earthkit.time.timesteps.week_from_range`. Both functions take a forecast +base time (which can be a date, a datetime, a time, a weeek day, or a (week day, +time) tuple) and a starting day of the week. The first week starts on the first +full day matching the given starting day (or the first full day if not given), +at 00:00. + +.. code-block:: pycon + + >>> from datetime import date, datetime, time + >>> from earthkit.time.calendar import MONDAY, SUNDAY, THURSDAY + >>> from earthkit.time import range_from_week, week_from_range + >>> range_from_week(1) + (0, 168) + >>> range_from_week(2, time(12)) + (180, 348) + >>> range_from_week(1, THURSDAY, MONDAY) + (96, 264) + >>> range_from_week(1, (THURSDAY, time(12)), MONDAY) + (84, 252) + >>> range_from_week(3, date(2023, 11, 10), SUNDAY) # 2023-11-10 is a Friday + (384, 552) + >>> range_from_week(3, datetime(2023, 11, 10, 6), SUNDAY) + (378, 546) + >>> week_from_range((336, 504)) + 3 + >>> week_from_range((96, 264), THURSDAY, MONDAY) + 1 + +:func:`~earthkit.time.timesteps.week_from_range` also supports providing only +either end of the range: + +.. code-block:: pycon + + >>> from earthkit.time.calendar import MONDAY, SUNDAY, THURSDAY + >>> from earthkit.time import week_from_range + >>> week_from_range((None, 504)) + 3 + >>> week_from_range((96, None), THURSDAY, MONDAY) + 1 + + +Monthly dates and ranges +~~~~~~~~~~~~~~~~~~~~~~~~ + +Monthly date computations can be done with +:func:`~earthkit.time.timesteps.startdate_from_month` and +:func:`~earthkit.time.timesteps.month_from_startdate`. Both functions take a +forecast base date (which can be a date, or a ``(year, month)`` tuple). +:func:`~earthkit.time.timesteps.startdate_from_month` optionally takes a start +day for the month (in the reverse case, the start day is inferred from the +date). + +.. code-block:: pycon + + >>> from datetime import date + >>> from earthkit.time import startdate_from_month, month_from_startdate + >>> startdate_from_month(1, (2026, 1)) + datetime.date(2026, 1, 1) + >>> startdate_from_month(3, date(2024, 1, 15)) + datetime.date(2024, 4, 1) + >>> startdate_from_month(5, date(2022, 1, 1), 15) + datetime.date(2022, 5, 15) + >>> startdate_from_month(6, date(2021, 8, 1)) + datetime.date(2022, 1, 1) + >>> month_from_startdate((2026, 1), (2026, 1)) + 1 + >>> month_from_startdate(date(2024, 1, 1), (2024, 3)) + 3 + >>> month_from_startdate(date(2020, 1, 15), date(2020, 8, 1)) + 7 + >>> month_from_startdate(date(2019, 1, 15), date(2019, 8, 15)) + 8 + >>> month_from_startdate(date(2018, 1, 1), date(2018, 9, 15)) + 9 + >>> month_from_startdate((2017, 4), date(2018, 1, 1)) + 10 + +Conversion between month numbers and the corresponding step ranges can be done +with :func:`~earthkit.time.timesteps.range_from_month` and +:func:`~earthkit.time.timesteps.month_from_range`. Both functions take a +forecast base date (which can be a date, or a ``(year, month)`` tuple) and a +starting day of the month. + +.. code-block:: pycon + + >>> from datetime import date + >>> from earthkit.time import range_from_month, month_from_range + >>> range_from_month(1, (2026, 1)) + (0, 744) + >>> range_from_month(2, date(2025, 1, 1)) + (744, 1416) + >>> range_from_month(4, date(2023, 1, 15)) + (2544, 3288) + >>> range_from_month(5, date(2022, 1, 15), 15) + (2880, 3624) + >>> month_from_range((744, 1416), (2025, 1)) + 2 + >>> month_from_range((2544, 3288), date(2023, 1, 15)) + 4 + >>> month_from_range((2880, 3624), date(2022, 1, 15), 15) + 5 + +:func:`~earthkit.time.timesteps.month_from_range` also supports providing only +either end of the range: + +.. code-block:: pycon + + >>> from datetime import date + >>> from earthkit.time import month_from_range + >>> month_from_range((744, None), (2025, 1)) + 2 + >>> month_from_range((None, 3288), date(2023, 1, 15)) + 4 + >>> month_from_range((2880, None), date(2022, 1, 15), 15) + 5 diff --git a/src/earthkit/time/__init__.py b/src/earthkit/time/__init__.py index 4513031..b19d4a8 100644 --- a/src/earthkit/time/__init__.py +++ b/src/earthkit/time/__init__.py @@ -7,6 +7,24 @@ YearlySequence, create_sequence, ) +from .timesteps import ( + day_from_range, + day_from_start, + expand_range, + hours_from_delta, + month_from_range, + month_from_startdate, + parse_range, + range_from_day, + range_from_month, + range_from_week, + regular_ranges, + start_from_day, + startdate_from_month, + startdate_from_week, + week_from_range, + week_from_startdate, +) __version__ = "0.1.8" @@ -21,4 +39,20 @@ "Sequence", "WeeklySequence", "YearlySequence", + "parse_range", + "regular_ranges", + "expand_range", + "hours_from_delta", + "start_from_day", + "day_from_start", + "range_from_day", + "day_from_range", + "startdate_from_week", + "week_from_startdate", + "range_from_week", + "week_from_range", + "startdate_from_month", + "month_from_startdate", + "range_from_month", + "month_from_range", ] diff --git a/src/earthkit/time/timesteps.py b/src/earthkit/time/timesteps.py new file mode 100644 index 0000000..936d892 --- /dev/null +++ b/src/earthkit/time/timesteps.py @@ -0,0 +1,461 @@ +from datetime import date, datetime, time, timedelta +from typing import Iterable, Iterator, Tuple, Union + +from earthkit.time.calendar import MonthInYear, Weekday, month_length + + +def parse_range(arg: str) -> Tuple[int, int]: + """Convert a "start-end" step range to a tuple""" + start, sep, end = arg.partition("-") + if not sep: + end = start + try: + return (int(start), int(end)) + except ValueError: + raise ValueError(f"Invalid step range: {arg!r}") + + +def regular_ranges( + start: int, end: int, width: int, interval: int +) -> Iterator[Tuple[int, int]]: + """Iterate over regularly-spaced step ranges + + This generates ``(a, b)`` pairs where ``a`` is ``start``, ``start + + interval``, ``start + 2 * interval``, ... and ``b = a + width``, such that + the last pair satisfies ``b <= end``. + + Example + ------- + >>> list(regular_ranges(24, 48, 12, 6)) + [(24, 36), (30, 42), (36, 48)] + """ + for step in range(start, end - width + 1, interval): + yield (step, step + width) + + +def expand_range( + steprange: Union[str, Tuple[int, int]], interval: int, include_start: bool = True +) -> Iterable[int]: + """Iterate over regularly-spaced steps within a range + + Examples + -------- + >>> list(expand_range("0-24", 6)) + [0, 6, 12, 18, 24] + >>> list(expand_range((48, 96), 24, include_start=False)) + [72, 96] + """ + if isinstance(steprange, str): + steprange = parse_range(steprange) + start, end = steprange + if not include_start: + start += interval + return range(start, end + 1, interval) + + +def hours_from_delta(delta: timedelta) -> int: + """Convert a :class:`~datetime.timedelta` to a whole number of hours""" + return round(delta.total_seconds() / 3600) + + +def _daily_shift( + base: Union[datetime, time, int] = 0, + daystart: Union[time, int] = 0, +) -> int: + if isinstance(base, datetime): + base = base.time() + if isinstance(base, time): + base = base.hour + if isinstance(daystart, time): + daystart = daystart.hour + return (daystart - base) % 24 + + +def start_from_day( + day: int, base: datetime, daystart: Union[time, int] = 0 +) -> datetime: + """Compute the start of the given day + + Day 1 starts on the first time step with valid time ``daystart``. + + Examples + -------- + >>> from datetime import datetime, time + >>> start_from_day(1, datetime(2026, 9, 5)) + datetime.datetime(2026, 9, 5, 0, 0) + >>> start_from_day(2, datetime(2026, 9, 5), 12) + datetime.datetime(2026, 9, 6, 12, 0) + >>> start_from_day(1, datetime(2026, 9, 5, 12), time(6)) + datetime.datetime(2026, 9, 6, 6, 0) + """ + shift = _daily_shift(base, daystart) + return base + timedelta(days=day - 1, hours=shift) + + +def day_from_start(base: datetime, start: datetime) -> int: + """Compute the forecast day starting at the given datetime + + Examples + -------- + >>> day_from_start(datetime(2026, 9, 5), datetime(2026, 9, 7)) + 3 + >>> day_from_start(datetime(2026, 9, 5, 12), datetime(2026, 9, 6, 6)) + 1 + """ + return (start - base).days + 1 + + +def range_from_day( + day: int, base: Union[datetime, time, int] = 0, daystart: Union[time, int] = 0 +) -> Tuple[int, int]: + """Compute the step range corresponding to a given day + + Day 1 starts on the first time step with valid time ``daystart``. + """ + shift = _daily_shift(base, daystart) + start = shift + 24 * (day - 1) + end = start + 24 + return (start, end) + + +def day_from_range( + steprange: Tuple[Union[int, None], Union[int, None]], + base: Union[datetime, time, int] = 0, + daystart: Union[time, int] = 0, +) -> int: + """Compute the day number corresponding to the given step or step range + + This is the exact inverse of :func:`range_from_day`. + """ + shift = _daily_shift(base, daystart) + start, end = steprange + if start is None: + if end is None: + raise ValueError("At least one end of the range must be provided") + start = end - 24 + elif end is None: + end = start + 24 + + if end - start != 24: + raise ValueError(f"Range '{steprange[0]}-{steprange[1]}' is not one day long") + day, rem = divmod((end - shift), 24) + if rem != 0: + raise ValueError( + f"Range '{steprange[0]}-{steprange[1]}' does not align on a day" + ) + return day + + +_WEEK_IN_HOURS = 7 * 24 + + +def _weekly_shift( + base: Union[date, datetime, time, Weekday, Tuple[Weekday, time], None] = None, + weekstart: Union[Weekday, None] = None, +) -> int: + if base is None: + if weekstart is None: + return 0 + raise ValueError("Cannot compute non-trivial week shift without a base time") + + if isinstance(base, Weekday): + basewd = base + basehour = 0 + elif isinstance(base, datetime): + basewd = base.weekday() + basehour = base.hour + elif isinstance(base, date): + basewd = base.weekday() + basehour = 0 + elif isinstance(base, time): + basewd = None + basehour = base.hour + elif isinstance(base, tuple): + basewd, basetime = base + assert isinstance(basewd, Weekday) + assert isinstance(basetime, time) + basehour = basetime.hour + else: + raise TypeError(f"Invalid type for `base`: {type(base)!r}") + + if weekstart is None: + return (-basehour) % 24 + elif isinstance(weekstart, Weekday): + if basewd is None: + raise ValueError( + "Cannot compute non-trivial week shift without a base week day" + ) + shift = ((weekstart - basewd) % 7) * 24 - basehour + if shift < 0: + shift += _WEEK_IN_HOURS + return shift + else: + raise TypeError(f"Invalid type for `weekstart`: {type(weekstart)!r}") + + +def startdate_from_week( + week: int, base: Union[date, datetime], weekstart: Union[Weekday, None] = None +) -> date: + """Compute the date of the first day of a forecast week + + If no time component is present in ``base``, it is assumed to be 00:00. If + no week start is given, the first week is assumed to start on the first time + step with valid time 00:00. + + This always returns a :class:`~datetime.date`. + + Examples + -------- + >>> from datetime import date, datetime + >>> from earthkit.time.calendar import WEDNESDAY + >>> startdate_from_week(2, date(2022, 1, 11)) + datetime.date(2022, 1, 18) + >>> startdate_from_week(1, datetime(2025, 7, 22, 12), WEDNESDAY) # 2025-07-22 is a Tuesday + datetime.date(2025, 7, 23) + """ + shift = _weekly_shift(base, weekstart) + start = base + timedelta(hours=shift + _WEEK_IN_HOURS * (week - 1)) + if isinstance(start, datetime): + start = start.date() + return start + + +def week_from_startdate(base: Union[date, datetime], start: date) -> int: + """Compute the forecast week starting on the given date + + Examples + -------- + >>> from datetime import date, datetime + >>> week_from_startdate(date(2022, 1, 11), date(2022, 1, 18)) + 2 + >>> week_from_startdate(datetime(2025, 7, 22, 12), datetime(2025, 7, 23)) + 1 + """ + if isinstance(base, datetime): + start = datetime.combine(start, time()) + delta = (start - base).days + else: + delta = (start - base).days + return delta // 7 + 1 + + +def range_from_week( + week: int, + base: Union[date, datetime, time, Weekday, Tuple[Weekday, time], None] = None, + weekstart: Union[Weekday, None] = None, +) -> Tuple[int, int]: + """Compute the step range corresponding to a given week + + If ``base`` is ``None``, ``weekstart`` must be ``None`` as well, + corresponding to the first week starting at step 0. If ``base`` is a + :class:`~datetime.time`, ``weekstart`` must be ``None``. If no time + component is present in ``base``, it is assumed to be 00:00. + + If no week start is given, the first week is assumed to start on the first + time step with valid time 00:00. + + Examples + -------- + >>> from datetime import date, datetime, time + >>> from earthkit.time.calendar import MONDAY, SUNDAY, THURSDAY + >>> range_from_week(1) + (0, 168) + >>> range_from_week(2, time(12)) + (180, 348) + >>> range_from_week(1, THURSDAY, MONDAY) + (96, 264) + >>> range_from_week(1, (THURSDAY, time(12)), MONDAY) + (84, 252) + >>> range_from_week(3, date(2023, 11, 10), SUNDAY) # 2023-11-10 is a Friday + (384, 552) + >>> range_from_week(3, datetime(2023, 11, 10, 6), SUNDAY) + (378, 546) + """ + shift = _weekly_shift(base, weekstart) + start = shift + _WEEK_IN_HOURS * (week - 1) + end = start + _WEEK_IN_HOURS + return (start, end) + + +def week_from_range( + steprange: Tuple[Union[int, None], Union[int, None]], + base: Union[date, datetime, time, Weekday, Tuple[Weekday, time], None] = None, + weekstart: Union[Weekday, None] = None, +) -> int: + """Compute the week number corresponding to the given step or step range + + This is the exact inverse of :func:`range_from_week`. + """ + shift = _weekly_shift(base, weekstart) + start, end = steprange + if start is None: + if end is None: + raise ValueError("At least one end of the range must be provided") + start = end - _WEEK_IN_HOURS + elif end is None: + end = start + _WEEK_IN_HOURS + + if end - start != _WEEK_IN_HOURS: + raise ValueError(f"Range '{steprange[0]}-{steprange[1]}' is not one week long") + week, rem = divmod((end - shift), _WEEK_IN_HOURS) + if rem != 0: + raise ValueError( + f"Range '{steprange[0]}-{steprange[1]}' does not align on a week" + ) + return week + + +def _month_to_date( + arg: Union[date, MonthInYear, Tuple[int, int]], day: int = 1 +) -> date: + if isinstance(arg, date): + return arg + if isinstance(arg, MonthInYear): + return date(arg.year, arg.month, day) + if isinstance(arg, tuple): + year, month = arg + return date(year, month, day) + raise TypeError(f"Cannot convert {type(arg)} to a date") + + +def startdate_from_month( + month: int, + base: Union[date, MonthInYear, Tuple[int, int]], + mstart: int = 1, +) -> date: + """Compute the date of the first day of a forecast month + + Parameters + ---------- + month: int + Forecast month (first month has index 1) + base: :class:`~datetime.date`, :class:`~earthkit.time.calendar.MonthInYear`, ``(year, month)`` tuple + Forecast base time. The first day is assumed to be 1 if not provided. + mstart: int + Starting day for the month + + Examples + -------- + >>> from datetime import date + >>> startdate_from_month(1, (2026, 1)) + datetime.date(2026, 1, 1) + >>> startdate_from_month(3, date(2024, 1, 15)) + datetime.date(2024, 4, 1) + >>> startdate_from_month(5, date(2022, 1, 1), 15) + datetime.date(2022, 5, 15) + >>> startdate_from_month(6, date(2021, 8, 1)) + datetime.date(2022, 1, 1) + """ + base = _month_to_date(base, mstart) + if base.day > mstart: + month += 1 + dyear, vmonth = divmod(base.month + month - 2, 12) + vmonth += 1 + return base.replace(year=base.year + dyear, month=vmonth, day=mstart) + + +def month_from_startdate( + base: Union[date, MonthInYear, Tuple[int, int]], + start: Union[date, MonthInYear, Tuple[int, int]], +) -> int: + """Compute the forecast month starting on the given date + + The first month has number 1 + + Parameters + ---------- + base: :class:`~datetime.date`, :class:`~earthkit.time.calendar.MonthInYear`, ``(year, month)`` tuple + Forecast base time. The first day is assumed to be 1 if not provided. + start: :class:`~datetime.date`, :class:`~earthkit.time.calendar.MonthInYear`, ``(year, month)`` tuple + Month start date. The first day is assumed to be 1 if not provided. + + Examples + -------- + >>> from datetime import date + >>> month_from_startdate((2026, 1), (2026, 1)) + 1 + >>> month_from_startdate(date(2024, 1, 1), (2024, 3)) + 3 + >>> month_from_startdate(date(2020, 1, 15), date(2020, 8, 1)) + 7 + >>> month_from_startdate(date(2019, 1, 15), date(2019, 8, 15)) + 8 + >>> month_from_startdate(date(2018, 1, 1), date(2018, 9, 15)) + 9 + >>> month_from_startdate((2017, 4), date(2018, 1, 1)) + 10 + """ + mstart = start.day if isinstance(start, date) else 1 + base = _month_to_date(base, mstart) + start = _month_to_date(start, mstart) + dyear = start.year - base.year + dmonth = start.month - base.month + if base.day > start.day: + dmonth -= 1 + return dyear * 12 + dmonth + 1 + + +def range_from_month( + month: int, + base: Union[date, MonthInYear, Tuple[int, int]], + mstart: int = 1, +) -> Tuple[int, int]: + """Compute the step range corresponding to a forecast month + + Parameters + ---------- + month: int + Forecast month (first month has index 1) + base: :class:`~datetime.date`, :class:`~earthkit.time.calendar.MonthInYear`, ``(year, month)`` tuple + Forecast base date. The first day is assumed to be 1 if not provided. + mstart: int + Starting day for the month + + Examples + -------- + >>> from datetime import date + >>> range_from_month(1, (2026, 1)) + (0, 744) + >>> range_from_month(2, date(2025, 1, 1)) + (744, 1416) + >>> range_from_month(4, date(2023, 1, 15)) + (2544, 3288) + >>> range_from_month(5, date(2022, 1, 15), 15) + (2880, 3624) + """ + base = _month_to_date(base, mstart) + valid = startdate_from_month(month, base, mstart) + start = hours_from_delta(valid - base) + end = start + month_length(valid.year, valid.month) * 24 + return (start, end) + + +def month_from_range( + steprange: Tuple[Union[int, None], Union[int, None]], + base: Union[date, MonthInYear, Tuple[int, int]], + mstart: int = 1, +) -> int: + """Compute the forecast month corresponding to the given step or step range + + This is the exact inverse of :func:`range_from_month`. + """ + base = _month_to_date(base, mstart) + startstep, endstep = steprange + if startstep is not None: + start = base + timedelta(hours=startstep) + if start.day != mstart: + raise ValueError( + f"Range '{steprange[0]}-{steprange[1]}' does not align on a forecast month" + ) + if endstep is not None: + end = base + timedelta(hours=endstep) + if end.day != start.day: + raise ValueError( + f"Range '{steprange[0]}-{steprange[1]}' is not one month long" + ) + return month_from_startdate(base, start) + else: + if endstep is None: + raise ValueError("At least one end of the range must be provided") + end = base + timedelta(hours=endstep) + return month_from_startdate(base, end) - 1 diff --git a/tests/test_timesteps.py b/tests/test_timesteps.py new file mode 100644 index 0000000..f5edc76 --- /dev/null +++ b/tests/test_timesteps.py @@ -0,0 +1,455 @@ +from contextlib import nullcontext +from datetime import date, datetime, time, timedelta +from typing import List, Optional, Tuple, Union + +import pytest + +from earthkit.time.calendar import ( + FRIDAY, + MONDAY, + SATURDAY, + SUNDAY, + THURSDAY, + TUESDAY, + WEDNESDAY, + MonthInYear, + Weekday, +) +from earthkit.time.timesteps import ( + day_from_range, + day_from_start, + expand_range, + hours_from_delta, + month_from_range, + month_from_startdate, + parse_range, + range_from_day, + range_from_month, + range_from_week, + regular_ranges, + start_from_day, + startdate_from_month, + startdate_from_week, + week_from_range, + week_from_startdate, +) + + +@pytest.mark.parametrize( + "arg, expected", + [ + pytest.param("24-48", (24, 48), id="range"), + pytest.param("12", (12, 12), id="instant"), + pytest.param("abc", None, id="invalid"), + ], +) +def test_parse_range(arg: str, expected: Optional[Tuple[int, int]]): + context = ( + pytest.raises(ValueError, match="^Invalid step range:") + if expected is None + else nullcontext() + ) + with context: + assert parse_range(arg) == expected + + +@pytest.mark.parametrize( + "start, end, width, interval, expected", + [ + pytest.param( + 0, + 120, + 24, + 24, + [(0, 24), (24, 48), (48, 72), (72, 96), (96, 120)], + id="0/120/24/24", + ), + pytest.param(24, 48, 12, 6, [(24, 36), (30, 42), (36, 48)], id="24/48/12/6"), + pytest.param( + 120, 240, 0, 48, [(120, 120), (168, 168), (216, 216)], id="120/240/0/48" + ), + ], +) +def test_regular_ranges( + start: int, end: int, width: int, interval: int, expected: List[Tuple[int, int]] +): + ranges = regular_ranges(start, end, width, interval) + assert list(ranges) == expected + + +@pytest.mark.parametrize( + "steprange, interval, include_start, expected", + [ + pytest.param("0-24", 6, None, [0, 6, 12, 18, 24], id="0-24/6"), + pytest.param("12-36", 12, True, [12, 24, 36], id="12-36/12-start"), + pytest.param("48-96", 24, False, [72, 96], id="48-96/24-nostart"), + pytest.param( + (120, 168), 12, None, [120, 132, 144, 156, 168], id="tup-120-168/12" + ), + pytest.param( + (240, 360), + 24, + True, + [240, 264, 288, 312, 336, 360], + id="tup-240-360/24-start", + ), + pytest.param((0, 48), 12, False, [12, 24, 36, 48], id="tup-0-48/12-nostart"), + ], +) +def test_expand_range( + steprange: Union[str, Tuple[int, int]], + interval: int, + include_start: Optional[bool], + expected: List[int], +): + steps = ( + expand_range(steprange, interval) + if include_start is None + else expand_range(steprange, interval, include_start) + ) + assert list(steps) == expected + + +@pytest.mark.parametrize( + "delta, expected", + [ + pytest.param(timedelta(hours=5), 5, id="5h"), + pytest.param(timedelta(days=2), 48, id="2d"), + pytest.param(timedelta(days=1, hours=12), 36, id="1d12h"), + pytest.param(timedelta(hours=120), 120, id="120h"), + ], +) +def test_hours_from_delta(delta: timedelta, expected: int): + assert hours_from_delta(delta) == expected + + +@pytest.mark.parametrize( + "day, base, daystart, expected", + [ + pytest.param( + 1, datetime(2020, 3, 28, 0), 0, datetime(2020, 3, 28, 0), id="1/0/0" + ), + pytest.param( + 2, + datetime(2026, 7, 13, 12), + time(), + datetime(2026, 7, 15, 0), + id="2/12/time0", + ), + pytest.param( + 3, datetime(2024, 9, 19, 0), 6, datetime(2024, 9, 21, 6), id="3/0/6" + ), + pytest.param( + 4, + datetime(2022, 12, 9, 18), + time(14), + datetime(2022, 12, 13, 14), + id="4/18/time14", + ), + ], +) +def test_start_from_day( + day: int, base: datetime, daystart: Union[time, int], expected: datetime +): + assert start_from_day(day, base, daystart) == expected + + +@pytest.mark.parametrize( + "base, start, expected", + [ + pytest.param(datetime(2020, 3, 28, 0), datetime(2020, 3, 28, 0), 1, id="1/0/0"), + pytest.param( + datetime(2026, 7, 13, 12), datetime(2026, 7, 15, 0), 2, id="2/12/time0" + ), + pytest.param(datetime(2024, 9, 19, 0), datetime(2024, 9, 21, 6), 3, id="3/0/6"), + pytest.param( + datetime(2022, 12, 9, 18), datetime(2022, 12, 13, 14), 4, id="4/18/time14" + ), + ], +) +def test_day_from_start(base: datetime, start: datetime, expected: int): + assert day_from_start(base, start) == expected + + +day_range_params = [ + pytest.param(1, 0, 0, 0, id="1/0/0"), + pytest.param(4, 12, 0, 12, id="4/12/0"), + pytest.param(3, time(6), 12, 6, id="4/time6/12"), + pytest.param(2, datetime(2024, 6, 22, 18), time(12), 18, id="2/datetime18/time12"), +] + + +@pytest.mark.parametrize("day, base, daystart, exp_shift", day_range_params) +def test_range_from_day( + day: int, + base: Union[datetime, time, int], + daystart: Union[time, int], + exp_shift: int, +): + assert range_from_day(day, base, daystart) == ( + exp_shift + 24 * (day - 1), + exp_shift + 24 * day, + ) + + +@pytest.mark.parametrize( + "has_start, has_end", + [ + pytest.param(True, True, id=""), + pytest.param(True, False, id="noend"), + pytest.param(False, True, id="nostart"), + ], +) +@pytest.mark.parametrize("day, base, daystart, shift", day_range_params) +def test_day_from_range( + day: int, + base: Union[datetime, time, int], + daystart: Union[time, int], + shift: int, + has_start: bool, + has_end: bool, +): + start = (day - 1) * 24 + shift + end = day * 24 + shift + steprange = ( + start if has_start else None, + end if has_end else None, + ) + assert day_from_range(steprange, base, daystart) == day + + +def test_day_from_range_invalid(): + with pytest.raises(ValueError, match="Range '.+' is not one day long"): + day_from_range((0, 168)) + + with pytest.raises(ValueError, match="Range '.+' does not align on a day"): + day_from_range((1, 25)) + + with pytest.raises( + ValueError, match="At least one end of the range must be provided" + ): + day_from_range((None, None)) + + +@pytest.mark.parametrize( + "week, base, weekstart, expected", + [ + pytest.param(1, date(2021, 7, 19), None, date(2021, 7, 19), id="1/date/None"), + pytest.param( + 2, datetime(2022, 3, 14, 15), None, date(2022, 3, 22), id="2/datetime/None" + ), + pytest.param(3, date(2023, 5, 17), MONDAY, date(2023, 6, 5), id="3/date/wday"), + pytest.param( + 4, + datetime(2024, 2, 29, 22), + WEDNESDAY, + date(2024, 3, 27), + id="4/datetime/wday", + ), + ], +) +def test_startdate_from_week( + week: int, + base: Union[date, datetime], + weekstart: Union[Weekday, None], + expected: date, +): + assert startdate_from_week(week, base, weekstart) == expected + + +@pytest.mark.parametrize( + "base, start, expected", + [ + pytest.param(date(2021, 7, 19), date(2021, 7, 19), 1, id="1/date"), + pytest.param(datetime(2022, 3, 14, 15), date(2022, 3, 22), 2, id="2/datetime"), + pytest.param(date(2023, 5, 17), date(2023, 6, 5), 3, id="3/date"), + pytest.param(datetime(2024, 2, 29, 22), date(2024, 3, 27), 4, id="4/datetime"), + ], +) +def test_week_from_startdate(base: Union[date, datetime], start: date, expected: int): + assert week_from_startdate(base, start) == expected + + +week_range_params = [ + pytest.param(1, None, None, 0, 0, id="1/None/None"), + pytest.param(5, None, None, 0, 0, id="5/None/None"), + pytest.param(2, WEDNESDAY, None, 0, 0, id="2/wday/None"), + pytest.param(3, date(2022, 7, 19), None, 0, 0, id="2/date/None"), + pytest.param(4, datetime(2023, 3, 14, 15), None, 0, 9, id="4/datetime/None"), + pytest.param(5, time(13), None, 0, 11, id="5/time/None"), + pytest.param(1, (TUESDAY, time(9)), None, 0, 15, id="1/wday-time/None"), + pytest.param(4, date(2025, 5, 17), MONDAY, 2, 0, id="4/date/wday"), + pytest.param(3, THURSDAY, SUNDAY, 3, 0, id="3/wday/wday"), + pytest.param(2, datetime(2024, 2, 29, 22), WEDNESDAY, 5, 2, id="2/datetime/wday"), + pytest.param(4, (SATURDAY, time(20)), FRIDAY, 5, 4, id="4/wday-time/wday"), +] + + +@pytest.mark.parametrize("week, base, weekstart, expd, exph", week_range_params) +def test_range_from_week( + week: int, + base: Union[date, datetime, time, Weekday, Tuple[Weekday, time], None], + weekstart: Union[Weekday, None], + expd: int, + exph: int, +): + assert range_from_week(week, base, weekstart) == ( + ((week - 1) * 7 + expd) * 24 + exph, + (week * 7 + expd) * 24 + exph, + ) + + +@pytest.mark.parametrize( + "has_start, has_end", + [ + pytest.param(True, True, id=""), + pytest.param(True, False, id="noend"), + pytest.param(False, True, id="nostart"), + ], +) +@pytest.mark.parametrize( + "week, base, weekstart, shift_days, shift_hours", week_range_params +) +def test_week_from_range( + week: int, + base: Union[date, datetime, time, Weekday, Tuple[Weekday, time], None], + weekstart: Union[Weekday, None], + shift_days: int, + shift_hours: int, + has_start: bool, + has_end: bool, +): + start = ((week - 1) * 7 + shift_days) * 24 + shift_hours + end = (week * 7 + shift_days) * 24 + shift_hours + steprange = ( + start if has_start else None, + end if has_end else None, + ) + assert week_from_range(steprange, base, weekstart) == week + + +def test_week_from_range_invalid(): + with pytest.raises(ValueError, match="Range '.+' is not one week long"): + week_from_range((0, 24)) + + with pytest.raises(ValueError, match="Range '.+' does not align on a week"): + week_from_range((1, 7 * 24 + 1)) + + with pytest.raises( + ValueError, match="At least one end of the range must be provided" + ): + week_from_range((None, None)) + + +@pytest.mark.parametrize( + "month, base, mstart, expected", + [ + pytest.param(1, (2026, 1), 1, date(2026, 1, 1), id="1/202601-tup"), + pytest.param(2, MonthInYear(2025, 1), 1, date(2025, 2, 1), id="2/202501-month"), + pytest.param(3, date(2024, 1, 15), 1, date(2024, 4, 1), id="3/20240115-date"), + pytest.param( + 4, date(2023, 1, 15), 15, date(2023, 4, 15), id="4/20230115-date/15" + ), + pytest.param( + 5, date(2022, 1, 1), 15, date(2022, 5, 15), id="5/20220101-date/15" + ), + pytest.param(6, date(2021, 8, 1), 1, date(2022, 1, 1), id="6/20210801-date"), + ], +) +def test_startdate_from_month( + month: int, + base: Union[date, MonthInYear, Tuple[int, int]], + mstart: int, + expected: date, +): + assert startdate_from_month(month, base, mstart) == expected + + +@pytest.mark.parametrize( + "base, start, expected", + [ + pytest.param((2026, 1), (2026, 1), 1, id="202601-tup/202601-tup"), + pytest.param(MonthInYear(2025, 1), (2025, 2), 2, id="202501-month/202502-tup"), + pytest.param(date(2024, 1, 1), (2024, 3), 3, id="20240101/202403-tup"), + pytest.param( + date(2023, 1, 1), MonthInYear(2023, 4), 4, id="20230101/202304-month" + ), + pytest.param(date(2022, 1, 1), date(2022, 5, 1), 5, id="20220101/20230401"), + pytest.param( + MonthInYear(2021, 1), + MonthInYear(2021, 6), + 6, + id="202101-month/202104-month", + ), + pytest.param(date(2020, 1, 15), date(2020, 8, 1), 7, id="20200115/20200801"), + pytest.param(date(2019, 1, 15), date(2019, 8, 15), 8, id="20190115/20190815"), + pytest.param(date(2018, 1, 1), date(2018, 9, 15), 9, id="20180101/20180915"), + pytest.param(date(2017, 4, 1), date(2018, 1, 1), 10, id="20170401/20180101"), + ], +) +def test_month_from_startdate( + base: Union[date, MonthInYear, Tuple[int, int]], + start: Union[date, MonthInYear, Tuple[int, int]], + expected: int, +): + assert month_from_startdate(base, start) == expected + + +month_range_params = [ + pytest.param(1, (2026, 1), 1, (0, 744), id="1/202601-tup"), + pytest.param(2, MonthInYear(2025, 1), 1, (744, 1416), id="2/202501-month"), + pytest.param(3, date(2024, 1, 1), 1, (1440, 2184), id="3/20240101-date"), + pytest.param(4, date(2023, 1, 15), 1, (2544, 3288), id="4/20230115-date"), + pytest.param(5, date(2022, 1, 15), 15, (2880, 3624), id="5/20201115-date/15"), +] + + +@pytest.mark.parametrize( + "month, base, mstart, expected", + month_range_params, +) +def test_range_from_month( + month: int, + base: Union[date, MonthInYear, Tuple[int, int]], + mstart: int, + expected: Tuple[int, int], +): + assert range_from_month(month, base, mstart) == expected + + +incomplete_month_range_params = [ + pytest.param(1, date(2021, 1, 1), 1, (0, None), id="1/20210101-date/noend"), + pytest.param(2, date(2020, 1, 15), 1, (1104, None), id="2/20200115-date/noend"), + pytest.param( + 3, date(2019, 1, 1), 15, (None, 2496), id="3/20190101-date/15/nostart" + ), + pytest.param(1, date(2016, 2, 1), 1, (None, 696), id="1/20160201-date/1/nostart"), +] + + +@pytest.mark.parametrize( + "month, base, mstart, steprange", + month_range_params + incomplete_month_range_params, +) +def test_month_from_range( + month: int, + base: Union[date, MonthInYear, Tuple[int, int]], + mstart: int, + steprange: Tuple[int, int], +): + assert month_from_range(steprange, base, mstart) == month + + +def test_month_from_range_invalid(): + with pytest.raises( + ValueError, match="Range '.+' does not align on a forecast month" + ): + month_from_range((3216, 3960), (2022, 1)) + + with pytest.raises(ValueError, match="Range '.+' is not one month long"): + month_from_range((0, 168), (2022, 1)) + + with pytest.raises( + ValueError, match="At least one end of the range must be provided" + ): + month_from_range((None, None), (2022, 1))