allow shifting more than 1 year

Signed-off-by: Victor Garcia Reolid <victor@seita.nl>

timely_beliefs/sensors/func_store/knowledge_horizons.py

@@ -32,44 +32,37 @@ def at_date(
 
 
 
-def at_date_annually(
+def x_years_ago_at_date(
     event_start: datetime | pd.DatetimeIndex,
     day: int,
     month : int,
+    x : int = 1,
     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.
+    """Compute the sensor's knowledge horizon to represent the event could be known since some date, `x` years ago.
     
-    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.
+    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 x:               The number of years to shift the reference date to.
     :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)
+    def x_years_ago_at_date_datetime(_event_start : datetime, day : int, month : int) -> timedelta:
+        anchor = dict(year=_event_start.year - x, month=month, day=day, hour=0, minute=0, second=0, microsecond=0)
 
-        if delta_this_year > timedelta(0):
-            return delta_this_year
-        else:
-            return delta_previous_year
+        return _event_start - _event_start.replace(**anchor)
 
     if isinstance(event_start, datetime):
-        return at_date_annually_datetime(event_start, day, month)
+        return x_years_ago_at_date_datetime(event_start, day, month)
     else:
-        return event_start.map(lambda _event_start: at_date_annually_datetime(_event_start,day,month))
+        return event_start.map(lambda _event_start: x_years_ago_at_date_datetime(_event_start,day,month))
 
 def ex_post(
     event_resolution: timedelta,

timely_beliefs/sensors/func_store/test_knowledge_horizons.py

@@ -6,7 +6,7 @@
 
 from timely_beliefs.sensors.func_store.knowledge_horizons import (
     at_date,
-    at_date_annually,
+    x_years_ago_at_date,
     ex_ante,
     ex_post,
     x_days_ago_at_y_oclock,
@@ -55,58 +55,71 @@ def test_fixed_knowledge_time():
     )
 
 
-def test_at_date_annually():
+def test_x_years_ago_at_date():
     """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
+    # Events that occur before the reference
     # year 2024 is leap
-    assert at_date_annually(
+    assert x_years_ago_at_date(
         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(
+    assert x_years_ago_at_date(
         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(
+    assert x_years_ago_at_date(
         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(
+    # Events that occur after the reference
+    assert x_years_ago_at_date(
         event_start=datetime(2021, 11, 21, 3, tzinfo=utc),
         **knowledge_func_params
-    ) == timedelta(days=1, hours=3)
+    ) == timedelta(days=366, hours=3) # 365 + 1
 
-    assert at_date_annually(
+    assert x_years_ago_at_date(
         event_start=datetime(2021, 11, 21, 4, tzinfo=utc),
         **knowledge_func_params
-    ) == timedelta(days=1, hours=4)
+    ) == timedelta(days=366, hours=4) # 365 + 1
+
+    assert x_years_ago_at_date(
+        event_start=datetime(2020, 11, 21, 4, tzinfo=utc),
+        **knowledge_func_params
+    ) == timedelta(days=367, hours=4) # 366 (leap year) + 1
 
     # Test a Daylight Savings Transition
     knowledge_func_params_dst = dict(month=3, day=23) # DST transition 2024
-    assert at_date_annually(
+    assert x_years_ago_at_date(
         event_start=datetime(2024, 3, 25, 0, tzinfo=timezone("Europe/Amsterdam")),
         **knowledge_func_params_dst
-    ) == timedelta(days=2)
+    ) == timedelta(days=368)
 
     # 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(
+        x_years_ago_at_date(
             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)]),
+        pd.TimedeltaIndex([timedelta(days=365, hours=1), timedelta(days=364, hours=2), timedelta(days=364, hours=2), timedelta(days=366, hours=3), timedelta(days=366, hours=4)]),
     )
 
+    knowledge_func_params_2_years = dict(month=11, day=20, x=2)    
+
+    # Check years parameter
+    assert x_years_ago_at_date(
+        event_start=datetime(2024, 11, 19, 1, tzinfo=utc),
+        **knowledge_func_params_2_years
+    ) == timedelta(days=2*365, hours=1) # 366 days - 1
+
 
 def test_dst():
     """Check definition of knowledge horizon for events known x days ago at y o'clock in some timezone z,