✨ introduce apply_cuts function

corrai/base/math.py

@@ -20,6 +20,19 @@
 }
 
 
+def apply_cuts(obj: pd.Series | pd.DataFrame, cuts):
+    if cuts is None:
+        return obj
+    if len(obj) == 0:
+        return ValueError("DataFrame is empty, cannot apply cuts.")
+    if obj.index.tz is not None:
+        obj.index = obj.index.tz_localize(None)
+    mask = pd.Series(False, index=obj.index)
+    for start, end in cuts:
+        mask |= (obj.index >= start) & (obj.index <= end)
+    return obj.loc[mask]
+
+
 def aggregate_time_series(
     results: pd.Series,
     indicator: str,
@@ -28,6 +41,7 @@ def aggregate_time_series(
     reference_time_series: pd.Series = None,
     freq: str | pd.Timedelta | dt.timedelta = None,
     prefix: str = "aggregated",
+    cuts: list[tuple[str, str]] | None = None,
 ) -> pd.DataFrame:
     """
     Aggregate time series data using a specified statistical or error metric.
@@ -78,6 +92,10 @@ def aggregate_time_series(
     prefix : str, default="aggregated"
         Prefix to use for naming the output column when `freq` is not specified.
 
+    cuts : list[tuple[str, str]], optional
+        List of (start, end) time intervals (timezone-aware or unaware).
+        If provided, aggregation is performed only on data within these intervals.
+
     Returns
     -------
     pandas.DataFrame
@@ -141,8 +159,11 @@ def aggregate_time_series(
         check_datetime_index(df)
     agg_df = pd.concat([df[indicator].rename(i) for i, df in results.items()], axis=1)
 
+    agg_df = apply_cuts(agg_df, cuts)
+
     if reference_time_series is not None:
         check_datetime_index(reference_time_series)
+        reference_time_series = apply_cuts(reference_time_series, cuts)
         if not agg_df.shape[0] == reference_time_series.shape[0]:
             raise ValueError(
                 "Cannot perform aggregation, Dataframes in results and "

tests/base/test_math.py

@@ -1,10 +1,76 @@
 import pandas as pd
 import numpy as np
 
-from corrai.base.math import aggregate_time_series
+from corrai.base.math import aggregate_time_series, apply_cuts
 
 
 class TestMath:
+    def test_apply_cuts_series(self):
+        cuts = [
+            ("2009-01-01 00:00:00", "2009-01-01 2:00:00"),
+            ("2009-01-01 05:00:00", "2009-01-01 07:00:00"),
+        ]
+
+        index = pd.date_range("2009-01-01", freq="h", periods=8)
+
+        s = pd.Series(range(len(index)), index=index)
+
+        result = apply_cuts(s, cuts)
+
+        mask = pd.Series(False, index=s.index)
+        for start, end in cuts:
+            mask |= (s.index >= start) & (s.index <= end)
+
+        pd.testing.assert_series_equal(
+            result,
+            s.loc[mask],
+        )
+
+    def test_apply_cuts_frame(self):
+        cuts = [
+            ("2009-01-01 00:00:00", "2009-01-01 2:00:00"),
+            ("2009-01-01 05:00:00", "2009-01-01 07:00:00"),
+        ]
+
+        index = pd.date_range("2009-01-01", freq="h", periods=8)
+
+        df = pd.DataFrame(
+            {"a": range(len(index)), "b": range(100, 100 + len(index))},
+            index=index,
+        )
+
+        result = apply_cuts(df, cuts)
+
+        mask = pd.Series(False, index=df.index)
+        for start, end in cuts:
+            mask |= (df.index >= start) & (df.index <= end)
+
+        pd.testing.assert_frame_equal(
+            result,
+            df.loc[mask],
+        )
+
+    def test_apply_cuts_none(self):
+        index = pd.date_range("2009-01-01", freq="h", periods=8)
+        s = pd.Series(range(len(index)), index=index)
+        df = pd.DataFrame(
+            {"a": range(len(index)), "b": range(100, 100 + len(index))},
+            index=index,
+        )
+
+        result_series = apply_cuts(s, None)
+        result_frame = apply_cuts(df, None)
+
+        pd.testing.assert_series_equal(
+            result_series,
+            s,
+        )
+
+        pd.testing.assert_frame_equal(
+            result_frame,
+            df,
+        )
+
     def test_aggregate_time_series(self):
         sim_res = pd.Series(
             [