SeasonalImputer¶

class SeasonalImputer(BaseActualTransformer):
    """Seasonal decomposition-based imputation for missing values.

    Imputes missing values by leveraging seasonal patterns in the data.
    Missing values are replaced with the mean (or median) of the non-missing
    observations that fall at the same position within each seasonal cycle.
    This is simple modular averaging by season, not a trend/seasonal
    decomposition.

    Parameters
    ----------
    period : int
        Seasonal period (e.g., 7 for weekly, 12 for monthly with annual seasonality).
        Must be >= 2.
    fill_method : {"seasonal_mean", "seasonal_median"}, default="seasonal_mean"
        Method to compute seasonal values:
        - "seasonal_mean": Use mean of same-season observations
        - "seasonal_median": Use median of same-season observations

    Attributes
    ----------
    seasonal_values_ : dict
        Dictionary mapping column names to seasonal value arrays of shape (period,).

    Examples
    --------
    >>> import polars as pl
    >>> from datetime import datetime
    >>> import numpy as np
    >>> from yohou.preprocessing import SeasonalImputer

    >>> # Weekly data with missing values
    >>> X = pl.DataFrame({
    ...     "time": [datetime(2020, 1, i) for i in range(1, 15)],
    ...     "value": [10.0, 20.0, 30.0, 25.0, 15.0, 5.0, 8.0, np.nan, 21.0, 31.0, np.nan, 16.0, 6.0, 9.0],
    ... })
    >>> imputer = SeasonalImputer(period=7)
    >>> imputer.fit(X)
    SeasonalImputer(period=7)
    >>> X_imputed = imputer.transform(X)
    >>> X_imputed["value"].null_count()
    0

    See Also
    --------
    - [`SimpleTimeImputer`][yohou.preprocessing.imputation.SimpleTimeImputer] : Interpolation-based imputation.
    - [`SimpleImputer`][yohou.preprocessing.imputation.SimpleImputer] : Simple constant-strategy imputation.

    """

    _valid_fill_methods = {"seasonal_mean", "seasonal_median"}

    _parameter_constraints: dict = {
        "period": [Interval(numbers.Integral, 2, None, closed="left")],
        "fill_method": [StrOptions(_valid_fill_methods)],
    }

    _tags = {"stateful": False, "invertible": False}

    def __init__(
        self,
        period: int,
        fill_method: str = "seasonal_mean",
    ):
        self.period = period
        self.fill_method = fill_method

    def _season_index_expr(self) -> pl.Expr:
        """Build a polars expression for the time-anchored season index.

        The index counts whole sampling steps between each timestamp and the
        first fitted timestamp, modulo ``period``. This keeps the seasonal
        bucket of any timestamp stable regardless of where a transform slice
        starts.
        """
        if self._step_seconds_ is None:
            # Irregular interval: fall back to row position within the frame.
            return pl.arange(0, pl.len()) % self.period
        offset = (pl.col("time") - self._first_time_).dt.total_seconds() / self._step_seconds_
        return (offset.round().cast(pl.Int64) % self.period).cast(pl.Int64)

    @_fit_context(prefer_skip_nested_validation=True)
    def fit(self, X: pl.DataFrame, y: pl.DataFrame | None = None, **params) -> "SeasonalImputer":
        """Fit the imputer by computing seasonal values.

        Parameters
        ----------
        X : pl.DataFrame
            Input time series with a ``"time"`` column (datetime) and one or
            more numeric columns.
        y : pl.DataFrame or None, default=None
            Ignored.  Present for API compatibility.
        **params : dict
            Metadata to route to nested estimators.

        Returns
        -------
        self
            The fitted transformer instance.

        """
        X = validate_transformer_data(self, X=X, reset=True)
        BaseActualTransformer.fit(self, X, y, **params)

        # Anchor the season index to the first observed timestamp so that
        # transform on any sub-range of the series uses the same seasonal
        # bucket as fit did for the same timestamps.
        self._first_time_ = X["time"][0]
        step = interval_to_timedelta(self.interval_)
        self._step_seconds_ = step.total_seconds() if step is not None else None

        # Compute seasonal values for each column
        data_cols = [c for c in X.columns if c != "time"]
        self.seasonal_values_: dict[str, np.ndarray] = {}

        # Add season index
        X_with_season = X.with_columns(self._season_index_expr().alias("_season_idx"))

        # Aggregate all columns per season in a single grouped pass. Nulls and
        # NaNs are excluded so empty seasons remain null (mapped to NaN below).
        def _stat(col: str) -> pl.Expr:
            """Per-season mean or median of ``col``, ignoring nulls and NaNs."""
            expr = pl.col(col).fill_nan(None).drop_nulls()
            return (expr.mean() if self.fill_method == "seasonal_mean" else expr.median()).alias(col)

        season_stats = X_with_season.group_by("_season_idx").agg([_stat(c) for c in data_cols])

        # Index aggregated rows by season for direct lookup.
        stats_by_season = {row["_season_idx"]: row for row in season_stats.to_dicts()}

        for col_name in data_cols:
            seasonal_vals = np.full(self.period, np.nan)
            for season_idx in range(self.period):
                row = stats_by_season.get(season_idx)
                if row is not None and row[col_name] is not None:
                    seasonal_vals[season_idx] = row[col_name]
            self.seasonal_values_[col_name] = seasonal_vals

        return self

    def _transform(self, X: pl.DataFrame) -> pl.DataFrame:
        """Impute missing values using seasonal patterns.

        Parameters
        ----------
        X : pl.DataFrame
            Validated input time series.

        Returns
        -------
        pl.DataFrame
            Imputed time series.

        """
        # Get data columns
        data_cols = [c for c in X.columns if c != "time"]

        # Time-anchored season index for each row of this slice.
        season_idx_arr = X.select(self._season_index_expr().alias("_season_idx"))["_season_idx"].to_numpy()

        # Impute each column
        result_cols = {"time": X["time"]}

        for col_name in data_cols:
            values = X[col_name].to_numpy().copy()
            seasonal_vals = self.seasonal_values_.get(col_name)

            if seasonal_vals is not None:
                # Find null/nan indices
                null_mask = np.isnan(values) | (X[col_name].is_null().to_numpy())

                # Replace with seasonal values in one vectorised assignment.
                values[null_mask] = seasonal_vals[season_idx_arr[null_mask]]

            # A position whose season bucket was empty at fit time stays NaN;
            # convert back to a proper polars null rather than emitting a float
            # NaN, which would violate the data contract.
            result_cols[col_name] = pl.Series(values).fill_nan(None)

        return pl.DataFrame(result_cols)

    def get_feature_names_out(self, input_features: list[str] | None = None) -> list[str]:
        """Get output feature names for transformation.

        Parameters
        ----------
        input_features : list of str or None, default=None
            Column names of the input features.  If ``None``, uses the
            feature names seen during ``fit``.

        Returns
        -------
        list of str
            Output feature names after transformation.

        """
        check_is_fitted(self, ["feature_names_in_"])
        input_features = _check_feature_names_in(self, input_features)
        return list(input_features)