Merge pull request #18 from NERC-CEH/feature/FPM-503-add_mean_sum_agg

Add MeanSum and Sum aggregations

Author: simonstanley

src/time_stream/aggregation.py

@@ -46,6 +46,10 @@ def expr(self, columns: List[str]) -> List[pl.Expr]:
         """Return the Polars expressions for this aggregation."""
         pass
 
+    def post_expr(self, columns: List[str]) -> List[pl.Expr]:
+        """Return additional Polars expressions to be applied after the aggregation."""
+        return []
+
     @classmethod
     def get(cls, aggregation: Union[str, Type["AggregationFunction"], "AggregationFunction"]) -> "AggregationFunction":
         """Factory method to get an aggregation function instance from string names, class types, or existing instances.
@@ -158,6 +162,7 @@ def apply(
         with_columns_expressions = []
         with_columns_expressions.append(self._expected_count_expr(self.ts, aggregation_period))
         with_columns_expressions.extend(self._missing_data_expr(self.ts, columns, missing_criteria))
+        with_columns_expressions.extend(self.post_expr(columns))
 
         # Do the with_column methods
         for with_column in with_columns_expressions:
@@ -313,6 +318,34 @@ def expr(self, columns: List[str]) -> List[pl.Expr]:
         return [pl.col(col).mean().alias(f"mean_{col}") for col in columns]
 
 
+@register_aggregation
+class Sum(AggregationFunction):
+    """An aggregation class to calculate the sum (total) of values within each aggregation period."""
+
+    name = "sum"
+
+    def expr(self, columns: List[str]) -> List[pl.Expr]:
+        """Return the `Polars` expression for calculating the mean in an aggregation period."""
+        return [pl.col(col).sum().alias(f"sum_{col}") for col in columns]
+
+
+@register_aggregation
+class MeanSum(AggregationFunction):
+    """An aggregation class to calculate the mean sum (averaged total) of values within each aggregation period.
+    This will estimate the sum when values are missing according how many values are expected in the period."""
+
+    name = "mean_sum"
+
+    def expr(self, columns: List[str]) -> List[pl.Expr]:
+        """To calculate the mean sum the expression must return the mean, and be multiplied by the expected
+        counts, which is calculated after in the post_expr method."""
+        return [pl.col(col).mean().alias(f"mean_sum_{col}") for col in columns]
+
+    def post_expr(self, columns: List[str]) -> List[pl.Expr]:
+        """Multiply the mean by the expected count to get the mean sum."""
+        return [(pl.col(f"mean_sum_{col}") * pl.col(f"expected_count_{self.ts.time_name}")) for col in columns]
+
+
 @register_aggregation
 class Min(AggregationFunction):
     """An aggregation class to find the minimum of values within each aggregation period."""

tests/time_stream/test_aggregation.py

@@ -9,7 +9,7 @@
 
 from time_stream.period import Period
 from time_stream.base import TimeSeries
-from time_stream.aggregation import AggregationFunction, Max, Mean, Min
+from time_stream.aggregation import AggregationFunction, Max, Mean, Min, MeanSum, Sum
 
 
 def generate_time_series(resolution: Period, periodicity: Period, length: int, missing_data: bool=False) -> TimeSeries:
@@ -129,31 +129,31 @@ def setUp(self):
         self.mock_ts.time_name = "timestamp"
 
     @parameterized.expand([
-        ("mean", Mean), ("min", Min), ("max", Max)
+        ("mean", Mean), ("min", Min), ("max", Max), ("mean_sum", MeanSum), ("sum", Sum)
     ])
     def test_get_with_string(self, get_input, expected):
         """Test AggregationFunction.get() with string input."""
         agg = AggregationFunction.get(get_input)
         self.assertIsInstance(agg, expected)
 
     @parameterized.expand([
-        (Mean, Mean), (Min, Min), (Max, Max)
+        (Mean, Mean), (Min, Min), (Max, Max), (MeanSum, MeanSum), (Sum, Sum)
     ])
     def test_get_with_class(self, get_input, expected):
         """Test AggregationFunction.get() with class input."""
         agg = AggregationFunction.get(get_input)
         self.assertIsInstance(agg, expected)
 
     @parameterized.expand([
-        (Mean(), Mean), (Min(), Min), (Max(), Max)
+        (Mean(), Mean), (Min(), Min), (Max(), Max), (MeanSum(), MeanSum), (Sum(), Sum)
     ])
     def test_get_with_instance(self, get_input, expected):
         """Test AggregationFunction.get() with instance input."""
         agg = AggregationFunction.get(get_input)
         self.assertIsInstance(agg, expected)
 
     @parameterized.expand([
-        "Mean", "MIN", "mAx", "123"
+        "Mean", "MIN", "mAx", "123", "meansum", "sUm"
     ])
     def test_get_with_invalid_string(self, get_input):
         """Test AggregationFunction.get() with invalid string."""
@@ -232,7 +232,14 @@ class TestSimpleAggregations(unittest.TestCase):
 
         ("hourly_to_daily_min", ts_PT1H_2days, Min, P1D, "value", [datetime(2025, 1, 1), datetime(2025, 1, 2)],
          [24, 24], {"value": [0, 24]}, [datetime(2025, 1, 1), datetime(2025, 1, 2)]),
+
+        ("hourly_to_daily_mean_sum", ts_PT1H_2days, MeanSum, P1D, "value", [datetime(2025, 1, 1), datetime(2025, 1, 2)],
+         [24, 24], {"value": [276, 852]}, None),
+        
+        ("hourly_to_daily_sum", ts_PT1H_2days, Sum, P1D, "value", [datetime(2025, 1, 1), datetime(2025, 1, 2)],
+         [24, 24], {"value": [276, 852]}, None),
     ])
+
     def test_microsecond_to_microsecond(
             self, _, input_ts, aggregator, target_period, column, timestamps, counts, values, timestamps_of
     ):
@@ -251,6 +258,12 @@ def test_microsecond_to_microsecond(
 
         ("hourly_to_monthly_min", ts_PT1H_2month, Min, P1M, "value", [datetime(2025, 1, 1), datetime(2025, 2, 1)],
          [744, 672], {"value": [0, 744]}, [datetime(2025, 1, 1), datetime(2025, 2, 1)]),
+        
+        ("hourly_to_monthly_mean_sum", ts_PT1H_2month, MeanSum, P1M, "value", [datetime(2025, 1, 1), datetime(2025, 2, 1)],
+         [744, 672], {"value": [276396, 725424]}, None),
+        
+        ("hourly_to_monthly_sum", ts_PT1H_2month, Sum, P1M, "value", [datetime(2025, 1, 1), datetime(2025, 2, 1)],
+         [744, 672], {"value": [276396, 725424]}, None),
     ])
     def test_microsecond_to_month(
             self, _, input_ts, aggregator, target_period, column, timestamps, counts, values, timestamps_of
@@ -269,6 +282,12 @@ def test_microsecond_to_month(
 
         ("monthly_to_yearly_min", ts_P1M_2years, Min, P1Y, "value", [datetime(2025, 1, 1), datetime(2026, 1, 1)],
          [12, 12], {"value": [0, 12]}, [datetime(2025, 1, 1), datetime(2026, 1, 1)]),
+
+        ("monthly_to_yearly_mean_sum", ts_P1M_2years, MeanSum, P1Y, "value", [datetime(2025, 1, 1), datetime(2026, 1, 1)],
+         [12, 12], {"value": [66, 210]}, None),
+        
+        ("monthly_to_yearly_sum", ts_P1M_2years, Sum, P1Y, "value", [datetime(2025, 1, 1), datetime(2026, 1, 1)],
+         [12, 12], {"value": [66, 210]}, None),
     ])
     def test_month_to_month(
             self, _, input_ts, aggregator, target_period, column, timestamps, counts, values, timestamps_of
@@ -292,6 +311,14 @@ def test_month_to_month(
          [datetime(2025, 1, 1), datetime(2025, 1, 2)], [24, 24],
          {"value": [0, 24], "value_plus1": [1, 25], "value_times2": [0, 48]},
          [datetime(2025, 1, 1), datetime(2025, 1, 2)]),
+        
+        ("multi_column_mean_sum", ts_PT1H_2days, MeanSum, P1D, ["value", "value_plus1", "value_times2"],
+         [datetime(2025, 1, 1), datetime(2025, 1, 2)], [24, 24],
+         {"value": [276, 852], "value_plus1": [300, 876], "value_times2": [552, 1704]}, None),
+        
+        ("multi_column_sum", ts_PT1H_2days, Sum, P1D, ["value", "value_plus1", "value_times2"],
+         [datetime(2025, 1, 1), datetime(2025, 1, 2)], [24, 24],
+         {"value": [276, 852], "value_plus1": [300, 876], "value_times2": [552, 1704]}, None),
     ])
     def test_multi_column(
             self, _, input_ts, aggregator, target_period, column, timestamps, counts, values, timestamps_of
@@ -492,6 +519,27 @@ def test_missing_criteria_available(self, _, valid, criteria):
         assert_frame_equal(result.df, expected_df, check_dtype=False, check_column_order=False, check_exact=False)
 
 
+class TestMeanSumWithMissingData(unittest.TestCase):
+    """Tests the MeanSum aggregation with time series that has missing data."""
+    def setUp(self):
+        self.input_ts = ts_PT1H_2days_missing
+        self.target_period = P1D
+        self.column = "value"
+        self.timestamps = [datetime(2025, 1, 1), datetime(2025, 1, 2)]
+        self.expected_counts = [24, 24]
+        self.actual_counts = [20, 21]
+        self.values = {"value": [280.8, 853.71]}
+
+    def test_mean_sum_with_missing_data(self):
+        """Test MeanSum aggregation with time series that has missing data."""
+        expected_df = generate_expected_df(
+            self.timestamps, MeanSum, self.column, self.values, self.expected_counts,
+            self.actual_counts
+        )
+        result = MeanSum().apply(self.input_ts, self.target_period, self.column)
+        assert_frame_equal(result.df, expected_df, check_dtype=False, check_column_order=False, check_exact=False)
+
+
 class TestPaddedAggregations(unittest.TestCase):
     """Tests that aggregations work as expected with padded time series."""