Skip to content

Commit

Permalink
Add support for sub-hourly frequencies
Browse files Browse the repository at this point in the history
  • Loading branch information
b8raoult committed Aug 27, 2024
1 parent 9998fe9 commit 30b371c
Show file tree
Hide file tree
Showing 2 changed files with 121 additions and 0 deletions.
2 changes: 2 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,8 @@ dynamic = [
"version",
]
dependencies = [
"isodate",

"pyyaml",
"tomli", # Only needed before 3.11
"tqdm",
Expand Down
119 changes: 119 additions & 0 deletions src/anemoi/utils/dates.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@

import calendar
import datetime
import re

import isodate

from .hindcasts import HindcastDatesTimes

Expand Down Expand Up @@ -66,6 +69,122 @@ def as_datetime(date):
raise ValueError(f"Invalid date type: {type(date)}")


def _compress_dates(dates):
dates = sorted(dates)
if len(dates) < 3:
yield dates
return

prev = first = dates.pop(0)
curr = dates.pop(0)
delta = curr - prev
while curr - prev == delta:
prev = curr
if not dates:
break
curr = dates.pop(0)

yield (first, prev, delta)
if dates:
yield from _compress_dates([curr] + dates)


def compress_dates(dates):
dates = [as_datetime(_) for _ in dates]
result = []

for n in _compress_dates(dates):
if isinstance(n, list):
result.extend([str(_) for _ in n])
else:
result.append(" ".join([str(n[0]), "to", str(n[1]), "by", str(n[2])]))

return result


def print_dates(dates):
print(compress_dates(dates))


def frequency_to_string(frequency):
# TODO: use iso8601
frequency = frequency_to_timedelta(frequency)

total_seconds = frequency.total_seconds()
assert int(total_seconds) == total_seconds, total_seconds
total_seconds = int(total_seconds)

seconds = total_seconds

days = seconds // (24 * 3600)
seconds %= 24 * 3600
hours = seconds // 3600
seconds %= 3600
minutes = seconds // 60
seconds %= 60

if days > 0 and hours == 0 and minutes == 0 and seconds == 0:
return f"{days}d"

if days == 0 and hours > 0 and minutes == 0 and seconds == 0:
return f"{hours}h"

if days == 0 and hours == 0 and minutes > 0 and seconds == 0:
return f"{minutes}m"

if days == 0 and hours == 0 and minutes == 0 and seconds > 0:
return f"{seconds}s"

if days > 0:
return f"{total_seconds}s"

return str(frequency)


def frequency_to_timedelta(frequency):
# TODO: use iso8601 or check pytimeparse

if isinstance(frequency, datetime.timedelta):
return frequency

if isinstance(frequency, int):
return datetime.timedelta(hours=frequency)

assert isinstance(frequency, str), (type(frequency), frequency)

try:
return frequency_to_timedelta(int(frequency))
except ValueError:
pass

if re.match(r"^\d+[hdms]$", frequency, re.IGNORECASE):
unit = frequency[-1].lower()
v = int(frequency[:-1])
unit = {"h": "hours", "d": "days", "s": "seconds", "m": "minutes"}[unit]
return datetime.timedelta(**{unit: v})

m = frequency.split(":")
if len(m) == 2:
return datetime.timedelta(hours=int(m[0]), minutes=int(m[1]))

if len(m) == 3:
return datetime.timedelta(hours=int(m[0]), minutes=int(m[1]), seconds=int(m[2]))

# ISO8601
try:
return isodate.parse_duration(frequency)
except isodate.isoerror.ISO8601Error:
pass

raise ValueError(f"Cannot convert frequency {frequency} to timedelta")


def normalize_date(x):
if isinstance(x, str):
return no_time_zone(datetime.datetime.fromisoformat(x))
return x


DOW = {
"monday": 0,
"tuesday": 1,
Expand Down

0 comments on commit 30b371c

Please sign in to comment.