scripts/make_split_diagram.py

"""Figure for Methodology 3.5: the three chronological splitting schemes.

Dates are computed the way the experiments compute them, not typed in:
  A  purged walk-forward (regime ladder, scripts/run_ablation_study.py):
     the longest contiguous run of 31,234 hours ending 16 Oct 2026, cut into
     6 equal blocks; fold k trains on blocks 1..k (last fifth = validation),
     skips 24 hours, and tests block k+1
  B  hold-out (deployed regime model, train_vae_transformer_artifacts.py):
     the same run, first 85% for training, 24 hours purged, last 15% tested
  C  quarterly walk-forward (price range, backtest_improvements.py): each
     quarter Jan 2023 - Sep 2026 tested by a fit on the two years ending
     72 hours before it

Usage:
    python scripts/make_split_diagram.py
"""
from __future__ import annotations

from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.patches import Patch

ROOT = Path(__file__).resolve().parents[1]
SURFACE, INK, INK_2, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#e6e5e1"
TRAIN, VALID, TEST, GAP = "#2a78d6", "#a9cbf0", "#eb6834", "#0b0b0b"

END = pd.Timestamp("2026-10-16 23:00")
ROWS = 31_234
START = END - pd.Timedelta(hours=ROWS - 1)
H = pd.Timedelta(hours=1)


def bar(ax, y, a, b, color, h=0.62, **kw):
    if color == GAP:
        # 24-72 hours is invisible on a six-year axis: draw it as a marker
        # a few weeks wide, centred on the real gap, and say so in the legend.
        mid = a + (b - a) / 2
        a, b = mid - pd.Timedelta(days=9), mid + pd.Timedelta(days=9)
        h = h * 1.25
    ax.barh(y, (b - a) / pd.Timedelta(days=1), left=mdates.date2num(a), height=h,
            color=color, edgecolor=SURFACE, linewidth=0.6, zorder=3 if color == GAP else 2, **kw)


def main() -> None:
    plt.rcParams.update({"font.family": "DejaVu Sans", "text.color": INK})
    fig, axes = plt.subplots(3, 1, figsize=(13, 9.6), facecolor=SURFACE,
                             gridspec_kw={"height_ratios": [5, 1.4, 5.6]})
    fig.subplots_adjust(left=0.13, right=0.98, top=0.86, bottom=0.07, hspace=0.75)
    lo, hi = pd.Timestamp("2020-07-01"), pd.Timestamp("2028-03-01")

    # A: purged walk-forward, 5 folds
    a = axes[0]
    edges = [START + (END - START) * k / 6 for k in range(7)]
    for f in range(5):
        y = 4 - f
        tr_end = edges[f + 1] - 24 * H                    # 24h purge before the test block
        cut = START + (tr_end - START) * 0.8
        bar(a, y, START, cut, TRAIN)
        bar(a, y, cut, tr_end, VALID)
        bar(a, y, tr_end, edges[f + 1], GAP, h=0.62)
        bar(a, y, edges[f + 1], edges[f + 2], TEST)
        a.text(mdates.date2num(edges[f + 2]) + 12, y, f"test {edges[f + 1]:%b %Y} – {edges[f + 2]:%b %Y}",
               va="center", fontsize=8.5, color=INK_2)
    a.set_yticks(range(5))
    a.set_yticklabels([f"fold {5 - i}" for i in range(5)], fontsize=9)
    a.set_title("A  Purged walk-forward, 5 folds — comparing the 14 regime models",
                loc="left", fontsize=11.5, fontweight="bold", pad=20)
    a.text(0, 1.015, "Each model is tested on five different later periods; the training block grows each time. "
                    "Its last fifth (validation) is used for early stopping and the decision cut-off.",
           transform=a.transAxes, fontsize=8.8, color=INK_2, va="bottom")

    # B: hold-out
    b = axes[1]
    cut = START + (END - START) * 0.85
    bar(b, 0, START, cut - 24 * H, TRAIN)
    bar(b, 0, cut - 24 * H, cut, GAP)
    bar(b, 0, cut, END, TEST)
    b.text(mdates.date2num(START) + 20, 0, "train on the first 85%", va="center", fontsize=8.5, color="white")
    b.text(mdates.date2num(END) + 12, 0, f"test the last 15%\n({cut:%b} – {END:%b %Y}, 4,686 h)",
           va="center", fontsize=8.5, color=INK_2)
    b.set_yticks([0])
    b.set_yticklabels(["deployed model"], fontsize=9)
    b.set_ylim(-0.6, 0.6)
    b.set_title("B  Hold-out — the final VAE + Transformer that runs live", loc="left",
                fontsize=11.5, fontweight="bold", pad=20)
    b.text(0, 1.04, "One fit on all earlier data, scored once on the most recent period it never saw.",
           transform=b.transAxes, fontsize=8.8, color=INK_2, va="bottom")

    # C: quarterly walk-forward
    c = axes[2]
    quarters = pd.date_range("2023-01-01", "2026-07-01", freq="QS")
    shown = list(quarters[:3]) + [None] + list(quarters[-2:])
    labels = []
    for i, q in enumerate(shown):
        y = len(shown) - 1 - i
        if q is None:
            c.text(mdates.date2num(pd.Timestamp("2024-09-01")), y, "⋮   one row per quarter, 15 in total   ⋮",
                   va="center", ha="center", fontsize=9, color=INK_2)
            labels.append("")
            continue
        q_end = q + pd.offsets.QuarterBegin(1, startingMonth=1)
        bar(c, y, q - pd.DateOffset(years=2) - 72 * H, q - 72 * H, TRAIN)
        bar(c, y, q - 72 * H, q, GAP)
        bar(c, y, q, q_end, TEST)
        labels.append(f"Q{(q.month - 1) // 3 + 1} {q.year}")
    c.set_yticks(range(len(shown)))
    c.set_yticklabels(labels[::-1], fontsize=9)
    c.set_title("C  Quarterly walk-forward — the price range, exactly as it refits live",
                loc="left", fontsize=11.5, fontweight="bold", pad=20)
    c.text(0, 1.015, "Every quarter from Jan 2023 to Sep 2026 is forecast by a model fitted on the two years "
                    "before it (≈ 32,700 forecasts per horizon).",
           transform=c.transAxes, fontsize=8.8, color=INK_2, va="bottom")

    for ax in axes:
        ax.set_facecolor(SURFACE)
        ax.set_xlim(mdates.date2num(lo), mdates.date2num(hi))
        ax.xaxis.set_major_locator(mdates.YearLocator())
        ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y"))
        ax.grid(axis="x", color=GRID, linewidth=0.8)
        ax.set_axisbelow(True)
        for s in ("top", "right", "left"):
            ax.spines[s].set_visible(False)
        ax.spines["bottom"].set_color(GRID)
        ax.tick_params(colors=INK_2, labelsize=8.5, length=0)

    handles = [Patch(color=TRAIN, label="training (past)"), Patch(color=VALID, label="validation"),
               Patch(color=GAP, label="gap / purge: 24h (A, B) or 72h (C), drawn wider to be visible"), Patch(color=TEST, label="test (later, unseen)")]
    fig.legend(handles=handles, loc="upper left", bbox_to_anchor=(0.13, 0.925), ncol=4, frameon=False, fontsize=9)
    fig.suptitle("Figure 3.2  Three chronological splitting schemes: train on the past, test on the future",
                 x=0.02, ha="left", fontsize=14, fontweight="bold", y=0.975)
    out = ROOT / "outputs" / "figures" / "fig20_data_splitting.png"
    fig.savefig(out, dpi=170, facecolor=SURFACE)
    plt.close(fig)
    print(f"written: {out.relative_to(ROOT)}  (contiguous run {START:%Y-%m-%d} – {END:%Y-%m-%d})")


if __name__ == "__main__":
    main()