tsagentkit 1.0.2__py3-none-any.whl

tsagentkit/models/baselines.py

@@ -0,0 +1,207 @@
+"""Baseline model implementations backed by statsforecast."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pandas as pd
+
+from tsagentkit.contracts import ModelArtifact
+from tsagentkit.utils import normalize_quantile_columns, quantile_col_name
+
+
+def _normalize_name(name: str) -> str:
+    return name.lower().replace("_", "").replace("-", "")
+
+
+def _get_statsforecast_models() -> dict[str, type]:
+    try:
+        from statsforecast import models as sf_models
+    except ImportError as exc:
+        raise ImportError(
+            "statsforecast is required for baseline models. "
+            "Install with: uv sync (or pip install statsforecast)"
+        ) from exc
+
+    model_map: dict[str, type] = {}
+    for key, attr in {
+        "naive": "Naive",
+        "seasonalnaive": "SeasonalNaive",
+        "historicaverage": "HistoricAverage",
+        "theta": "Theta",
+        "windowaverage": "WindowAverage",
+        "movingaverage": "WindowAverage",
+        "seasonalwindowaverage": "SeasonalWindowAverage",
+        "autoets": "AutoETS",
+        "ets": "AutoETS",
+    }.items():
+        model_cls = getattr(sf_models, attr, None)
+        if model_cls is not None:
+            model_map[key] = model_cls
+
+    croston_cls = None
+    for name in ("CrostonClassic", "CrostonOptimized", "CrostonSBA", "CrostonTSB"):
+        croston_cls = getattr(sf_models, name, None)
+        if croston_cls is not None:
+            break
+    if croston_cls is not None:
+        model_map["croston"] = croston_cls
+
+    return model_map
+
+
+def is_baseline_model(model_name: str) -> bool:
+    """Return True if model_name is a supported baseline."""
+    normalized = _normalize_name(model_name)
+    return normalized in _get_statsforecast_models()
+
+
+def _build_model(
+    model_name: str,
+    config: dict[str, Any],
+) -> tuple[Any, str]:
+    model_map = _get_statsforecast_models()
+    normalized = _normalize_name(model_name)
+
+    if normalized not in model_map:
+        raise ValueError(f"Unknown baseline model: {model_name}")
+
+    model_cls = model_map[normalized]
+    model_key = model_cls.__name__
+
+    kwargs: dict[str, Any] = {}
+    if normalized == "seasonalnaive":
+        kwargs["season_length"] = int(config.get("season_length", 1))
+    elif normalized in {"windowaverage", "movingaverage"}:
+        kwargs["window_size"] = int(config.get("window_size", 3))
+    elif normalized == "seasonalwindowaverage":
+        kwargs["season_length"] = int(config.get("season_length", 1))
+        kwargs["window_size"] = int(config.get("window_size", 3))
+    elif normalized in {"autoets", "ets"}:
+        kwargs["season_length"] = int(config.get("season_length", 1))
+    return model_cls(**kwargs), model_key
+
+
+def fit_baseline(
+    model_name: str,
+    dataset: Any,
+    config: dict[str, Any],
+) -> ModelArtifact:
+    """Fit a baseline model using statsforecast."""
+    from statsforecast import StatsForecast
+
+    model, model_key = _build_model(model_name, config)
+
+    try:
+        sf = StatsForecast(
+            models=[model],
+            freq=dataset.task_spec.freq,
+            n_jobs=1,
+        )
+        sf.fit(dataset.df)
+    except TypeError:
+        sf = StatsForecast(
+            df=dataset.df,
+            models=[model],
+            freq=dataset.task_spec.freq,
+            n_jobs=1,
+        )
+        sf.fit()
+
+    return ModelArtifact(
+        model=sf,
+        model_name=model_name,
+        config=config,
+        metadata={
+            "baseline_model": model_key,
+            "freq": dataset.task_spec.freq,
+        },
+    )
+
+
+def _extract_point_column(forecast_df: pd.DataFrame) -> str:
+    value_cols = [c for c in forecast_df.columns if c not in {"unique_id", "ds"}]
+    point_cols = [c for c in value_cols if "lo-" not in c and "hi-" not in c]
+    if not point_cols:
+        raise ValueError("No point forecast column found in statsforecast output.")
+    if len(point_cols) == 1:
+        return point_cols[0]
+    # Prefer plain yhat if present
+    if "yhat" in point_cols:
+        return "yhat"
+    return point_cols[0]
+
+
+def _level_for_quantile(q: float) -> int:
+    return int(round((1 - 2 * abs(q - 0.5)) * 100))
+
+
+def _find_interval_column(
+    forecast_df: pd.DataFrame,
+    level: int,
+    kind: str,
+) -> str | None:
+    suffix = f"{kind}-{level}"
+    matches = [c for c in forecast_df.columns if c.endswith(suffix)]
+    if matches:
+        return matches[0]
+    return None
+
+
+def predict_baseline(
+    model_artifact: ModelArtifact,
+    dataset: Any | None,
+    horizon: int,
+    quantiles: list[float] | None = None,
+) -> pd.DataFrame:
+    """Generate baseline forecasts using statsforecast."""
+    sf = model_artifact.model
+    levels: list[int] = []
+
+    if quantiles:
+        for q in quantiles:
+            if q == 0.5:
+                continue
+            levels.append(_level_for_quantile(q))
+        levels = sorted({lvl for lvl in levels if lvl > 0})
+
+    try:
+        forecast_df = sf.forecast(h=horizon, level=levels or None)
+    except TypeError:
+        if dataset is None:
+            raise
+        forecast_df = sf.forecast(df=dataset.df, h=horizon, level=levels or None)
+    if "unique_id" not in forecast_df.columns or "ds" not in forecast_df.columns:
+        forecast_df = forecast_df.reset_index()
+
+    point_col = _extract_point_column(forecast_df)
+    if point_col != "yhat":
+        forecast_df = forecast_df.rename(columns={point_col: "yhat"})
+
+    # Normalize quantile columns
+    if quantiles:
+        for q in quantiles:
+            col_name = quantile_col_name(q)
+            if q == 0.5:
+                forecast_df[col_name] = forecast_df["yhat"]
+                continue
+            level = _level_for_quantile(q)
+            kind = "lo" if q < 0.5 else "hi"
+            interval_col = _find_interval_column(forecast_df, level, kind)
+            if interval_col is None:
+                forecast_df[col_name] = forecast_df["yhat"]
+            else:
+                forecast_df[col_name] = forecast_df[interval_col]
+
+    forecast_df = normalize_quantile_columns(forecast_df)
+
+    # Keep standard column order
+    if "model" not in forecast_df.columns:
+        forecast_df["model"] = model_artifact.model_name
+    cols = [
+        "unique_id",
+        "ds",
+        "model",
+        "yhat",
+    ] + [c for c in forecast_df.columns if c.startswith("q")]
+    return forecast_df[cols]

tsagentkit/models/sktime.py

@@ -0,0 +1,307 @@
+"""Sktime forecaster adapter with covariate support."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import UTC
+from typing import Any
+
+import pandas as pd
+
+from tsagentkit.contracts import ETaskSpecIncompatible, ForecastResult, ModelArtifact
+from tsagentkit.time import make_future_index
+from tsagentkit.utils import normalize_quantile_columns
+
+
+@dataclass(frozen=True)
+class SktimeModelBundle:
+    """Container for per-series sktime forecasters."""
+
+    model_name: str
+    forecasters: dict[str, Any]
+    exog_columns: list[str]
+    static_columns: list[str]
+    future_columns: list[str]
+
+
+def _require_sktime() -> None:
+    try:
+        import sktime  # noqa: F401
+    except ImportError as exc:
+        raise ImportError(
+            "sktime is required for sktime adapters. Install with: uv sync --extra features"
+        ) from exc
+
+
+def _make_forecaster(model_key: str, season_length: int | None) -> Any:
+    _require_sktime()
+
+    from sktime.forecasting.naive import NaiveForecaster
+
+    if model_key in {"naive", "last"}:
+        return NaiveForecaster(strategy="last")
+    if model_key in {"seasonal_naive", "seasonal"}:
+        sp = season_length or 1
+        return NaiveForecaster(strategy="last", sp=sp)
+
+    raise ValueError(f"Unsupported sktime model key: {model_key}")
+
+
+def _get_exog_columns(
+    covariates: Any | None,
+    plan: Any,
+) -> tuple[list[str], list[str]]:
+    static_cols: list[str] = []
+    future_cols: list[str] = []
+
+    if covariates is None:
+        return static_cols, future_cols
+
+    if getattr(plan, "use_static", True) and covariates.static_x is not None:
+        static_cols = [c for c in covariates.static_x.columns if c != "unique_id"]
+
+    if getattr(plan, "use_future_known", True) and covariates.future_x is not None:
+        future_cols = [
+            c for c in covariates.future_x.columns if c not in {"unique_id", "ds"}
+        ]
+
+    return static_cols, future_cols
+
+
+def _build_train_exog(
+    dataset: Any,
+    covariates: Any,
+    uid: str,
+    ds_index: pd.Index,
+    static_cols: list[str],
+    future_cols: list[str],
+) -> pd.DataFrame | None:
+    if not static_cols and not future_cols:
+        return None
+
+    exog = pd.DataFrame(index=ds_index)
+
+    if future_cols:
+        panel = dataset.panel_with_covariates
+        if panel is None:
+            raise ETaskSpecIncompatible(
+                "Future-known covariates require panel_with_covariates for history.",
+                context={"unique_id": uid},
+            )
+        hist = panel[panel["unique_id"] == uid][["ds"] + future_cols].copy()
+        hist = hist.set_index("ds")
+        exog = exog.join(hist[future_cols], how="left")
+
+    if static_cols:
+        static_row = covariates.static_x
+        if static_row is None:
+            raise ETaskSpecIncompatible(
+                "Static covariates requested but none provided.",
+                context={"unique_id": uid},
+            )
+        static_row = static_row[static_row["unique_id"] == uid]
+        if static_row.empty:
+            raise ETaskSpecIncompatible(
+                "Static covariate row missing for series.",
+                context={"unique_id": uid},
+            )
+        for col in static_cols:
+            value = static_row.iloc[0][col]
+            exog[col] = value
+
+    if exog.isna().any().any():
+        raise ETaskSpecIncompatible(
+            "Missing exogenous covariate values in training data.",
+            context={"unique_id": uid},
+        )
+
+    return exog
+
+
+def _build_future_exog(
+    covariates: Any,
+    uid: str,
+    static_cols: list[str],
+    future_cols: list[str],
+) -> pd.DataFrame | None:
+    if not static_cols and not future_cols:
+        return None
+
+    if covariates.future_index is None:
+        raise ETaskSpecIncompatible(
+            "Future index is required for exogenous forecasting.",
+            context={"unique_id": uid},
+        )
+
+    future_index = covariates.future_index
+    future_index = future_index[future_index["unique_id"] == uid].copy()
+    future_index = future_index.set_index("ds")
+
+    exog = pd.DataFrame(index=future_index.index)
+
+    if future_cols:
+        future = covariates.future_x
+        if future is None:
+            raise ETaskSpecIncompatible(
+                "Future-known covariates requested but none provided.",
+                context={"unique_id": uid},
+            )
+        future = future[future["unique_id"] == uid][["ds"] + future_cols].copy()
+        future = future.set_index("ds")
+        exog = exog.join(future[future_cols], how="left")
+
+    if static_cols:
+        static_row = covariates.static_x
+        if static_row is None:
+            raise ETaskSpecIncompatible(
+                "Static covariates requested but none provided.",
+                context={"unique_id": uid},
+            )
+        static_row = static_row[static_row["unique_id"] == uid]
+        if static_row.empty:
+            raise ETaskSpecIncompatible(
+                "Static covariate row missing for series.",
+                context={"unique_id": uid},
+            )
+        for col in static_cols:
+            value = static_row.iloc[0][col]
+            exog[col] = value
+
+    if exog.isna().any().any():
+        raise ETaskSpecIncompatible(
+            "Missing exogenous covariate values for forecast horizon.",
+            context={"unique_id": uid},
+        )
+
+    return exog
+
+
+def fit_sktime(
+    model_name: str,
+    dataset: Any,
+    plan: Any,
+    covariates: Any | None = None,
+) -> ModelArtifact:
+    """Fit per-series sktime forecasters with exogenous covariates."""
+    model_key = model_name.split("sktime-", 1)[-1]
+    static_cols, future_cols = _get_exog_columns(covariates, plan)
+    forecasters: dict[str, Any] = {}
+
+    for uid in dataset.series_ids:
+        series_df = dataset.get_series(uid).copy()
+        series_df = series_df.sort_values("ds")
+        y = series_df.set_index("ds")["y"].astype(float)
+
+        x_train = None
+        if covariates is not None:
+            x_train = _build_train_exog(
+                dataset=dataset,
+                covariates=covariates,
+                uid=uid,
+                ds_index=y.index,
+                static_cols=static_cols,
+                future_cols=future_cols,
+            )
+
+        forecaster = _make_forecaster(model_key, dataset.task_spec.season_length)
+        forecaster.fit(y, X=x_train)
+        forecasters[uid] = forecaster
+
+    bundle = SktimeModelBundle(
+        model_name=model_name,
+        forecasters=forecasters,
+        exog_columns=sorted(static_cols + future_cols),
+        static_columns=static_cols,
+        future_columns=future_cols,
+    )
+
+    return ModelArtifact(
+        model=bundle,
+        model_name=model_name,
+        config={
+            "model_key": model_key,
+            "static_columns": static_cols,
+            "future_columns": future_cols,
+        },
+    )
+
+
+def predict_sktime(
+    dataset: Any,
+    artifact: ModelArtifact,
+    spec: Any,
+    covariates: Any | None = None,
+) -> ForecastResult:
+    """Predict with per-series sktime forecasters."""
+    _require_sktime()
+
+    from sktime.forecasting.base import ForecastingHorizon
+
+    bundle: SktimeModelBundle = artifact.model
+    rows: list[dict[str, Any]] = []
+
+    if covariates is None:
+        future_index = make_future_index(dataset.df, spec.horizon, spec.freq)
+    else:
+        future_index = covariates.future_index
+        if future_index is None:
+            future_index = make_future_index(dataset.df, spec.horizon, spec.freq)
+
+    for uid, forecaster in bundle.forecasters.items():
+        future_dates = future_index[future_index["unique_id"] == uid]["ds"]
+        future_dates = pd.to_datetime(future_dates).sort_values()
+        fh = ForecastingHorizon(pd.DatetimeIndex(future_dates), is_relative=False)
+
+        x_future = None
+        if covariates is not None and (bundle.static_columns or bundle.future_columns):
+            x_future = _build_future_exog(
+                covariates=covariates,
+                uid=uid,
+                static_cols=bundle.static_columns,
+                future_cols=bundle.future_columns,
+            )
+
+        y_pred = forecaster.predict(fh, X=x_future)
+        if isinstance(y_pred, pd.DataFrame):
+            y_pred = y_pred.iloc[:, 0]
+
+        for ds, value in y_pred.items():
+            rows.append(
+                {
+                    "unique_id": uid,
+                    "ds": pd.Timestamp(ds),
+                    "model": artifact.model_name,
+                    "yhat": float(value),
+                }
+            )
+
+    result_df = pd.DataFrame(rows)
+    result_df = normalize_quantile_columns(result_df)
+
+    provenance = _basic_provenance(dataset, spec, artifact)
+    return ForecastResult(
+        df=result_df,
+        provenance=provenance,
+        model_name=artifact.model_name,
+        horizon=spec.horizon,
+    )
+
+
+def _basic_provenance(dataset: Any, spec: Any, artifact: ModelArtifact) -> Any:
+    from datetime import datetime
+
+    from tsagentkit.contracts import Provenance
+    from tsagentkit.utils import compute_data_signature
+
+    return Provenance(
+        run_id=f"sktime_{datetime.now(UTC).strftime('%Y%m%d_%H%M%S')}",
+        timestamp=datetime.now(UTC).isoformat(),
+        data_signature=compute_data_signature(dataset.df),
+        task_signature=spec.model_hash(),
+        plan_signature=artifact.signature,
+        model_signature=artifact.signature,
+        metadata={"adapter": "sktime"},
+    )
+
+
+__all__ = ["SktimeModelBundle", "fit_sktime", "predict_sktime"]

tsagentkit/monitoring/__init__.py

@@ -0,0 +1,51 @@
+"""Monitoring module for drift detection, coverage, and model stability.
+
+Provides utilities for detecting data drift, monitoring prediction stability,
+checking quantile coverage, and triggering alerts when necessary.
+"""
+
+from __future__ import annotations
+
+from tsagentkit.monitoring.alerts import (
+    Alert,
+    AlertCondition,
+    AlertManager,
+    create_default_coverage_alerts,
+    create_default_drift_alerts,
+)
+from tsagentkit.monitoring.coverage import CoverageCheck, CoverageMonitor
+from tsagentkit.monitoring.drift import DriftDetector
+from tsagentkit.monitoring.report import (
+    CalibrationReport,
+    DriftReport,
+    FeatureDriftResult,
+    StabilityReport,
+    TriggerResult,
+)
+from tsagentkit.monitoring.stability import StabilityMonitor
+from tsagentkit.monitoring.triggers import RetrainTrigger, TriggerEvaluator, TriggerType
+
+__all__ = [
+    # Coverage monitoring
+    "CoverageMonitor",
+    "CoverageCheck",
+    # Alerts
+    "AlertManager",
+    "AlertCondition",
+    "Alert",
+    "create_default_coverage_alerts",
+    "create_default_drift_alerts",
+    # Drift detection
+    "DriftDetector",
+    "DriftReport",
+    "FeatureDriftResult",
+    # Stability monitoring
+    "StabilityMonitor",
+    "StabilityReport",
+    "CalibrationReport",
+    # Retrain triggers
+    "TriggerEvaluator",
+    "RetrainTrigger",
+    "TriggerType",
+    "TriggerResult",
+]