Caractéristiques d’actions sans fuite de données, issues de modèles

# 04_model_based_features.py


```py
# ---
# jupyter:
#   jupytext:
#     cell_metadata_filter: tags,-all
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.19.3
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# %% [markdown]
# # US Equities Panel: Model-Based Features
#
# Every feature in [`03_financial_features`](03_financial_features.ipynb) is a function of
# past bars: hand it a row's history and it returns the same value whatever else the panel
# contains. A feature on this page is a function of *parameters estimated from* bars, so the
# estimation window is part of what the feature knows, and a parameter fitted once on the
# whole sample carries the whole sample into every row it touches - including the rows a
# model will later be scored on.
#
# The discipline that removes it is an **estimation schedule**. Each model spends a burn-in
# of history, is fitted on everything before the block it is about to speak for, emits values
# over that block, and is then re-estimated on everything up to the start of the next one.
# A value at a date is therefore a function of that date's own past and of parameters
# estimated strictly earlier, at every date rather than only after some window closed. The
# schedule is what bounds an estimate here; a cross-validation fold selects rows and bounds
# nothing, so this artifact carries one value per stock-date whichever fold later reads it.
#
# Three transforms are built that way, each explained where it is used:
#
# 1. **A regime distance.** Recent months of market-wide return are compared against two
#    reference months learned from the history before them, and each date is given how far it
#    sits from the nearer of them. Section 2.
# 2. **A fractionally differenced price.** A price level is differenced to a fractional order,
#    which keeps part of what the level knows where a return keeps none of it. The weights
#    follow from the order alone, so nothing here is estimated and there is no schedule to
#    put it on. Section 3.
# 3. **A conditional volatility.** Every stock with enough history gets its own volatility
#    model, re-estimated on the schedule and run forward between estimates; a stock too short
#    to pay the burn-in takes a market-level fit. Section 4.
#
# ## Learning objectives
#
# By the end of this notebook you will be able to:
#
# - Tell apart the two date ranges any fitted feature has - the range its parameters were
#   estimated from, and the range it produces values over - and keep the first one entirely
#   before the second
# - Read an estimation schedule off a calendar: what the burn-in costs, how often the
#   parameters are refreshed, and where re-estimation stops so the holdout is never fitted on
# - Chart how a model's fitted parameters move as the schedule advances, and use that to
#   decide how often the model is worth re-estimating
# - Measure what a differencing order costs in the memory it discards and buys in the
#   stationarity it gains, rather than adopting the number a library defaults to
# - Score how well a single column ranks stocks against their later returns, using only test
#   rows, correcting the standard error for the persistence of the series and the test's
#   threshold for the number of columns tried
#
# ## Book reference, prerequisites and artifacts
#
# Chapter 9, Sections 9.1 (Stationarity), 9.3 (Volatility), 9.5 (Regimes). Assumes
# [`02_labels`](02_labels.ipynb) and [`03_financial_features`](03_financial_features.ipynb)
# have been run.
#
# Reads the adjusted daily panel through `load_us_equities()`, `config/setup.yaml` for the
# estimation schedule, the fold design and the holdout boundary, and the primary label file
# written by [`02_labels`](02_labels.ipynb) for the ranking check in Section 7. Writes
# `features/model_based.parquet`, which the model stages join to the stage-03 matrix on
# `(symbol, timestamp)`, alongside a small companion file recording what was written - the
# digest sidecar Section 6 describes.

# %%
"""US Equities Panel: Model-Based Features."""

import gc
import multiprocessing
import os
from concurrent.futures import ProcessPoolExecutor
from datetime import date

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import polars as pl
import yaml
from arch import arch_model
from IPython.display import display
from ml4t.diagnostic.evaluation.stats import benjamini_hochberg_fdr
from ml4t.diagnostic.metrics import compute_ic_hac_stats, cross_sectional_ic_series
from ml4t.diagnostic.splitters.calendar import TradingCalendar
from ml4t.engineer.features.fdiff import ffdiff, get_ffd_weights
from numpy.typing import NDArray
from statsmodels.tsa.stattools import adfuller

from case_studies.utils.artifact_digest import read_digest, value_digest
from case_studies.utils.coverage import assert_sessions_complete
from case_studies.utils.cv_window import modeling_fold_boundaries
from case_studies.utils.temporal import (
    fit_wasserstein_kmeans,
    garch11_conditional_volatility,
    lift_stream,
    refit_boundaries,
    walk_forward_feature,
    wasserstein_distance_1d,
    write_model_based,
)
from data import load_us_equities
from utils.artifact_specs import resolve_label_horizon
from utils.cv_splits import select_folds
from utils.data_quality import top_entities
from utils.paths import display_path, get_case_study_dir
from utils.reproducibility import set_global_seeds
from utils.style import COLORS, FIGSIZE, add_message_title, show_with_alt

CASE_DIR = get_case_study_dir("us_equities_panel")
FEATURES_DIR = CASE_DIR / "features"

# The eligibility screen, carried by 02_labels and 03_financial_features from the same three
# constants on the same columns, so all three stages screen one universe.
MIN_ADV_USD = 1_000_000
MIN_PRICE = 5.0
ADV_WINDOW = 21

# Transform parameters. These define the transforms rather than the strategy, so they are
# declared here; everything that defines the strategy is bound from setup.yaml below.
FFD_D = 0.4  # equity-class default; Section 3 measures what it costs and buys
FFD_D_GRID = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.7, 1.0]
FFD_THRESHOLD = 1e-5

FDR_ALPHA = 0.05

FloatArray = NDArray[np.float64]

# %% [markdown]
# ### The values a run can be given
#
# These are the only ones a caller overrides, so they sit in their own cell where Papermill
# can reach them, and nothing below re-assigns them. What each decides:
#
# - **`START_DATE`** is the first session the price panel is read from. It has to match the
#   date `02_labels` ran from, because Section 1 asserts that the panel this notebook reads
#   digests to the value recorded against the label file Section 7 scores against.
# - **`MAX_FOLDS`** keeps only the *n* earliest cross-validation folds. Zero, the default,
#   keeps all of them. The folds select the rows Section 7 scores and bound no estimate, so a
#   shortened run measures the same features over a shorter span.
# - **`MAX_SYMBOLS`** caps how many stocks are given their own volatility model, taking the
#   ones with the most return history through `top_entities`. Zero, the default, fits every
#   stock that clears the burn-in. It does **not** reduce the panel, and that is deliberate:
#   the regime and market volatility features are computed from the cross-sectional median,
#   and a median over five stocks is not a market. The cost is that a capped run here and a
#   capped `05_evaluation` rank over different frames - complete return histories against the
#   rows that survive eligibility, the fold windows and the label join - so the two sets can
#   differ, and a stock 05 scores but 04 did not fit carries the market-level volatility
#   through the coalesce in Section 5. Nothing here detects that: Section 7's variation table
#   only reports `garch_cond_vol` as market-level when *every* stock left on the date carries
#   the broadcast, and a partial overlap puts fitted and substituted values in one column that
#   still varies across the cross-section. Set the two caps to the same value only after
#   checking that they select the same names, or leave this one at zero.
# - **`XS_MIN_STOCKS`** is the narrowest cross-section a daily return distribution is
#   summarized from. The clustering in Section 2 reads the median of that distribution, and a
#   median over a handful of names is not a market. It belongs here rather than with the
#   transform constants above because it is a property of the panel rather than of the
#   transform: a run over fewer stocks has to lower it or every date is dropped and the
#   clustering has nothing to fit on.
# - **`GARCH_MIN_OBS`** and **`REGIME_MIN_OBS`** override the two burn-ins `setup.yaml`
#   declares - 504 sessions of a stock's own returns before its variance model is fitted, and
#   756 sessions of market history before the first clustering is. Zero, the default, takes
#   the declared value. A run over a shorter history has to lower them or no series clears the
#   burn-in and every block is left empty.
#
# `SEED` fixes the one random step in the notebook, the initialization of the Wasserstein
# clustering in Section 2.

# %% tags=["parameters"]
CASE_STUDY_ID = "us_equities_panel"
START_DATE = "1990-01-01"
MAX_FOLDS = 0
MAX_SYMBOLS = 0
XS_MIN_STOCKS = 50
GARCH_MIN_OBS = 0
REGIME_MIN_OBS = 0
SEED = 42

# %%
set_global_seeds(SEED)

# %% [markdown]
# ## Configuration
#
# The estimation schedule, the fold design, the holdout boundary and the primary label come
# from `config/setup.yaml`. The label's horizon is what binds Section 7: an IC series scored
# on a one-session forward return needs its Newey-West lag set from that horizon, and the
# validation window it may be scored over ends one session before the holdout opens rather
# than on the holdout date.
#
# The horizon is stated in sessions and the buffer in calendar days because that is how the
# splitter takes them. A buffer of one day is the gap the walk-forward design leaves between
# the last training session of a fold and the first session it is scored on, so that the
# outcome of the last training decision is already known when the validation window opens.
#
# The schedule is two numbers per model: how much history it spends before its first estimate,
# and how many sessions an estimate speaks for before the next one replaces it. Both are
# declared rather than searched, and Section 4b measures what the cadence buys.

# %%
SETUP = yaml.safe_load((CASE_DIR / "config" / "setup.yaml").read_text())

PRIMARY_LABEL = SETUP["labels"]["primary"]
LABEL_HORIZON = int(resolve_label_horizon(CASE_STUDY_ID, PRIMARY_LABEL, SETUP).rstrip("Dd"))
LABEL_BUFFER = SETUP["labels"]["buffer"]
HOLDOUT_START = str(SETUP["evaluation"]["holdout_start"])
END_DATE = str(SETUP["evaluation"]["holdout_end"])
CALENDAR = SETUP["evaluation"]["calendar"]

_REGIME = SETUP["model_based"]["regime"]
_GARCH = SETUP["model_based"]["garch"]
N_CLUSTERS = int(_REGIME["n_clusters"])
WASSERSTEIN_WINDOW = int(_REGIME["window"])
WASSERSTEIN_OVERLAP = int(_REGIME["overlap"])
REGIME_BURNIN = REGIME_MIN_OBS or int(_REGIME["burnin"])
REGIME_REFIT_EVERY = int(_REGIME["refit_every"])
GARCH_BURNIN = GARCH_MIN_OBS or int(_GARCH["burnin"])
GARCH_REFIT_EVERY = int(_GARCH["refit_every"])

print(
    f"Regime model: {N_CLUSTERS} states clustered from {WASSERSTEIN_WINDOW}-session windows "
    f"overlapping by {WASSERSTEIN_OVERLAP}, after a {REGIME_BURNIN}-session burn-in and "
    f"re-estimated every {REGIME_REFIT_EVERY} sessions."
)
print(
    f"Volatility model: one GARCH(1,1) per stock after a {GARCH_BURNIN}-session burn-in, "
    f"re-estimated every {GARCH_REFIT_EVERY} sessions."
)
print(
    f"Section 7 scores against {PRIMARY_LABEL}, the return over the next {LABEL_HORIZON} "
    f"session(s), so its Newey-West lag is set from {LABEL_HORIZON}."
)
print(
    f"The walk-forward design leaves {LABEL_BUFFER} between a fold's last training session "
    "and the first session it is scored on."
)
print(
    f"Everything from {HOLDOUT_START} to {END_DATE} is held out: no parameter here is "
    "estimated from it, and no number here is measured on it."
)
print(
    f"A stock is eligible on a date when its printed close is above ${MIN_PRICE:.0f} and its "
    f"dollar volume has averaged above ${MIN_ADV_USD:,} over the previous {ADV_WINDOW} "
    "sessions - the screen 02_labels and 03_financial_features apply."
)

# %% [markdown]
# ## Why this panel is given regime and volatility features
#
# The strategy this case study builds ranks stocks cross-sectionally and holds the top names
# against the bottom ones. A ranking like that earns steadily for long stretches and then
# gives several years back in a few weeks, and the weeks it gives them back in are the ones
# where the market turns sharply after a decline - Daniel and Moskowitz (2016) call these
# momentum crashes and show they cluster where volatility is high and the market is
# rebounding. A model that only sees each stock's own price history has no way to tell those
# weeks apart from any other.
#
# So the three transforms fitted below each supply something a per-stock price feature cannot:
#
# - **Where the whole cross-section currently sits.** The Wasserstein clustering in Section 2
#   compares the recent month of market-wide returns against two reference months learned from
#   the training window, and reports how close the match is. Its useful output is the *distance*
#   rather than the state, because a crash happens while the market is between states.
# - **How turbulent each stock is right now.** The GARCH fit in Section 4 gives each stock a
#   conditional volatility that responds to its own recent moves, which is the quantity the
#   crash literature conditions on.
# - **A price level that is still usable as a regressor.** Fractional differencing in Section 3
#   keeps part of what the level of a price knows, which a return has thrown away entirely.
#
# None of the three is a trading rule. They are inputs a model in the later stages can
# condition on, and whether conditioning on them helps is a question for `05_evaluation` and
# the model notebooks, not for this page.

# %% [markdown]
# ## 1. Load the panel and screen it
#
# Two screens run here, and they are the ones
# [`02_labels`](02_labels.ipynb) and [`03_financial_features`](03_financial_features.ipynb)
# already run, rebuilt from the same constants on the same columns so that all three stages
# describe one universe.
#
# **Sessions are numbered first.** The archive carries a small number of stray prints on dates
# the exchange held no market. A date that was never open is not a date a position can be taken
# on, and `get_sessions` identifies them: a date that maps to itself is a session, and a stray
# print maps to a neighbouring one. Dropping them and numbering what is left gives a counter
# whose difference between two rows is a count of sessions rather than a count of rows. Every
# window on this page needs it - a variance recursion, a fractional-difference convolution and
# a rolling turnover average all read their input in order and treat consecutive elements as
# consecutive sessions.
#
# **Then eligibility**, on three conditions: a printed close above \$5, dollar volume
# `close * volume` averaging above \$1M over the previous month, and that month being an
# unbroken run of sessions rather than whatever twenty-one rows the stock happens to have. The
# first two legs read figures the tape carried on the day, so neither depends on a corporate
# action that had not happened yet, and Section B of [`02_labels`](02_labels.ipynb) derives why
# the adjusted close cannot serve for either. The third is what stops a stock returning from a
# halt qualifying on volume it traded before the halt.
#
# **Eligibility is applied only to what is emitted, never to what the transforms read.** On the
# eligible frame a per-stock window would count *eligible* rows, so a stock that falls below a
# threshold for two years and recovers would have its convolution and its variance recursion
# reach straight across the excursion as though those were consecutive sessions. Both run on
# the full session panel; the eligible frame decides only which rows leave this notebook.
#
# The digest of the panel read here has to equal the one [`02_labels`](02_labels.ipynb)
# recorded against the label file this notebook scores against in Section 7; the assertion
# below is what makes the two files comparable rather than merely both present.

# %%
# The six columns this notebook reads, and `lazy=True` so the projection reaches the parquet
# scan rather than a frame that has already been read. The archive carries fourteen columns
# over 14.5M rows, and the eight not named here - the unadjusted open, high and low, their
# adjusted counterparts, the dividend and the split ratio - are read nowhere on this page.
# Selecting after an eager load still materializes all eight: 1.40 GB against 0.53 GB, at a
# peak of 2.88 GB against 1.66 GB, measured on 2026-09-10.
#
# The digest below is the reason this is safe to assert rather than merely likely.
# `value_digest` hashes the columns it is given and nothing else, and all five it is given
# are among the six: loading both ways on 2026-09-10 returned 81db7d7920165013 either way,
# which is the value `02_labels` recorded against the label file.
READ_COLS = ["symbol", "timestamp", "close", "volume", "adj_close", "adj_volume"]

raw_df = (
    load_us_equities(start_date=START_DATE, end_date=END_DATE, lazy=True)
    .select(READ_COLS)
    .collect()
)

if raw_df.schema["timestamp"] == pl.Datetime:
    raw_df = raw_df.with_columns(pl.col("timestamp").dt.date().alias("timestamp"))

raw_df = raw_df.sort(["symbol", "timestamp"])

MARKET_DATA_DIGEST = value_digest(raw_df, ["symbol", "timestamp", "close", "volume", "adj_close"])
LABEL_INPUT_DIGEST = read_digest(CASE_DIR / "labels" / f"{PRIMARY_LABEL}.parquet")["inputs"][
    "market_data"
]
print(f"market_data digest: {MARKET_DATA_DIGEST}")
assert MARKET_DATA_DIGEST == LABEL_INPUT_DIGEST, (
    f"the labels were written against market_data {LABEL_INPUT_DIGEST} and this stage read "
    f"{MARKET_DATA_DIGEST}. Re-run 02_labels before scoring features against its output."
)

# %%
# The session counter, built exactly as 02_labels and 03_financial_features build it.
_dates = raw_df.select("timestamp").unique().sort("timestamp")
_settling_session = pl.Series(
    TradingCalendar(CALENDAR)
    .get_sessions(pd.DatetimeIndex(_dates["timestamp"].to_list(), tz="UTC"))
    .to_numpy()
).cast(pl.Date)
_sessions = (
    _dates.filter(_settling_session == pl.col("timestamp"))
    .with_row_index("session")
    .with_columns(pl.col("session").cast(pl.Int64))
)
_archive_rows = raw_df.height
raw_df = raw_df.join(_sessions, on="timestamp", how="inner").sort(["symbol", "timestamp"])
print(
    f"{_sessions.height:,} of {_dates.height:,} dates in the archive are {CALENDAR} sessions; "
    f"the other {_dates.height - _sessions.height} carry stray prints and take "
    f"{_archive_rows - raw_df.height:,} rows with them"
)

# The archive is missing one session the exchange held, `2017-11-08`, a Wednesday. It is absent
# UPSTREAM - the raw archive carries no row on it - so no stage of this case study drops it, and
# over the whole archive, 1962-01-02 to 2018-03-27, it is the only NYSE session of 14,156 that is
# absent. It is declared here so it passes deliberately, and so a second one refuses instead.
#
# The filter above answers one direction only: which archive dates the exchange never held. The
# other direction, a session the exchange DID hold that the archive never printed, leaves no row
# to test - nothing raises, every query succeeds, and one day's rows are simply gone. That is how
# this one survived until two counts of an unrelated quantity came out one apart.
KNOWN_ABSENT_SESSIONS = [date(2017, 11, 8)]

_declared = assert_sessions_complete(
    _sessions["timestamp"].to_list(),
    calendar=CALENDAR,
    known_absent=KNOWN_ABSENT_SESSIONS,
    source="04_model_based_features session index",
)
print(
    f"Every {CALENDAR} session between {_sessions['timestamp'].min()} and "
    f"{_sessions['timestamp'].max()} is in the archive except {_declared}, which is declared "
    "above"
)

# %%
raw_df = raw_df.with_columns(
    (pl.col("adj_close") / pl.col("adj_close").shift(1).over("symbol") - 1).alias("returns"),
    (pl.col("close") * pl.col("volume")).alias("dollar_volume"),
)
raw_df = raw_df.with_columns(
    pl.col("dollar_volume").rolling_mean(ADV_WINDOW).over("symbol").alias("adv_21d"),
    (pl.col("session") - pl.col("session").shift(ADV_WINDOW - 1) == ADV_WINDOW - 1)
    .over("symbol")
    .alias("adv_covered"),
)

ELIGIBLE = pl.col("adv_covered") & (pl.col("close") > MIN_PRICE) & (pl.col("adv_21d") > MIN_ADV_USD)
df = raw_df.filter(ELIGIBLE)

print(
    f"{len(raw_df):,} session rows on {raw_df['symbol'].n_unique():,} symbols, "
    f"{raw_df['timestamp'].min()} to {raw_df['timestamp'].max()}"
)
print(
    f"{len(df):,} of them on {df['symbol'].n_unique():,} symbols pass all three conditions and "
    "are eligible to be emitted"
)

# %% [markdown]
# Those two totals are sums over twenty-eight years, and what every transform below actually
# works with is one day's slice of the panel. The figure is that slice through time: how many
# stocks are eligible on each session, with the two thresholds that read the count drawn across
# it.
#
# It is worth looking at before anything is fitted, because two of the decisions on this page
# are decisions about that count. The clustering in Section 2 takes a median across the slice
# and skips any date holding fewer than `XS_MIN_STOCKS` names, so where that line sits relative
# to the curve says whether the threshold ever binds. And the count rises for most of the
# sample before turning down, which is why two windows of the same length in sessions are not
# comparable in how many stocks they saw.

# %%
_coverage = df.group_by("timestamp").len().sort("timestamp")

fig, ax = plt.subplots(figsize=FIGSIZE["single"])
ax.plot(
    _coverage["timestamp"],
    _coverage["len"],
    color=COLORS["blue"],
    lw=0.8,
    label="eligible on the session",
)
ax.axhline(
    XS_MIN_STOCKS,
    color=COLORS["neutral"],
    ls=":",
    lw=1.0,
    label=f"{XS_MIN_STOCKS}: below this a date is not summarized",
)
ax.set_ylim(0, None)
ax.set_xlabel("Date")
ax.set_ylabel("Eligible stocks")
ax.legend(frameon=False, fontsize=7, loc="upper left")
add_message_title(
    ax,
    "The panel these transforms fit on grows for two decades, then turns down",
    subtitle="Stocks passing the price and dollar-volume screen on each session",
)
show_with_alt(
    fig,
    "A single line counts the stocks eligible on each session across the sample. It climbs "
    "for about two decades, drops sharply in the 2008 crisis, recovers to its highest point "
    "and then declines over the final years. One flat reference line sits far below it, "
    "marking the count below which a date is not summarized; the eligible-stock line stays "
    "above it throughout.",
)

# %% [markdown]
# ## 1b. What bounds an estimate, and what selects a row
#
# **A schedule bounds an estimate.** Each model below spends a burn-in, is fitted on the
# sessions before the block it is about to speak for, emits values over that block, and is
# re-estimated on everything up to the start of the next one. No session is ever used to fit
# the model that describes it, at any position in the series rather than only after some
# window closed. Past `holdout_start` nothing is re-estimated: the last estimate made on
# development sessions carries the rest, because a coefficient refitted on a held-out session
# is a parameter estimated from the holdout however causal the recursion around it looks.
#
# **A fold selects rows.** The walk-forward folds are resolved here because Section 7 scores
# each column over the sessions the folds validate on, and because the figure below is worth
# looking at before anything is fitted. They enter no fit and they are not a key of the
# artifact: a stock-date carries one value whichever fold later reads it.
#
# **They are resolved from the label file, through the same call the model stages use.** A
# walk-forward splitter counts backward from the holdout boundary in rows of whatever frame it
# is handed, and it seals the end of each validation window by the horizon of the label being
# predicted. Both of those are properties of the label file, not of the price panel, and
# `modeling_fold_boundaries` reads the label file's own date index and its own configured
# buffer and horizon. It is what `load_modeling_dataset` calls on the other side of the join,
# so the sessions this notebook scores on are the sessions a model is validated on rather
# than a second set that happens to carry the same numbers.
#
# **Both ends of a window are inclusive**: `train_end` is the last session a fold trains on
# and `val_end` the last session it is scored on.

# %%
holdout_start = date.fromisoformat(HOLDOUT_START)
holdout_end = date.fromisoformat(END_DATE)

# The trading calendar the folds are counted on, taken from the label file itself.
SESSIONS = sorted(
    pl.read_parquet(CASE_DIR / "labels" / f"{PRIMARY_LABEL}.parquet")["timestamp"]
    .unique()
    .to_list()
)

# Ordered by the sessions they score, so the printout below reads chronologically. The
# `MAX_FOLDS` reduction underneath names the fold ids it keeps rather than taking a head
# slice of that order, which is a count that says nothing about which windows it kept.
folds = sorted(
    (
        {
            "fold": split["fold"],
            "train_start": split["train_start"],
            "train_end": split["train_end"],
            "test_start": split["val_start"],
            "test_end": split["val_end"],
        }
        for split in modeling_fold_boundaries(CASE_STUDY_ID, PRIMARY_LABEL)
    ),
    key=lambda f: f["test_start"],
)
if MAX_FOLDS > 0:
    folds = select_folds(folds, range(MAX_FOLDS))

print(f"{len(folds)} cross-validation folds, which select rows and bound no estimate:")
for f in folds:
    print(
        f"  Fold {f['fold']}: trained on {f['train_start']} to {f['train_end']}, "
        f"scored on {f['test_start']} to {f['test_end']}"
    )

# The one condition Section 7 rests on, asserted rather than described: a validation window
# that crept past the boundary would still be scored and still print a table.
for f in folds:
    assert f["train_end"] < f["test_start"], f"fold {f['fold']} trains into its own validation"
    assert f["test_end"] < holdout_start, (
        f"fold {f['fold']} is scored through {f['test_end']}, past the holdout opening "
        f"{holdout_start}"
    )
print(f"  every validation window closes before {holdout_start}")


# %%
def in_validation_windows(column: str = "timestamp") -> pl.Expr:
    """True on a session inside any cross-validation fold's validation window."""
    spans = [
        (pl.col(column) >= pl.lit(f["test_start"]).cast(pl.Date))
        & (pl.col(column) <= pl.lit(f["test_end"]).cast(pl.Date))
        for f in folds
    ]
    expr = spans[0]
    for span in spans[1:]:
        expr = expr | span
    return expr


# %% [markdown]
# Each row of the figure is one fold: the filled bar is the span it trains on, the open bar
# the span it is scored on, and the dashed rule is the date the holdout opens. The windows
# roll back one year at a time and no open bar crosses the rule, so every column Section 7
# scores is scored on development sessions.
#
# There is no estimation window on this figure, because a fold does not carry one. Section
# 4b draws the schedule that bounds the estimates.

# %%
fig, ax = plt.subplots(figsize=FIGSIZE["single_wide"])
for row, f in enumerate(folds):
    tr0, tr1 = f["train_start"], f["train_end"]
    te0, te1 = f["test_start"], f["test_end"]
    ax.barh(row, (tr1 - tr0).days, left=tr0, height=0.62, color=COLORS["blue"], alpha=0.85)
    ax.barh(
        row,
        (te1 - te0).days,
        left=te0,
        height=0.62,
        facecolor="none",
        edgecolor=COLORS["neutral"],
        linewidth=1.2,
    )
ax.axvline(holdout_start, color=COLORS["copper"], ls="--", lw=1.4)
ax.set_yticks(range(len(folds)))
ax.set_yticklabels([str(f["fold"]) for f in folds], fontsize=7)
ax.invert_yaxis()
ax.set_xlabel("Date")
ax.set_ylabel("Fold")
add_message_title(
    ax,
    "Every session a column is scored on lies before the date the holdout opens",
    subtitle="Filled: the fold's training span. Outlined: the span it is scored on",
)
show_with_alt(
    fig,
    "One horizontal row per cross-validation fold. Each row is a long filled bar for the "
    "window the fold trains on, followed by a short outlined bar for the window it is scored "
    "on. The pairs step later in time down the rows. A dashed vertical line marks the date "
    "the holdout opens, and no bar of either kind crosses it.",
)

# %% [markdown]
# ## 2. Wasserstein regime distance
#
# At each date the median return is taken across every eligible stock trading that day. That
# one number per date is how the centre of the whole cross-section moves, and it is the series
# everything in this section reads.
#
# The method compares one recent month of that series against reference months learned from
# the training window, and it needs a way to say how far apart two months are. Two months of
# returns are two collections of twenty-one numbers, and the natural comparison is not
# value-by-value in date order - the same month reordered is the same market - but
# distribution against distribution. The **Wasserstein distance** measures exactly that: sort
# both collections, pair the smallest with the smallest and the largest with the largest, and
# average how far each pair has to move. It answers "how much return would have to be shifted,
# and how far, to turn one month into the other".
#
# With a distance in hand, ordinary k-means applies. **k-means** repeatedly assigns each
# window to its nearest of $k$ reference windows and then recomputes each reference as the
# centre of the windows assigned to it, until the references stop moving. Those references are
# called **centroids**, and with $k=2$ the two the algorithm settles on separate the calm,
# mildly positive months from the falling, turbulent ones. Two states is the coarsest split
# that can express that distinction, and it is the one the momentum-crash literature works in.
#
# The centroids are re-estimated on the schedule: fitted on the market history up to a
# boundary, held fixed while the next quarter of windows is scored against them, then fitted
# again on everything up to the following boundary. Every stock carries the same value on a
# date, because the series being clustered is market-wide.
#
# The estimator itself - the lifting, the distance, the barycenter and the k-means around them -
# is `lift_stream`, `wasserstein_distance_1d`, `wasserstein_barycenter_1d` and
# `fit_wasserstein_kmeans` in `case_studies/utils/temporal.py`, beside the HMM helpers the other
# case studies fit their regimes with. This notebook composes them on the schedule below.
# [`09_model_based_features/12_wasserstein_regimes`](../../09_model_based_features/12_wasserstein_regimes.ipynb)
# builds the same four objects from nothing, for a reader who wants to see the algorithm rather
# than use it.


# %% [markdown]
# ### The series the clustering reads
#
# One median per date, over the eligible stocks that traded that date. Dates whose
# cross-section is thinner than `XS_MIN_STOCKS` are dropped rather than summarized, because a
# median over a handful of names describes those names and not the market.

# %%
xs_stats = (
    df.filter(pl.col("returns").is_not_null())
    .group_by("timestamp")
    .agg(
        pl.col("returns").median().alias("xs_median_ret"),
        pl.col("returns").count().alias("n_stocks"),
    )
    .sort("timestamp")
    .filter(pl.col("n_stocks") >= XS_MIN_STOCKS)
)

market_ret = xs_stats["xs_median_ret"].to_numpy()
dates = xs_stats["timestamp"].to_list()

print(
    f"{len(xs_stats):,} dates carry a cross-section of at least {XS_MIN_STOCKS} eligible "
    f"stocks and are summarized; the median date carries "
    f"{int(xs_stats['n_stocks'].median()):,}"
)

# %% [markdown]
# ### Fitting the centroids on the schedule
#
# The clustering is one walk over the market series. It spends `REGIME_BURNIN` sessions
# before its first fit, scores the next `REGIME_REFIT_EVERY` sessions against the centroids
# that fit produced, and then re-estimates on everything up to the start of the block after
# that. The reference a window is compared against was therefore learned from sessions that
# close before the window opens, at every session rather than only inside a fold.
#
# What the walk emits per date is the assigned cluster, the distance to the nearer and the
# farther centroid, their ratio, and how differently the window's best and worst days sit
# against the centroid it matched.
#
# k-means labels are arbitrary - which of the two states the algorithm happens to call zero
# depends on where it started - so after each fit the two are reordered by their mean, and
# state zero is always the lower-return one. Without that step a downstream model would see
# the same market condition under one number before a refit and the other number after it.


# %%
def fit_regime_centroids(train: FloatArray) -> FloatArray:
    """Cluster the training prefix into `N_CLUSTERS` reference windows, lowest mean first."""
    lifted = lift_stream(train[:, 0], WASSERSTEIN_WINDOW, WASSERSTEIN_OVERLAP)
    _, centroids = fit_wasserstein_kmeans(
        lifted.sorted_segments, n_clusters=N_CLUSTERS, random_state=SEED
    )
    return centroids[np.argsort([c.mean() for c in centroids])]


def assign_regime_features(centroids: FloatArray, prefix: FloatArray) -> FloatArray:
    """Score every session of a prefix against fitted centroids.

    One row per input row, so the walk can keep the block it asked for. A session is scored
    on the `WASSERSTEIN_WINDOW` sessions strictly before it, which is why the first
    `WASSERSTEIN_WINDOW` rows carry no value: there is no complete window behind them.
    """
    series = prefix[:, 0]
    out = np.full((len(series), 5), np.nan, dtype=float)
    if len(series) <= WASSERSTEIN_WINDOW:
        return out

    # Window ending at t-1 and starting at t-WASSERSTEIN_WINDOW, for every t from the window
    # length onwards. Sorted, because the distance below compares distributions rather than
    # dates: the same month reordered is the same market.
    windows = np.sort(
        np.lib.stride_tricks.sliding_window_view(series[:-1], WASSERSTEIN_WINDOW), axis=1
    )
    distances = np.stack(
        [wasserstein_distance_1d(windows, centroids[k][None, :]) for k in range(len(centroids))],
        axis=1,
    )
    cluster = distances.argmin(axis=1)
    nearest = centroids[cluster]
    min_dist = distances.min(axis=1)
    max_dist = distances.max(axis=1)
    tail_div = np.abs(windows[:, -5:] - nearest[:, -5:]).mean(axis=1) - np.abs(
        windows[:, :5] - nearest[:, :5]
    ).mean(axis=1)

    out[WASSERSTEIN_WINDOW:] = np.column_stack(
        [cluster, min_dist, max_dist, min_dist / (max_dist + 1e-10), tail_div]
    )
    return out


# %% [markdown]
# The walk is driven once over the whole market series. `freeze_after` is the count of
# development sessions: past it the last pre-holdout centroids are reused rather than
# re-estimated, so no clustering reads a held-out session. The burn-in has to be long enough
# to lift into at least the number of windows k-means needs, which is asserted rather than
# assumed - a shortened run that fell below it would return centroids fitted on two windows
# and report nothing about it.

# %%
REGIME_COLS = [
    "wass_cluster",
    "wass_dist_min",
    "wass_dist_max",
    "wass_dist_ratio",
    "wass_tail_div",
]

_step = WASSERSTEIN_WINDOW - WASSERSTEIN_OVERLAP
_min_windows = 2 * N_CLUSTERS - 1
assert WASSERSTEIN_WINDOW + _min_windows * _step <= REGIME_BURNIN, (
    f"a {REGIME_BURNIN}-session burn-in lifts into "
    f"{max(0, (REGIME_BURNIN - WASSERSTEIN_WINDOW) // _step + 1)} windows, fewer than the "
    f"{_min_windows + 1} k-means needs to separate {N_CLUSTERS} centroids"
)

regime_fits: list[dict] = []


def _fit_and_record(train: FloatArray) -> FloatArray:
    centroids = fit_regime_centroids(train)
    regime_fits.append(
        {
            "fit_end_session": dates[len(train) - 1],
            "n_fit": len(train),
            "stress_centroid_mean": float(centroids[0].mean()),
            "normal_centroid_mean": float(centroids[-1].mean()),
            "centroid_separation": float(np.abs(centroids[-1] - centroids[0]).mean()),
        }
    )
    return centroids


regime_freeze_after = int(sum(1 for d in dates if d < holdout_start))
regime_values = walk_forward_feature(
    market_ret.reshape(-1, 1),
    timestamps=dates,
    burnin=REGIME_BURNIN,
    refit_every=REGIME_REFIT_EVERY,
    fit=_fit_and_record,
    apply=assign_regime_features,
    n_features=len(REGIME_COLS),
    freeze_after=regime_freeze_after,
)

wass_df = (
    pl.DataFrame(
        {
            "timestamp": dates,
            **{col: regime_values[:, i] for i, col in enumerate(REGIME_COLS)},
        }
    )
    # `walk_forward_feature` marks the burn-in with `np.nan`, which polars keeps as a float
    # rather than a null: without this the burn-in rows survive the drop below and the
    # artifact's null counts read as zero on a column that has no value for three years.
    .with_columns(pl.col(col).fill_nan(None) for col in REGIME_COLS)
    .drop_nulls(subset=REGIME_COLS)
    .with_columns(pl.col("wass_cluster").cast(pl.Int64))
)
regime_fit_df = pl.DataFrame(regime_fits)

print(
    f"{len(regime_fits)} clusterings over {len(dates):,} sessions, the last fitted on "
    f"sessions through {regime_fit_df['fit_end_session'].max()}; "
    f"{wass_df.height:,} sessions carry a regime value, from {wass_df['timestamp'].min()}"
)
_cluster_counts = wass_df.group_by("wass_cluster").len().sort("wass_cluster")
for row in _cluster_counts.iter_rows(named=True):
    state = "lower-return" if row["wass_cluster"] == 0 else "higher-return"
    print(f"  state {row['wass_cluster']} ({state}): {row['len']:,}")

# The schedule is the provenance, so it is checked against the schedule rather than against
# prose. Every fit consumed a prefix ending before the block it spoke for, and no fit
# consumed a held-out session.
_scheduled = refit_boundaries(len(dates), REGIME_BURNIN, REGIME_REFIT_EVERY)
_estimated = [pair for pair in _scheduled if pair[0] <= regime_freeze_after]
assert len(regime_fits) == len(_estimated), (
    f"{len(regime_fits)} clusterings against a schedule of {len(_estimated)}"
)
assert regime_fit_df["fit_end_session"].max() < holdout_start, (
    f"a clustering read sessions through {regime_fit_df['fit_end_session'].max()}, inside "
    f"the holdout opening {holdout_start}"
)
print(
    f"  every fit ended before {holdout_start}; the {len(_scheduled) - len(_estimated)} "
    "blocks past it reuse the last development estimate"
)

# %% [markdown]
# ### What the clustering inferred, on the sessions a model is scored over
#
# The figure draws the quantity the feature actually carries, over the sessions the folds
# validate on. Every value on it was produced by centroids fitted before the window it
# scores, so what is plotted is a chain of out-of-sample assignments from a hundred different
# fits rather than one fit's view of the whole sample.
#
# The line is the trailing cross-sectional median return the assignment reads; the panel
# below it is the monthly share of sessions assigned to the low-return centroid. Nothing in
# the fitting procedure required those sessions to be the market's stressed ones.
#
# **The lower panel is still the weaker of the two outputs, and it is worth saying why.**
# State zero is whichever centroid has the lower mean *in the history that fit read*. That
# fixes the arbitrariness of k-means labelling within a fit; it does not make the number mean
# the same thing from one fit to the next, because a clustering estimated through the 2008
# decline and one estimated a decade later put their lower-return centroid in different
# places. What the panel shows is that the assignment nonetheless lands where a reader would
# expect it to: nearly every session of 2008 and 2002 is in the lower state and fewer than a
# tenth of 1995's are. `wass_dist_ratio` does not have the comparability problem, because it
# is a ratio of distances read against the fit that produced it. It answers a different
# question, though, and the difference matters: it says how firmly the window matches
# whichever centroid is nearest and discards which one that was, so a window sitting squarely
# in the calm state and one sitting squarely in the stressed state both drive it toward zero.
# A model that needs the direction still has to read the assignment.
#
# The assignment is aggregated to a monthly share rather than drawn as a daily strip. Sixteen
# years of daily flags give each session a fraction of a pixel, isolated days vanish, and the
# reader concludes the state stopped occurring when it did not.

# %%
_val_regime = wass_df.filter(in_validation_windows()).sort("timestamp")
_val_ret = xs_stats.join(_val_regime.select("timestamp", "wass_cluster"), on="timestamp").sort(
    "timestamp"
)
_smoothed = _val_ret.select(
    "timestamp",
    pl.col("xs_median_ret").rolling_mean(WASSERSTEIN_WINDOW).alias("trailing"),
    "wass_cluster",
).drop_nulls()

fig, (ax1, ax2) = plt.subplots(
    2,
    1,
    figsize=FIGSIZE["single"],
    sharex=True,
    height_ratios=[3, 1],
    gridspec_kw={"hspace": 0.22},
)
ax1.plot(_smoothed["timestamp"], _smoothed["trailing"], color=COLORS["blue"], lw=0.8)
ax1.axhline(0, color=COLORS["neutral"], lw=0.7)
ax1.set_ylabel("Trailing median return", fontsize=8)
ax1.locator_params(axis="y", nbins=4)
_monthly = (
    _smoothed.with_columns(pl.col("timestamp").dt.truncate("1mo").alias("month"))
    .group_by("month")
    .agg((pl.col("wass_cluster") == 0).mean().alias("share"))
    .sort("month")
)
ax2.fill_between(_monthly["month"], 0, _monthly["share"], color=COLORS["copper"], lw=0, step="mid")
ax2.set_ylim(0, 1)
ax2.set_yticks([0, 1])
ax2.set_ylabel("Share in the\nlower state", fontsize=7)
ax2.set_xlabel("Date")
add_message_title(
    ax1,
    "The lower-return state fills the years the market was falling",
    subtitle="Scored sessions only. Below: monthly share assigned to that state",
)
show_with_alt(
    fig,
    "Two stacked panels sharing a date axis over the scored sessions. The upper panel is a "
    "noisy trailing median return oscillating around zero, with its largest excursions in "
    "2008 and 2009. The lower panel is a filled area of the monthly share of sessions "
    "assigned to the lower-return state. It swings between the top and the bottom of the "
    "panel rather than trending: it is near the ceiling through 2000 to 2003 and again "
    "across 2008 and 2009, and close to the floor in the middle of the 1990s and again from "
    "2012 to 2014.",
)

_shaded = _smoothed.filter(pl.col("wass_cluster") == 0)
_runs = _smoothed.with_columns(
    (pl.col("wass_cluster").diff().fill_null(1) != 0).cum_sum().alias("run")
)
_run_lengths = _runs.filter(pl.col("wass_cluster") == 0).group_by("run").len()["len"]
print(
    f"scored sessions {_smoothed.height:,}, assigned to the lower-return state "
    f"{_shaded.height:,} ({_shaded.height / _smoothed.height:.0%}); mean trailing return "
    f"{_shaded['trailing'].mean():+.5f} in that state against "
    f"{_smoothed.filter(pl.col('wass_cluster') == 1)['trailing'].mean():+.5f} in the other"
)
print(
    f"  {_run_lengths.len():,} runs, median {_run_lengths.median():.0f} sessions and longest "
    f"{_run_lengths.max():,}; first assigned {_shaded['timestamp'].min()}, last "
    f"{_shaded['timestamp'].max()}, and the scored span runs to "
    f"{_smoothed['timestamp'].max()}"
)

# %% [markdown]
# `wass_dist_ratio` is the second thing the clustering yields: the distance to the nearer
# centroid divided by the distance to the farther one. A window sitting squarely inside one
# state drives it toward zero and a window equidistant from both drives it toward one, so the
# feature carries how *certain* the match is rather than which state it picked. That is the
# part a momentum model needs, because momentum crashes fall at the transitions rather than
# inside either state.
#
# The cost of clustering the median and nothing else is worth stating plainly: this reads a
# shift in the centre of the cross-section, and a market that keeps its centre while its tails
# widen looks unchanged to it. Reaching that would mean clustering quantile vectors rather than
# a scalar, which is a different transform and not a tuning of this one.

# %% [markdown]
# ## 3. Fractional differencing
#
# A log price is not stationary and a log return has thrown away everything the level knew.
# Fractional differencing (Hosking 1981; Lopez de Prado 2018) takes the difference to a
# non-integer order $d$, which puts a dial between the two: at $d=0$ the series is the level
# and at $d=1$ it is the first difference, and every value in between trades some memory for
# some stationarity. `FFD_D` is the equity-class default this notebook uses.
#
# **The default is measured here rather than quoted.** The cell below runs the whole grid
# `FFD_D_GRID` on a sample of stocks and reports, for each order, the correlation between
# the differenced series and the original log price - how much of the level's memory
# is retained - against the share of sampled stocks whose augmented Dickey-Fuller test rejects
# a unit root. Those are the two quantities the choice trades off, and neither is knowable
# without running it.
#
# **Nothing here is estimated, so there is no schedule to put it on.** The FFD weights are a
# closed-form function of $d$ and of the truncation threshold, so the transform carries no
# estimation window at all and is computed once over each stock's whole series. That makes it
# the useful contrast for the sections either side of it: the hazard this stage is about is
# *estimation*, not transformation, and a transform with no parameters has none of it.


# %%
def apply_ffd_per_symbol(
    data: pl.DataFrame, d: float = FFD_D, threshold: float = FFD_THRESHOLD
) -> pl.DataFrame:
    """Apply fractional differencing to log prices per symbol.

    Returns DataFrame with (symbol, date, ffd_log_price, ffd_log_volume).
    """
    results = []
    # Only the two columns the transform reads are partitioned. `partition_by` copies the
    # frame it is handed into one frame per symbol, so partitioning the caller's panel
    # would hold a second copy of every column on it - and this is called on the whole
    # panel, not the eligible subset.
    by_symbol = (
        data.select(["symbol", "timestamp", "adj_close", "adj_volume"])
        .sort(["symbol", "timestamp"])
        .partition_by("symbol", as_dict=True)
    )

    n_success = 0
    n_fail = 0

    for (sym,) in sorted(by_symbol):
        # Popped rather than read: the partition is dead once its result is appended, and
        # holding all of them to the end of the loop keeps a whole copy of the panel alive
        # alongside the results being built from it.
        sym_data = by_symbol.pop((sym,))

        if len(sym_data) < 100:
            n_fail += 1
            continue

        log_price = sym_data["adj_close"].log()
        # Floor volume at 1 to avoid log(0) = -inf
        log_vol = sym_data["adj_volume"].clip(lower_bound=1).log()

        try:
            ffd_price = ffdiff(log_price, d=d, threshold=threshold)
            ffd_vol = ffdiff(log_vol, d=d, threshold=threshold)

            sym_result = pl.DataFrame(
                {
                    "symbol": [sym] * len(sym_data),
                    "timestamp": sym_data["timestamp"],
                    "ffd_log_price": ffd_price,
                    "ffd_log_volume": ffd_vol,
                }
            ).drop_nulls()

            if len(sym_result) > 0:
                results.append(sym_result)
                n_success += 1
        except Exception:
            n_fail += 1

    print(f"  FFD: {n_success} symbols succeeded, {n_fail} failed/skipped")
    return pl.concat(results) if results else pl.DataFrame()


# %% [markdown]
# The sweep runs on a sample of stocks - every symbol with a long enough eligible history,
# taken at a fixed stride so the sample is not the alphabet's first few hundred names. The
# augmented Dickey-Fuller test asks whether a series has a unit root, which is the formal
# version of "wanders without returning"; what is reported is the *share* of sampled stocks
# whose test rejects that, because a single stock's test says very little and the question is
# whether the order works across the panel.
#
# **The sweep stops at the holdout boundary, on both counts.** It is a measurement that argues
# for a setting, so it is a development-time decision, and a development-time decision may not
# read a held-out bar. That governs which stocks it samples as much as which bars it reads: a
# sample drawn on history-length over the whole panel would let a stock's post-2016 record
# decide whether it is in the sample at all.

# %%
_ffd_dev = raw_df.filter(pl.col("timestamp") < holdout_start)
_ffd_symbols = (
    df.filter(pl.col("timestamp") < holdout_start)
    .group_by("symbol")
    .len()
    .filter(pl.col("len") >= 2000)
    .sort("symbol")["symbol"]
    .to_list()
)
_ffd_sample = _ffd_symbols[:: max(1, len(_ffd_symbols) // 120)][:120]
_ffd_panel = _ffd_dev.filter(pl.col("symbol").is_in(_ffd_sample)).sort(["symbol", "timestamp"])
_ffd_by_symbol = _ffd_panel.partition_by("symbol", as_dict=True)

grid_rows = []
for d in FFD_D_GRID:
    corrs, rejects = [], []
    for key in sorted(_ffd_by_symbol):
        _lp = _ffd_by_symbol[key]["adj_close"].log().drop_nulls()
        if len(_lp) < 500:
            continue
        _fd = ffdiff(_lp, d=d, threshold=FFD_THRESHOLD)
        _pair = pl.DataFrame({"level": _lp, "ffd": _fd}).drop_nulls()
        if _pair.height < 500 or _pair["ffd"].std() == 0:
            continue
        corrs.append(abs(float(np.corrcoef(_pair["level"], _pair["ffd"])[0, 1])))
        rejects.append(adfuller(_pair["ffd"].to_numpy(), autolag="AIC")[1] < FDR_ALPHA)
    grid_rows.append(
        {
            "d": d,
            "memory": float(np.mean(corrs)),
            "stationary_share": float(np.mean(rejects)),
            "n_symbols": len(corrs),
        }
    )

ffd_grid = pl.DataFrame(grid_rows)
# The sweep is over, and `_ffd_dev` is the panel before the holdout - most of the archive,
# and the largest thing on this page that nothing below reads. A name bound in a notebook
# stays bound until the kernel exits, so without this it is still live when the volatility
# fits fork, and every worker inherits it: 0.95 GB of the 3.56 GB the parent held at the
# fork, measured on the full panel on 2026-09-10.
del _ffd_dev, _ffd_panel, _ffd_by_symbol

print(
    f"{len(FFD_D_GRID)} differencing orders, each measured on the same "
    f"{ffd_grid['n_symbols'].max()} sampled stocks, on bars before {holdout_start}"
)
display(ffd_grid)

# %% [markdown]
# The two curves cross, and where they cross is the whole argument for a fractional order.
# Memory falls with $d$ and the share of stocks that pass the stationarity test rises with
# it; the first difference sits at the right-hand end, stationary and remembering nothing of
# the level.

# %%
_chosen = ffd_grid.filter(pl.col("d") == FFD_D).row(0, named=True)

fig, ax = plt.subplots(figsize=FIGSIZE["single"])
ax.plot(ffd_grid["d"], ffd_grid["memory"], color=COLORS["blue"], marker="o", ms=4, label="memory")
ax.plot(
    ffd_grid["d"],
    ffd_grid["stationary_share"],
    color=COLORS["copper"],
    marker="s",
    ms=4,
    label="share passing ADF",
)
ax.axvline(FFD_D, color=COLORS["neutral"], ls="--", lw=1.2)
ax.set_xlabel("Differencing order $d$")
ax.set_ylabel("Correlation with the log level / share of stocks")
ax.set_ylim(0, 1.05)
ax.legend(frameon=False, fontsize=8, loc="center right")
add_message_title(
    ax,
    "A fractional order keeps memory the first difference throws away",
    subtitle="Correlation with the log level, and the share of stocks rejecting a unit root",
)
show_with_alt(
    fig,
    "Two curves against the differencing order on the horizontal axis. One falls "
    "monotonically from one at order zero to near zero at order one, the correlation the "
    "series keeps with its own log level. The other rises from near zero to one and then "
    "runs flat, the share of stocks whose unit root is rejected. A dashed vertical line "
    "marks the order chosen; it stands where the rising curve has just reached its ceiling "
    "and the falling curve still retains a substantial part of its height.",
)

print(
    f"at d={FFD_D}: memory {_chosen['memory']:.3f}, {_chosen['stationary_share']:.1%} of "
    f"{_chosen['n_symbols']} sampled stocks reject a unit root | "
    f"at d={ffd_grid['d'].max()}: memory "
    f"{ffd_grid.filter(pl.col('d') == ffd_grid['d'].max())['memory'][0]:.3f}, "
    f"{ffd_grid.filter(pl.col('d') == ffd_grid['d'].max())['stationary_share'][0]:.1%}"
)
print(
    f"  the weight vector at d={FFD_D} truncates at "
    f"{len(get_ffd_weights(FFD_D, threshold=FFD_THRESHOLD))} lags"
)

# %% [markdown]
# ### Apply the transform to the panel
#
# On the complete price series per stock, for the reason Section 1 states: the weight vector
# reaches back hundreds of sessions, and on the screened frame those would be eligible rows
# rather than sessions.

# %%
print("Computing fractional differencing features...")
ffd_df = apply_ffd_per_symbol(raw_df)
print(f"FFD features: {len(ffd_df):,} rows, {ffd_df['symbol'].n_unique()} symbols")

# %% [markdown]
# ## 4. GARCH conditional volatility
#
# This is the section the stage is really about. A GARCH conditional volatility is not a
# function of a stock's past returns alone - it is a function of $(\mu, \omega, \alpha,
# \beta)$, and those come from a maximum-likelihood fit over some window. Fit them once over
# everything and every row's volatility knows the whole sample.
#
# So each stock gets its own walk over its own return history:
#
# 1. Spend `GARCH_BURNIN` returns, which carry no value and pay for the first estimate.
# 2. Fit GARCH(1,1) by maximum likelihood on the returns up to the start of the next block.
# 3. Run the variance recursion forward with those coefficients held fixed, keep the block's
#    own rows, and re-estimate on everything up to the start of the block after it.
# 4. Past the holdout boundary stop re-estimating and carry the last development coefficients
#    forward, so no coefficient is estimated from a held-out session.
#
# Every stock that clears the burn-in is fitted. A stock whose history is shorter takes the
# market-level fit, which is the same walk over the cross-sectional median return, so every
# emitted row carries a conditional volatility.
#
# The returns handed to every fit come from the **complete** per-symbol series. A variance
# recursion reads its input in order and treats consecutive elements as consecutive sessions;
# feeding it the eligible rows only would splice the two sides of an ineligible spell
# together and price the jump across it as one day's move.


# %%
# The specification comes from `setup.yaml::model_based.garch` rather than from this line,
# so the model a reader is told about and the model that is fitted are the same statement.
# These are `arch_model`'s own argument names and are passed through unchanged.
GARCH_KW = {
    "mean": str(_GARCH["mean"]),
    "vol": str(_GARCH["vol"]),
    "p": int(_GARCH["p"]),
    "o": int(_GARCH["o"]),
    "q": int(_GARCH["q"]),
    "dist": str(_GARCH["dist"]),
}

# `garch11_conditional_volatility` is a GARCH(1,1) recursion, optionally with the one
# asymmetry term. A declared parameter it cannot represent has to refuse here rather than be
# silently dropped on the way to the filter, which is the failure a declaration exists to
# prevent: the fit would estimate one model and the emitted column would carry another.
if GARCH_KW["vol"] != "GARCH" or GARCH_KW["p"] != 1 or GARCH_KW["q"] != 1:
    raise ValueError(
        f"model_based.garch declares vol={GARCH_KW['vol']} p={GARCH_KW['p']} q={GARCH_KW['q']}; "
        "the filter this notebook emits through is GARCH(1,1)"
    )
if GARCH_KW["o"] not in (0, 1):
    raise ValueError(f"model_based.garch declares o={GARCH_KW['o']}; the filter carries 0 or 1")
if GARCH_KW["mean"] not in ("Constant", "Zero"):
    raise ValueError(
        f"model_based.garch declares mean={GARCH_KW['mean']}; the filter subtracts a single "
        "constant, so only Constant and Zero are representable"
    )


def garch_walk(
    payload: tuple[str, FloatArray, int, list],
) -> tuple[str, FloatArray, list[dict]]:
    """One walk-forward GARCH per series: refit on schedule, filter forward, freeze at the
    holdout.

    Takes and returns percent returns' annualized conditional volatility in decimal, ``nan``
    over the burn-in and over any block whose fit did not converge, and one record per
    estimation so Section 4b can measure what re-estimating moved.
    """
    symbol, returns_pct, freeze_after, sessions = payload
    fits: list[dict] = []

    def fit(X_train: FloatArray) -> dict[str, float]:
        result = arch_model(X_train[:, 0], **GARCH_KW).fit(disp="off", show_warning=False)
        # `arch` returns a result whatever the optimizer did and only warns, which
        # `show_warning=False` then swallows. A parameter vector the search never converged
        # on is not an estimate, so it is rejected here and `on_fit_error="skip"` leaves the
        # block empty - the coverage table in Section 5 is where that cost shows up.
        if result.convergence_flag != 0:
            raise RuntimeError(
                f"the variance model did not converge on {symbol}: scipy flag "
                f"{result.convergence_flag}"
            )
        coefficients = {
            # A zero-mean specification estimates no `mu` at all, so the parameter vector has
            # no such entry to read; the filter still subtracts one and it is zero.
            "mu": float(result.params["mu"]) if GARCH_KW["mean"] == "Constant" else 0.0,
            "omega": float(result.params["omega"]),
            "alpha": float(result.params["alpha[1]"]),
            # The leverage coefficient under `o=1`, zero under the symmetric model. Read from
            # the fit rather than assumed, so a declared `o` reaches the emitted values.
            "gamma": float(result.params["gamma[1]"]) if GARCH_KW["o"] else 0.0,
            "beta": float(result.params["beta[1]"]),
            # The value that seeds the recursion, computed by `arch` from the ESTIMATION
            # window's residuals and nothing else. It has to be produced here, where only
            # training returns are in scope: the array `apply` receives runs to the end of
            # the block being emitted, so a seed derived there would read the block's own
            # sessions.
            "backcast": float(result.model.volatility.backcast(np.asarray(result.resid))),
        }
        fits.append({"symbol": symbol, "fit_end": len(X_train), **coefficients})
        return coefficients

    def apply(coefficients: dict[str, float], X_prefix: FloatArray) -> FloatArray:
        # `garch11_conditional_volatility` rather than the fitted result object's own
        # `conditional_volatility`, which is what `arch_model(