add at_date_annually

Signed-off-by: Victor Garcia Reolid <[email protected]>

timely_beliefs/sensors/func_store/knowledge_horizons.py

@@ -31,6 +31,46 @@ def at_date(
     return event_start - knowledge_time.astimezone(event_start.tzinfo)
 
 
+
+def at_date_annually(
+    event_start: datetime | pd.DatetimeIndex,
+    day: int,
+    month : int,
+    get_bounds: bool = False,
+) -> timedelta | pd.TimedeltaIndex | tuple[timedelta, timedelta]:
+    """Compute the sensor's knowledge horizon to represent the event could be known since some fixed date on the same year as the event_start.
+    
+    Note: if the event_start happens before the day and month, the knowledge time will be based on the previous year reference.
+
+    For example, it can be used for a tax rate that changes annually and with a known publication date.
+
+    :param event_start:     Start of the event, used as an anchor for determining the knowledge horizon.
+    :param day:             Reference day of the month of the annual date to compare against.
+    :param month:           The month of the annual date to compare against.
+    :param get_bounds:      If True, this function returns bounds on the possible return value.
+                            These bounds are normally useful for creating more efficient database queries when filtering by belief time.
+                            In this case, the knowledge horizon is unbounded.
+    """
+    if get_bounds:
+        return timedelta.min, timedelta.max
+
+    def at_date_annually_datetime(_event_start : datetime, day : int, month : int) -> timedelta:
+        current_year_anchor = dict(year=_event_start.year, month=month, day=day, hour=0, minute=0, second=0, microsecond=0)
+        previous_year_anchor = dict(year=_event_start.year-1, month=month, day=day, hour=0, minute=0, second=0, microsecond=0)
+
+        delta_this_year  =  _event_start - _event_start.replace(**current_year_anchor)
+        delta_previous_year  =  _event_start - _event_start.replace(**previous_year_anchor)
+
+        if delta_this_year > timedelta(0):
+            return delta_this_year
+        else:
+            return delta_previous_year
+
+    if isinstance(event_start, datetime):
+        return at_date_annually_datetime(event_start, day, month)
+    else:
+        return event_start.map(lambda _event_start: at_date_annually_datetime(_event_start,day,month))
+
 def ex_post(
     event_resolution: timedelta,
     ex_post_horizon: timedelta,

timely_beliefs/sensors/func_store/test_knowledge_horizons.py

@@ -2,10 +2,11 @@
 
 import pandas as pd
 from pandas.testing import assert_index_equal
-from pytz import utc
+from pytz import utc, timezone
 
 from timely_beliefs.sensors.func_store.knowledge_horizons import (
     at_date,
+    at_date_annually,
     ex_ante,
     ex_post,
     x_days_ago_at_y_oclock,
@@ -54,6 +55,59 @@ def test_fixed_knowledge_time():
     )
 
 
+def test_at_date_annually():
+    """Check definition of knowledge horizon for events known at a fixed date annually."""
+
+    knowledge_func_params = dict(month=11, day=20)    
+
+    # events that occur before the reference
+    # year 2024 is leap
+    assert at_date_annually(
+        event_start=datetime(2024, 11, 19, 1, tzinfo=utc),
+        **knowledge_func_params
+    ) == timedelta(days=365, hours=1) # 366 days - 1
+
+    # year 2025 is not leap, but 2024 is
+    assert at_date_annually(
+        event_start=datetime(2025, 11, 19, 2, tzinfo=utc),
+        **knowledge_func_params
+    ) == timedelta(days=364, hours=2) # 365 - 1
+
+    # year 2023 is not leap and 2022 either
+    assert at_date_annually(
+        event_start=datetime(2022, 11, 19, 2, tzinfo=utc),
+        **knowledge_func_params
+    ) == timedelta(days=364, hours=2) # 365 - 1
+
+    # events that occur after the reference
+    assert at_date_annually(
+        event_start=datetime(2021, 11, 21, 3, tzinfo=utc),
+        **knowledge_func_params
+    ) == timedelta(days=1, hours=3)
+
+    assert at_date_annually(
+        event_start=datetime(2021, 11, 21, 4, tzinfo=utc),
+        **knowledge_func_params
+    ) == timedelta(days=1, hours=4)
+
+    # Test a Daylight Savings Transition
+    knowledge_func_params_dst = dict(month=3, day=23) # DST transition 2024
+    assert at_date_annually(
+        event_start=datetime(2024, 3, 25, 0, tzinfo=timezone("Europe/Amsterdam")),
+        **knowledge_func_params_dst
+    ) == timedelta(days=2)
+
+    # Repeat test with pd.DatetimeIndex instead
+    event_start = pd.DatetimeIndex(["2024-11-19T01:00:00", "2025-11-19T02:00:00", "2022-11-19T02:00:00", "2021-11-21T03:00:00", "2021-11-21T04:00:00"], tz="utc")
+    assert_index_equal(
+        at_date_annually(
+            event_start=event_start,
+            **knowledge_func_params
+        ),
+        pd.TimedeltaIndex([timedelta(days=365, hours=1), timedelta(days=364, hours=2), timedelta(days=364, hours=2), timedelta(days=1, hours=3), timedelta(days=1, hours=4)]),
+    )
+
+
 def test_dst():
     """Check definition of knowledge horizon for events known x days ago at y o'clock in some timezone z,
     especially around daylight savings time (DST) transitions."""