TimeGPT vs. Fibonacci Median on Wikipedia Web Traffic Forecasting (Kaggle Web Traffic Time Series)

wiki_traffic_forecast.py

"""
TimeGPT vs. Fibonacci Median on Wikipedia Web Traffic Forecasting.

Loads train_2.csv, splits into train/validation, runs both methods on all
145K series, and prints a timing + accuracy comparison.

Usage:
    uv run python wiki_traffic_forecast.py
    uv run python wiki_traffic_forecast.py --model timegpt-2.1

Reference baseline: https://www.kaggle.com/code/safavieh/median-estimation-by-fibonacci-et-al-lb-44-9
"""
import argparse
import os
import time

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from dotenv import load_dotenv
from nixtla import NixtlaClient

load_dotenv()

LANGUAGES      = ["en", "ja", "de", "fr", "zh", "ru", "es"]
GOLDEN_WINDOWS = [6, 12, 18, 30, 48, 78, 126, 203, 329]


def train_val_split(df_wide: pd.DataFrame, h: int):
    date_cols  = [c for c in df_wide.columns if c != "Page"]
    df_train   = df_wide[["Page"] + date_cols[:-h]].copy()
    df_val     = df_wide[["Page"] + date_cols[-h:]].copy()
    return df_train, df_val


def wide_to_long(df_wide: pd.DataFrame) -> pd.DataFrame:
    date_cols = [c for c in df_wide.columns if c != "Page"]
    df_long   = df_wide.melt(id_vars="Page", value_vars=date_cols,
                             var_name="ds", value_name="y")
    df_long   = df_long.rename(columns={"Page": "unique_id"})
    df_long["ds"] = pd.to_datetime(df_long["ds"])
    df_long["y"]  = df_long["y"].fillna(0).clip(lower=0)
    return df_long.sort_values(["unique_id", "ds"]).reset_index(drop=True)


def smape(actual: np.ndarray, forecast: np.ndarray) -> float:
    denom = (np.abs(actual) + np.abs(forecast)) / 2
    mask  = denom > 0
    if not mask.any():
        return np.nan
    return float(np.mean(np.abs(actual[mask] - forecast[mask]) / denom[mask]) * 100)


def golden_median(values: np.ndarray) -> float:
    start = np.nonzero(values)[0]
    if len(start) == 0:
        return 0.0
    active = values[start[0]:]
    n = len(active)
    if n < GOLDEN_WINDOWS[0]:
        return float(np.median(active))
    medians = [np.median(active[-w:]) for w in GOLDEN_WINDOWS if w <= n]
    return float(np.median(medians))


def run_fibonacci_baseline(df_train_wide: pd.DataFrame, df_val_wide: pd.DataFrame):
    train_cols   = [c for c in df_train_wide.columns if c != "Page"]
    val_cols     = [c for c in df_val_wide.columns   if c != "Page"]
    train_matrix = df_train_wide[train_cols].fillna(0).clip(lower=0).values
    val_matrix   = df_val_wide[val_cols].astype(float).fillna(0).clip(lower=0).values
    h            = len(val_cols)

    t0      = time.time()
    results = [smape(val_matrix[i], np.full(h, max(golden_median(train_matrix[i]), 0.0)))
               for i in range(len(train_matrix))]
    return np.array(results), time.time() - t0


def run_timegpt(df_train_wide: pd.DataFrame, df_val_wide: pd.DataFrame,
                client: NixtlaClient, model: str, h: int):
    val_cols   = [c for c in df_val_wide.columns if c != "Page"]
    val_matrix = df_val_wide[val_cols].astype(float).fillna(0).clip(lower=0).values
    val_lookup = dict(zip(df_val_wide["Page"], val_matrix))

    t0      = time.time()
    results = {}
    for lang in LANGUAGES:
        mask    = df_train_wide["Page"].str.contains(f"_{lang}.wikipedia.org_", regex=False)
        df_lang = df_train_wide[mask].reset_index(drop=True)
        if len(df_lang) == 0:
            continue

        print(f"  [{lang}] {len(df_lang):,} series ...", flush=True)
        fcst = client.forecast(
            df=wide_to_long(df_lang), h=h, freq="D",
            time_col="ds", target_col="y", id_col="unique_id",
            model=model,
        )
        fcst["TimeGPT"] = fcst["TimeGPT"].clip(lower=0)

        for page, grp in fcst.groupby("unique_id"):
            actual = val_lookup.get(page)
            if actual is not None:
                forecast = grp.sort_values("ds")["TimeGPT"].values
                results[page] = smape(actual, forecast)

    return np.array(list(results.values())), time.time() - t0


def plot_results(smape_fib: np.ndarray, smape_tgpt: np.ndarray, output: str):
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    axes[0].boxplot(
        [smape_fib[~np.isnan(smape_fib)], smape_tgpt[~np.isnan(smape_tgpt)]],
        tick_labels=["Fibonacci\nMedian", "TimeGPT\n(zero-shot)"],
        patch_artist=True,
        boxprops=dict(facecolor="#d0e8f5"),
        medianprops=dict(color="#e55a2b", linewidth=2),
        showfliers=False,
    )
    axes[0].set_ylabel("SMAPE (%)")
    axes[0].set_title("SMAPE Distribution (lower is better)")

    means  = [np.nanmean(smape_fib), np.nanmean(smape_tgpt)]
    labels = ["Fibonacci\nMedian", "TimeGPT\n(zero-shot)"]
    bars   = axes[1].bar(labels, means, color=["#94c5de", "#e55a2b"], width=0.5)
    for bar, val in zip(bars, means):
        axes[1].text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.3,
                     f"{val:.1f}%", ha="center", va="bottom", fontweight="bold")
    axes[1].set_ylabel("Mean SMAPE (%)")
    axes[1].set_title("Mean SMAPE Comparison")
    axes[1].set_ylim(0, np.nanmax(means) * 1.2)

    plt.suptitle("Wikipedia Web Traffic Forecasting: Fibonacci Median vs. TimeGPT",
                 fontsize=13, fontweight="bold")
    plt.tight_layout()
    plt.savefig(output, dpi=150, bbox_inches="tight")
    plt.close()
    print(f"Plot saved → {output}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--train",   default="data/wiki/train_2.csv")
    parser.add_argument("--horizon", type=int, default=60)
    parser.add_argument("--model",   default="timegpt-1",
                        choices=["timegpt-1", "timegpt-1-long-horizon",
                                 "timegpt-2.1", "timegpt-2-mini"])
    parser.add_argument("--output",  default="wiki_traffic_results.png")
    args = parser.parse_args()

    api_key  = os.environ.get("NIXTLA_API_KEY")
    base_url = os.environ.get("NIXTLA_BASE_URL")
    if not api_key:
        raise SystemExit("NIXTLA_API_KEY not set. Add it to .env or export it.")
    client = (NixtlaClient(api_key=api_key, base_url=base_url)
              if base_url else NixtlaClient(api_key=api_key))

    print("=" * 60)
    print("  Wikipedia Web Traffic: TimeGPT vs. Fibonacci Median")
    print("=" * 60)
    print(f"  Model   : {args.model}")
    print(f"  Horizon : {args.horizon} days")
    print()

    print(f"Loading {args.train} ...")
    t0      = time.time()
    df_wide = pd.read_csv(args.train)
    print(f"  {len(df_wide):,} series loaded in {time.time() - t0:.1f}s")
    n = len(df_wide)

    df_train, df_val = train_val_split(df_wide, args.horizon)
    train_cols = [c for c in df_train.columns if c != "Page"]
    print(f"  Train : {len(train_cols)} days  |  Val : {args.horizon} days")
    print()

    print("Running Fibonacci median baseline ...")
    smape_fib, t_fib = run_fibonacci_baseline(df_train, df_val)
    print(f"  Time  : {t_fib:.1f}s")
    print(f"  SMAPE : {np.nanmean(smape_fib):.2f}%  (median {np.nanmedian(smape_fib):.2f}%)")
    print()

    print(f"Running TimeGPT ({args.model}) ...")
    smape_tgpt, t_tgpt = run_timegpt(df_train, df_val, client, args.model, args.horizon)
    print(f"  Time  : {t_tgpt:.1f}s  ({t_tgpt/60:.1f} min)")
    print(f"  SMAPE : {np.nanmean(smape_tgpt):.2f}%  (median {np.nanmedian(smape_tgpt):.2f}%)")
    print()

    print("=" * 60)
    print(f"  {'Method':<22} {'Series':>7}  {'Time':>10}  {'SMAPE':>8}")
    print(f"  {'-'*22}  {'-'*7}  {'-'*10}  {'-'*8}")
    print(f"  {'Fibonacci median':<22} {n:>7,}  {t_fib:>9.1f}s  {np.nanmean(smape_fib):>7.2f}%")
    print(f"  {'TimeGPT (zero-shot)':<22} {n:>7,}  {t_tgpt/60:>8.1f}m  {np.nanmean(smape_tgpt):>7.2f}%")
    print("=" * 60)
    print()

    plot_results(smape_fib, smape_tgpt, args.output)


if __name__ == "__main__":
    main()