diff --git a/corrai/base/math.py b/corrai/base/math.py index db8ae7d..22d201b 100644 --- a/corrai/base/math.py +++ b/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 " diff --git a/tests/base/test_math.py b/tests/base/test_math.py index f37c8de..c739015 100644 --- a/tests/base/test_math.py +++ b/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( [