feature(STEF-2717): save last valid rolling aggregate during training (#811)

* feature(STEF-2717): save last valid rolling aggregate during training

* add backwards compatibility for RollingAggregatesAdder and SampleWeighter

* remove backwards compatibility for SampleWeighter (other branch)

* remove unused imports

Author: lschilders

packages/openstef-models/src/openstef_models/transforms/time_domain/rolling_aggregates_adder.py

@@ -2,23 +2,19 @@
 #
 # SPDX-License-Identifier: MPL-2.0
 
-"""Transform for extracting trend features from time series data.
-
-This module provides functionality to compute trend-based features that capture
-long-term patterns and movements in time series data, helping improve forecasting
-accuracy by identifying underlying trends.
-"""
+"""Transform for adding rolling aggregate features to time series data."""
 
 import logging
 from datetime import timedelta
-from typing import Literal, cast, override
+from typing import Any, Literal, override
 
 import pandas as pd
 from pydantic import Field, PrivateAttr
 
 from openstef_core.base_model import BaseConfig
 from openstef_core.datasets import TimeSeriesDataset
 from openstef_core.datasets.validation import validate_required_columns
+from openstef_core.exceptions import NotFittedError
 from openstef_core.transforms import TimeSeriesTransform
 from openstef_core.types import LeadTime
 from openstef_core.utils import timedelta_to_isoformat
@@ -29,11 +25,13 @@
 class RollingAggregatesAdder(BaseConfig, TimeSeriesTransform):
     """Transform that adds rolling aggregate features to time series data.
 
-    This transform computes rolling aggregate statistics (e.g., mean, median, min, max)
+    Computes rolling aggregate statistics (e.g., mean, median, min, max)
     over a specified rolling window and adds these as new features to the dataset.
-    It is useful for capturing recent trends and patterns in the data.
+    It is useful for capturing recent trends and patterns in the data. Handles
+    missing target data during inference via a fallback strategy:
 
-    The rolling aggregates are computed on the specified columns of the dataset.
+    1. Forward-fill from last computed aggregate
+    2. Use last valid aggregate from training
 
     Example:
         >>> import pandas as pd
@@ -55,6 +53,7 @@ class RollingAggregatesAdder(BaseConfig, TimeSeriesTransform):
         ...     aggregation_functions=["mean", "max"],
         ...     horizons=[LeadTime.from_string("PT36H")],
         ... )
+        >>> transform.fit(dataset)
         >>> transformed_dataset = transform.transform(dataset)
         >>> result = transformed_dataset.data[['rolling_mean_load_PT2H', 'rolling_max_load_PT2H']]
         >>> print(result.round(1).head(3))
@@ -65,62 +64,97 @@ class RollingAggregatesAdder(BaseConfig, TimeSeriesTransform):
         2025-01-01 02:00:00                   115.0                  120.0
     """
 
-    feature: str = Field(
-        description="Feature to compute rolling aggregates for.",
-    )
-    horizons: list[LeadTime] = Field(
-        description="List of forecast horizons.",
-        min_length=1,
-    )
+    feature: str = Field(description="Feature to compute rolling aggregates for.")
+    horizons: list[LeadTime] = Field(description="List of forecast horizons.", min_length=1)
     rolling_window_size: timedelta = Field(
         default=timedelta(hours=24),
         description="Rolling window size for the aggregation.",
     )
     aggregation_functions: list[AggregationFunction] = Field(
         default_factory=lambda: ["median", "min", "max"],
-        description="List of aggregation functions to compute over the rolling window. ",
+        description="Aggregation functions to compute over the rolling window.",
     )
 
-    _logger: logging.Logger = PrivateAttr(default=logging.getLogger(__name__))
-
-    def _transform_pandas(self, df: pd.DataFrame) -> pd.DataFrame:
-        rolling_df = cast(
-            pd.DataFrame,
-            df[self.feature].dropna().rolling(window=self.rolling_window_size).agg(self.aggregation_functions),  # pyright: ignore[reportUnknownMemberType, reportCallIssue, reportArgumentType]
-        )
-        # Fill missing values with the last known value
-        rolling_df = rolling_df.reindex(df.index).ffill()
+    _logger: logging.Logger = PrivateAttr(default_factory=lambda: logging.getLogger(__name__))
+    _last_valid_aggregates: dict[str, float] = PrivateAttr(default_factory=dict[str, float])
+    _is_fitted: bool = PrivateAttr(default=False)
 
+    def _make_column_name(self, func: AggregationFunction) -> str:
         suffix = timedelta_to_isoformat(td=self.rolling_window_size)
-        rolling_df = rolling_df.rename(
-            columns={func: f"rolling_{func}_{self.feature}_{suffix}" for func in self.aggregation_functions}
-        )
+        return f"rolling_{func}_{self.feature}_{suffix}"
 
-        return pd.concat([df, rolling_df], axis=1)
+    def _compute_rolling_aggregates(self, series: pd.Series) -> pd.DataFrame:
+        return series.dropna().rolling(window=self.rolling_window_size).agg(self.aggregation_functions)  # type: ignore[return-value]
+
+    @override
+    def fit(self, data: TimeSeriesDataset) -> None:
+        """Compute and store last valid aggregates from training data for fallback."""
+        validate_required_columns(df=data.data, required_columns=[self.feature])
+
+        rolling_df = self._compute_rolling_aggregates(data.data[self.feature])
+
+        for func in self.aggregation_functions:
+            valid_rows = rolling_df[func].dropna()
+            if not valid_rows.empty:
+                self._last_valid_aggregates[self._make_column_name(func)] = float(valid_rows.iloc[-1])
+
+        self._is_fitted = True
 
     @override
     def transform(self, data: TimeSeriesDataset) -> TimeSeriesDataset:
-        if len(self.aggregation_functions) == 0:
-            self._logger.warning(
-                "No aggregation functions specified for RollingAggregatesAdder. Returning original data."
-            )
+        """Add rolling aggregate features, using fallbacks for missing values.
+
+        Returns:
+            Dataset with rolling aggregate feature columns added.
+
+        Raises:
+            NotFittedError: If fit() has not been called.
+        """
+        if not self.aggregation_functions:
+            self._logger.warning("No aggregation functions specified. Returning original data.")
             return data
 
         if len(self.horizons) > 1:
-            self._logger.warning(
-                "Multiple horizons for RollingAggregatesAdder is not yet supported. Returning original data."
-            )
+            self._logger.warning("Multiple horizons not yet supported. Returning original data.")
             return data
 
         validate_required_columns(df=data.data, required_columns=[self.feature])
-        return data.pipe_pandas(self._transform_pandas)
+
+        if not self._is_fitted:
+            raise NotFittedError(self.__class__.__name__)
+
+        # Compute rolling aggregates and apply fallback for missing values
+        result_df = self._compute_and_apply_fallback(data.data)
+        return data.copy_with(result_df)
+
+    def _compute_and_apply_fallback(self, df: pd.DataFrame) -> pd.DataFrame:
+        rolling_df = self._compute_rolling_aggregates(df[self.feature])
+        rolling_df = rolling_df.reindex(df.index).ffill()
+
+        # Rename columns and apply last valid fallback
+        column_mapping = {func: self._make_column_name(func) for func in self.aggregation_functions}
+        rolling_df = rolling_df.rename(columns=column_mapping)
+
+        for col in column_mapping.values():
+            if col in self._last_valid_aggregates:
+                rolling_df[col] = rolling_df[col].fillna(self._last_valid_aggregates[col])  # pyright: ignore[reportUnknownMemberType]
+
+            if rolling_df[col].isna().any():
+                self._logger.warning("Column '%s' has NaN values after fallback.", col)
+
+        return pd.concat([df, rolling_df], axis=1)
 
     @override
     def features_added(self) -> list[str]:
-        return [
-            f"rolling_{func}_{self.feature}_{timedelta_to_isoformat(self.rolling_window_size)}"
-            for func in self.aggregation_functions
-        ]
+        return [self._make_column_name(func) for func in self.aggregation_functions]
+
+    @override
+    def __setstate__(self, state: Any) -> None:  # TODO(#799): delete after stable release
+        if "_last_valid_aggregates" not in state["__pydantic_private__"]:
+            state["__pydantic_private__"]["_last_valid_aggregates"] = {}
+        if "_is_fitted" not in state["__pydantic_private__"]:
+            state["__pydantic_private__"]["_is_fitted"] = True
+        return super().__setstate__(state)
 
 
 __all__ = ["RollingAggregatesAdder"]

packages/openstef-models/tests/unit/transforms/time_domain/test_rolling_aggregates_adder.py

@@ -31,6 +31,7 @@ def test_rolling_aggregate_features_basic():
     )
 
     # Act
+    transform.fit(dataset)
     result = transform.transform(dataset)
 
     # Assert
@@ -73,6 +74,7 @@ def test_rolling_aggregate_features_with_nan():
     )
 
     # Act
+    transform.fit(dataset)
     result = transform.transform(dataset)
 
     # Assert
@@ -99,10 +101,123 @@ def test_rolling_aggregate_features_missing_column_raises_error():
     )
 
     # Act & Assert
+    with pytest.raises(MissingColumnsError, match="Missing required columns"):
+        transform.fit(dataset)
+
     with pytest.raises(MissingColumnsError, match="Missing required columns"):
         transform.transform(dataset)
 
 
+def test_rolling_aggregate_features_empty_feature_on_fit():
+    """Test that transform applies fallback strategy when feature is fully missing during inference."""
+    # Arrange
+    train_data = pd.DataFrame(
+        {"load": [np.nan, np.nan, np.nan]},
+        index=pd.date_range("2023-01-01 00:00:00", periods=3, freq="1h"),
+    )
+    train_dataset = TimeSeriesDataset(train_data, sample_interval=timedelta(hours=1))
+
+    transform = RollingAggregatesAdder(
+        feature="load",
+        rolling_window_size=timedelta(hours=2),
+        aggregation_functions=["mean"],
+        horizons=[LeadTime.from_string("PT36H")],
+    )
+
+    # Act
+    transform.fit(train_dataset)
+    result = transform.transform(train_dataset)
+
+    # Assert
+    assert "rolling_mean_load_PT2H" in result.data.columns
+    assert result.data["rolling_mean_load_PT2H"].isna().all()
+
+
+def test_rolling_aggregate_features_partial_missing_during_inference():
+    """Test that transform computes fresh aggregates when recent data is available."""
+    # Arrange - training data
+    train_data = pd.DataFrame(
+        {"load": [10.0, 20.0, 30.0]},
+        index=pd.date_range("2023-01-01 00:00:00", periods=3, freq="1h"),
+    )
+    train_dataset = TimeSeriesDataset(train_data, sample_interval=timedelta(hours=1))
+
+    # Inference data: some recent values available, then NaN for forecast horizon
+    test_data = pd.DataFrame(
+        {"load": [40.0, 50.0, np.nan, np.nan]},
+        index=pd.date_range("2023-01-01 03:00:00", periods=4, freq="1h"),
+    )
+    test_dataset = TimeSeriesDataset(test_data, sample_interval=timedelta(hours=1))
+
+    transform = RollingAggregatesAdder(
+        feature="load",
+        rolling_window_size=timedelta(hours=2),
+        aggregation_functions=["mean", "max"],
+        horizons=[LeadTime.from_string("PT36H")],
+    )
+
+    # Act
+    transform.fit(train_dataset)
+    result = transform.transform(test_dataset)
+
+    # Assert
+    assert not result.data["rolling_mean_load_PT2H"].isna().any()
+    assert not result.data["rolling_max_load_PT2H"].isna().any()
+
+    # First row: only 40 in window → mean=40, max=40
+    assert result.data["rolling_mean_load_PT2H"].iloc[0] == 40.0
+    assert result.data["rolling_max_load_PT2H"].iloc[0] == 40.0
+
+    # Second row: [40, 50] in window → mean=45, max=50
+    assert result.data["rolling_mean_load_PT2H"].iloc[1] == 45.0
+    assert result.data["rolling_max_load_PT2H"].iloc[1] == 50.0
+
+    # Third and fourth rows: NaN target, forward-fill from last computed
+    assert result.data["rolling_mean_load_PT2H"].iloc[2] == 45.0
+    assert result.data["rolling_max_load_PT2H"].iloc[2] == 50.0
+    assert result.data["rolling_mean_load_PT2H"].iloc[3] == 45.0
+    assert result.data["rolling_max_load_PT2H"].iloc[3] == 50.0
+
+
+def test_rolling_aggregate_fallback_uses_last_valid_from_training():
+    """Test fallback uses last valid aggregate from training when inference data is all NaN."""
+    # Arrange
+    train_data = pd.DataFrame(
+        {"load": [10.0, 20.0, 30.0, 40.0, 50.0]},
+        index=pd.date_range("2023-01-01 00:00:00", periods=5, freq="1h"),
+    )
+    train_dataset = TimeSeriesDataset(train_data, sample_interval=timedelta(hours=1))
+
+    # Inference data with no valid target values
+    test_data = pd.DataFrame(
+        {"load": [np.nan, np.nan, np.nan]},
+        index=pd.date_range("2023-01-01 03:00:00", periods=3, freq="1h"),
+    )
+    test_dataset = TimeSeriesDataset(test_data, sample_interval=timedelta(hours=1))
+
+    transform = RollingAggregatesAdder(
+        feature="load",
+        rolling_window_size=timedelta(hours=2),
+        aggregation_functions=["mean", "max"],
+        horizons=[LeadTime.from_string("PT36H")],
+    )
+
+    # Act
+    transform.fit(train_dataset)
+    result = transform.transform(test_dataset)
+
+    # Assert - all values filled with last valid aggregate from training
+    # Last valid from training: mean of [40, 50] = 45.0, max = 50.0
+    assert "rolling_mean_load_PT2H" in result.data.columns
+    assert "rolling_max_load_PT2H" in result.data.columns
+    assert not result.data["rolling_mean_load_PT2H"].isna().any()
+    assert not result.data["rolling_max_load_PT2H"].isna().any()
+
+    for i in range(3):
+        assert result.data["rolling_mean_load_PT2H"].iloc[i] == 45.0
+        assert result.data["rolling_max_load_PT2H"].iloc[i] == 50.0
+
+
 def test_rolling_aggregate_features_default_parameters():
     """Test transform works with default parameters."""
     # Arrange
@@ -118,6 +233,7 @@ def test_rolling_aggregate_features_default_parameters():
     )
 
     # Act
+    transform.fit(dataset)
     result = transform.transform(dataset)
 
     # Assert - default is 24-hour window with median, min, max