tsagentkit 1.0.2__py3-none-any.whl

tsagentkit/router/router.py

@@ -0,0 +1,271 @@
+"""Deterministic routing logic aligned to the PRD PlanSpec."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pandas as pd
+
+from tsagentkit.contracts import PlanSpec, RouteDecision, RouterConfig, RouterThresholds, TaskSpec
+from tsagentkit.time import normalize_pandas_freq
+
+if TYPE_CHECKING:
+    from tsagentkit.qa import QAReport
+    from tsagentkit.series import TSDataset
+
+
+def make_plan(
+    dataset: TSDataset,
+    task_spec: TaskSpec,
+    qa: QAReport | None = None,
+    router_config: RouterConfig | None = None,
+    use_tsfm: bool = True,
+    tsfm_preference: list[str] | None = None,
+) -> tuple[PlanSpec, RouteDecision]:
+    """Create a deterministic PlanSpec and RouteDecision for a dataset.
+
+    Returns:
+        Tuple of (PlanSpec, RouteDecision) containing the execution plan
+        and detailed routing decision information.
+    """
+    thresholds = (router_config or RouterConfig()).thresholds
+    stats, buckets = _compute_router_stats(dataset, task_spec, thresholds)
+
+    # Determine TSFM availability
+    available_tsfms: list[str] = []
+    if use_tsfm and _tsfm_allowed(dataset, thresholds):
+        from tsagentkit.models.adapters import AdapterRegistry
+
+        for name in (tsfm_preference or ["chronos", "moirai", "timesfm"]):
+            is_avail, _ = AdapterRegistry.check_availability(name)
+            if is_avail:
+                available_tsfms.append(name)
+
+    # Build candidate model list
+    if "intermittent" in buckets:
+        candidates = ["Croston", "Naive"]
+    elif "short_history" in buckets:
+        candidates = ["HistoricAverage", "Naive"]
+    else:
+        candidates = ["SeasonalNaive", "HistoricAverage", "Naive"]
+
+    if available_tsfms and "intermittent" not in buckets:
+        tsfm_models = [f"tsfm-{name}" for name in available_tsfms]
+        candidates = tsfm_models + candidates
+
+    plan = PlanSpec(
+        plan_name="default",
+        candidate_models=candidates,
+        use_static=True,
+        use_past=True,
+        use_future_known=True,
+        min_train_size=thresholds.min_train_size,
+        max_train_size=thresholds.max_points_per_series_for_tsfm,
+        interval_mode=task_spec.forecast_contract.interval_mode,
+        levels=task_spec.forecast_contract.levels,
+        quantiles=task_spec.forecast_contract.quantiles,
+        allow_drop_covariates=True,
+        allow_baseline=True,
+    )
+
+    # Build RouteDecision for audit trail
+    reasons = [
+        f"selected_models: {candidates}",
+        f"buckets: {buckets}",
+        f"tsfm_available: {bool(available_tsfms)}",
+    ]
+    if available_tsfms:
+        reasons.append(f"tsfm_models: {available_tsfms}")
+
+    route_decision = RouteDecision(
+        stats=stats,
+        buckets=buckets,
+        selected_plan=plan,
+        reasons=reasons,
+    )
+
+    return plan, route_decision
+
+
+def get_model_for_series(
+    unique_id: str,
+    dataset: TSDataset,
+    task_spec: TaskSpec,
+    thresholds: RouterThresholds | None = None,
+) -> str:
+    """Get recommended model for a specific series."""
+    thresholds = thresholds or RouterThresholds()
+    series_df = dataset.get_series(unique_id)
+    stats, buckets = _compute_series_stats(series_df, task_spec, thresholds)
+
+    if "intermittent" in buckets:
+        return "Croston"
+    if "short_history" in buckets:
+        return "HistoricAverage"
+    return "SeasonalNaive"
+
+
+def _compute_router_stats(
+    dataset: TSDataset,
+    task_spec: TaskSpec,
+    thresholds: RouterThresholds,
+) -> tuple[dict[str, float], list[str]]:
+    df = dataset.df
+    stats: dict[str, float] = {}
+    buckets: list[str] = []
+
+    lengths = df.groupby("unique_id").size()
+    min_len = int(lengths.min()) if not lengths.empty else 0
+    stats["min_series_length"] = float(min_len)
+    if min_len < thresholds.min_train_size:
+        buckets.append("short_history")
+
+    missing_ratio = _compute_missing_ratio(df, task_spec)
+    stats["missing_ratio"] = float(missing_ratio)
+    if missing_ratio > thresholds.max_missing_ratio:
+        buckets.append("sparse")
+
+    uid_col = task_spec.panel_contract.unique_id_col
+    ds_col = task_spec.panel_contract.ds_col
+    y_col = task_spec.panel_contract.y_col
+
+    intermittency = _compute_intermittency(df, thresholds, uid_col, ds_col, y_col)
+    stats.update(intermittency)
+    if intermittency.get("intermittent_series_ratio", 0.0) > 0:
+        buckets.append("intermittent")
+
+    season_conf = _seasonality_confidence(df, task_spec, uid_col, y_col)
+    stats["seasonality_confidence"] = float(season_conf)
+    if season_conf >= thresholds.min_seasonality_conf:
+        buckets.append("seasonal_candidate")
+
+    return stats, buckets
+
+
+def _compute_series_stats(
+    series_df: pd.DataFrame,
+    task_spec: TaskSpec,
+    thresholds: RouterThresholds,
+) -> tuple[dict[str, float], list[str]]:
+    stats: dict[str, float] = {}
+    buckets: list[str] = []
+
+    length = len(series_df)
+    stats["series_length"] = float(length)
+    if length < thresholds.min_train_size:
+        buckets.append("short_history")
+
+    missing_ratio = _compute_missing_ratio(series_df, task_spec)
+    stats["missing_ratio"] = float(missing_ratio)
+    if missing_ratio > thresholds.max_missing_ratio:
+        buckets.append("sparse")
+
+    uid_col = task_spec.panel_contract.unique_id_col
+    ds_col = task_spec.panel_contract.ds_col
+    y_col = task_spec.panel_contract.y_col
+
+    intermittency = _compute_intermittency(series_df, thresholds, uid_col, ds_col, y_col)
+    stats.update(intermittency)
+    if intermittency.get("intermittent_series_ratio", 0.0) > 0:
+        buckets.append("intermittent")
+
+    season_conf = _seasonality_confidence(series_df, task_spec, uid_col, y_col)
+    stats["seasonality_confidence"] = float(season_conf)
+    if season_conf >= thresholds.min_seasonality_conf:
+        buckets.append("seasonal_candidate")
+
+    return stats, buckets
+
+
+def _compute_missing_ratio(df: pd.DataFrame, task_spec: TaskSpec) -> float:
+    if df.empty:
+        return 0.0
+    uid_col = task_spec.panel_contract.unique_id_col
+    ds_col = task_spec.panel_contract.ds_col
+
+    ratios = []
+    for uid in df[uid_col].unique():
+        series = df[df[uid_col] == uid].sort_values(ds_col)
+        if series.empty:
+            continue
+        full_range = pd.date_range(
+            start=series[ds_col].min(),
+            end=series[ds_col].max(),
+            freq=normalize_pandas_freq(task_spec.freq),
+        )
+        missing = len(full_range) - len(series)
+        ratio = missing / max(len(full_range), 1)
+        ratios.append(ratio)
+    return float(np.mean(ratios)) if ratios else 0.0
+
+
+def _compute_intermittency(
+    df: pd.DataFrame,
+    thresholds: RouterThresholds,
+    uid_col: str,
+    ds_col: str,
+    y_col: str,
+) -> dict[str, float]:
+    intermittent = 0
+    total = 0
+
+    for uid in df[uid_col].unique():
+        series = df[df[uid_col] == uid].sort_values(ds_col)
+        y = series[y_col].values
+        total += 1
+
+        non_zero_idx = np.where(y > 0)[0]
+        if len(non_zero_idx) <= 1:
+            adi = float("inf")
+            cv2 = float("inf")
+        else:
+            intervals = np.diff(non_zero_idx)
+            adi = float(np.mean(intervals)) if len(intervals) > 0 else float("inf")
+            non_zero_vals = y[non_zero_idx]
+            mean = np.mean(non_zero_vals) if len(non_zero_vals) > 0 else 0.0
+            std = np.std(non_zero_vals) if len(non_zero_vals) > 0 else 0.0
+            cv2 = float((std / mean) ** 2) if mean != 0 else float("inf")
+
+        if adi >= thresholds.max_intermittency_adi and cv2 >= thresholds.max_intermittency_cv2:
+            intermittent += 1
+
+    ratio = intermittent / total if total > 0 else 0.0
+    return {
+        "intermittent_series_ratio": ratio,
+        "intermittent_series_count": float(intermittent),
+    }
+
+
+def _seasonality_confidence(
+    df: pd.DataFrame,
+    task_spec: TaskSpec,
+    uid_col: str,
+    y_col: str,
+) -> float:
+    season_length = task_spec.season_length
+    if season_length is None or season_length <= 1:
+        return 0.0
+    confs: list[float] = []
+    for uid in df[uid_col].unique():
+        series = df[df[uid_col] == uid][y_col].values
+        if len(series) <= season_length:
+            continue
+        series = series - np.mean(series)
+        denom = np.dot(series, series)
+        if denom == 0:
+            continue
+        lagged = np.roll(series, season_length)
+        corr = np.dot(series[season_length:], lagged[season_length:]) / denom
+        confs.append(abs(float(corr)))
+    return float(np.mean(confs)) if confs else 0.0
+
+
+def _tsfm_allowed(dataset: TSDataset, thresholds: RouterThresholds) -> bool:
+    if dataset.n_series > thresholds.max_series_count_for_tsfm:
+        return False
+    max_points = dataset.df.groupby("unique_id").size().max()
+    return max_points <= thresholds.max_points_per_series_for_tsfm
+
+
+__all__ = ["make_plan", "get_model_for_series", "RouteDecision"]

tsagentkit/series/__init__.py

@@ -0,0 +1,26 @@
+"""Series module for tsagentkit.
+
+Provides time series data structures and operations.
+"""
+
+from .alignment import align_timezone, fill_gaps, resample_series
+from .dataset import TSDataset, build_dataset
+from .sparsity import SparsityClass, SparsityProfile, compute_sparsity_profile
+from .validation import normalize_panel_columns, validate_contract
+
+__all__ = [
+    # Dataset
+    "TSDataset",
+    "build_dataset",
+    # Sparsity
+    "SparsityProfile",
+    "SparsityClass",
+    "compute_sparsity_profile",
+    # Alignment
+    "align_timezone",
+    "resample_series",
+    "fill_gaps",
+    # Validation helpers (series layer)
+    "validate_contract",
+    "normalize_panel_columns",
+]

tsagentkit/series/alignment.py

@@ -0,0 +1,206 @@
+"""Time alignment and resampling utilities.
+
+Provides timezone unification and resampling for time series data.
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import pandas as pd
+
+
+def align_timezone(
+    df: pd.DataFrame,
+    target_tz: str | None = "UTC",
+    ds_col: str = "ds",
+) -> pd.DataFrame:
+    """Unify timezones across a dataset.
+
+    Converts all datetime values to the target timezone. Handles
+    timezone-aware and timezone-naive datetimes appropriately.
+
+    Args:
+        df: DataFrame with datetime column
+        target_tz: Target timezone (default: "UTC", None for naive)
+        ds_col: Name of datetime column (default: "ds")
+
+    Returns:
+        DataFrame with unified timezone
+
+    Raises:
+        ValueError: If ds_col is not found or not datetime
+    """
+    if ds_col not in df.columns:
+        raise ValueError(f"Column '{ds_col}' not found in DataFrame")
+
+    if not pd.api.types.is_datetime64_any_dtype(df[ds_col]):
+        raise ValueError(f"Column '{ds_col}' must be datetime type")
+
+    result = df.copy()
+
+    # Handle timezone
+    if target_tz is None:
+        # Make timezone-naive
+        if result[ds_col].dt.tz is not None:
+            result[ds_col] = result[ds_col].dt.tz_localize(None)
+    else:
+        # Convert to target timezone
+        if result[ds_col].dt.tz is None:
+            # Assume UTC for naive datetimes, then convert
+            result[ds_col] = result[ds_col].dt.tz_localize("UTC").dt.tz_convert(target_tz)
+        else:
+            result[ds_col] = result[ds_col].dt.tz_convert(target_tz)
+
+    return result
+
+
+def resample_series(
+    df: pd.DataFrame,
+    freq: str,
+    agg_func: Literal["sum", "mean", "last", "first", "max", "min"] = "sum",
+    ds_col: str = "ds",
+    unique_id_col: str = "unique_id",
+    y_col: str = "y",
+) -> pd.DataFrame:
+    """Resample time series to a new frequency.
+
+    Resamples each series independently to the target frequency using
+    the specified aggregation function.
+
+    Args:
+        df: DataFrame with time series data
+        freq: Target frequency (pandas freq string, e.g., 'D', 'H', 'M')
+        agg_func: Aggregation function (default: "sum")
+        ds_col: Name of datetime column (default: "ds")
+        unique_id_col: Name of series ID column (default: "unique_id")
+        y_col: Name of target column (default: "y")
+
+    Returns:
+        Resampled DataFrame
+
+    Raises:
+        ValueError: If required columns not found
+    """
+    required_cols = {ds_col, unique_id_col, y_col}
+    missing = required_cols - set(df.columns)
+    if missing:
+        raise ValueError(f"Missing required columns: {missing}")
+
+    # Ensure datetime
+    if not pd.api.types.is_datetime64_any_dtype(df[ds_col]):
+        raise ValueError(f"Column '{ds_col}' must be datetime type")
+
+    # Resample each series
+    resampled_frames: list[pd.DataFrame] = []
+
+    for uid in df[unique_id_col].unique():
+        series = df[df[unique_id_col] == uid].set_index(ds_col).sort_index()
+
+        # Select numeric columns for resampling
+        numeric_cols = series.select_dtypes(include=["number"]).columns.tolist()
+
+        if not numeric_cols:
+            continue
+
+        # Resample
+        resampler = series[numeric_cols].resample(freq)
+
+        # Apply aggregation
+        if agg_func == "sum":
+            resampled = resampler.sum()
+        elif agg_func == "mean":
+            resampled = resampler.mean()
+        elif agg_func == "last":
+            resampled = resampler.last()
+        elif agg_func == "first":
+            resampled = resampler.first()
+        elif agg_func == "max":
+            resampled = resampler.max()
+        elif agg_func == "min":
+            resampled = resampler.min()
+        else:
+            raise ValueError(f"Unknown aggregation function: {agg_func}")
+
+        # Add back unique_id
+        resampled[unique_id_col] = uid
+        resampled = resampled.reset_index()
+
+        resampled_frames.append(resampled)
+
+    if not resampled_frames:
+        # Return empty DataFrame with correct structure
+        return pd.DataFrame(columns=[unique_id_col, ds_col, y_col])
+
+    # Combine all series
+    result = pd.concat(resampled_frames, ignore_index=True)
+
+    # Reorder columns to standard order
+    cols = [unique_id_col, ds_col] + [c for c in result.columns if c not in {unique_id_col, ds_col}]
+
+    return result[cols]
+
+
+def fill_gaps(
+    df: pd.DataFrame,
+    freq: str,
+    method: Literal["interpolate", "forward", "backward", "zero"] = "interpolate",
+    ds_col: str = "ds",
+    unique_id_col: str = "unique_id",
+) -> pd.DataFrame:
+    """Fill gaps in time series data.
+
+    Identifies missing timestamps and fills them using the specified method.
+
+    Args:
+        df: DataFrame with time series data
+        freq: Expected frequency (pandas freq string)
+        method: Fill method (default: "interpolate")
+        ds_col: Name of datetime column (default: "ds")
+        unique_id_col: Name of series ID column (default: "unique_id")
+
+    Returns:
+        DataFrame with gaps filled
+    """
+    filled_frames: list[pd.DataFrame] = []
+
+    for uid in df[unique_id_col].unique():
+        series = df[df[unique_id_col] == uid].set_index(ds_col).sort_index()
+
+        # Create complete date range
+        full_range = pd.date_range(start=series.index.min(), end=series.index.max(), freq=freq)
+
+        # Reindex to include gaps
+        series_filled = series.reindex(full_range)
+
+        # Select only numeric columns for filling
+        numeric_cols = series_filled.select_dtypes(include=["number"]).columns.tolist()
+
+        # Fill missing values in numeric columns only
+        if method == "interpolate":
+            series_filled[numeric_cols] = series_filled[numeric_cols].interpolate(method="linear")
+        elif method == "forward":
+            series_filled[numeric_cols] = series_filled[numeric_cols].ffill()
+        elif method == "backward":
+            series_filled[numeric_cols] = series_filled[numeric_cols].bfill()
+        elif method == "zero":
+            series_filled[numeric_cols] = series_filled[numeric_cols].fillna(0)
+
+        # Add back unique_id
+        series_filled[unique_id_col] = uid
+        series_filled = series_filled.reset_index()
+        # Rename the datetime column (could be "index" or the original index name)
+        if "index" in series_filled.columns:
+            series_filled = series_filled.rename(columns={"index": ds_col})
+
+        filled_frames.append(series_filled)
+
+    if not filled_frames:
+        return df.copy()
+
+    result = pd.concat(filled_frames, ignore_index=True)
+
+    # Reorder columns
+    cols = [unique_id_col, ds_col] + [c for c in result.columns if c not in {unique_id_col, ds_col}]
+
+    return result[cols]