利用 ARIMA、傅里叶和状态模型构建 CME 期货持有收益特征

# 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]
# # CME Futures: Features That Are Themselves Model Output
#
# Every feature so far has been a formula applied to past prices. This notebook builds
# features of a second kind: it estimates a statistical model from past prices and then
# emits what that model says about each session as the feature. All three models start
# from the same quantity, the **carry** of a futures product - the price difference
# between the contract expiring soonest and the one expiring after it, which is what a
# trader holding a position earns or pays each time the position is rolled from one to
# the next.
#
# 1. **ARIMA** forecasts next session's carry from the recent path of carry, one forecast
#    per product per session.
# 2. **A rolling Fourier transform** measures which cycle lengths the carry of a product
#    has been oscillating at over the past year.
# 3. **A two-state hidden Markov model** reads one number per session - carry averaged
#    across the whole book - and infers which of two market states the book is in.
#
# It reads the raw CME settlement prices and the forward-return labels written by
# [`02_labels`](02_labels.ipynb), and it writes one artifact,
# `features/model_based.parquet`.
#
# **What you will be able to do after reading this**
#
# - Say why estimating a model on all your data and then using its output as a feature
#   gives you a number no one could have computed at the time, and recognise the shape of
#   that mistake in your own code.
# - Refresh a model's parameters on a declared schedule, so that the value for any
#   session is produced by an estimate made from sessions strictly earlier than it -
#   and see why a walk-forward period does not do that job on its own.
# - Run a hidden Markov model so that its answer for a given day uses that day and every
#   earlier day but no later day, and check by experiment that this is what it did.
# - Write the resulting features to a file that records which prices they came from, so a
#   model trained on them later can state which version of the features it read.
#
# **Book Reference**: Chapter 9, Sections 9.3-9.5
#
# **Prerequisites**: [`02_labels`](02_labels.ipynb). It writes the forward-return file
# this notebook reads, and the dates in that file are what the training and evaluation
# periods below are cut from. [`03_financial_features`](03_financial_features.ipynb) runs
# as a parallel branch on the same raw prices; the two feature sets are read together by
# the model notebooks in Chapter 11.

# %%
"""CME Futures: Temporal Feature Engineering."""

import multiprocessing
import os
import re
import time
import warnings
from concurrent.futures import ProcessPoolExecutor
from datetime import date

# Pin the start method to fork before any pool-using import: Python 3.14 defaults to
# forkserver, which re-executes this script in every worker process the ARIMA walk spawns.
if multiprocessing.get_start_method(allow_none=True) is None:
    multiprocessing.set_start_method("fork")

import numpy as np
import pandas as pd
import plotly.graph_objects as go
import polars as pl
from hmmlearn.hmm import GaussianHMM
from plotly.subplots import make_subplots
from statsmodels.tsa.arima.model import ARIMA
from threadpoolctl import threadpool_limits

from case_studies.utils.artifact_digest import value_digest
from case_studies.utils.artifact_quality import (
    label_universe,
    quality_report,
    render_quality_report,
)
from case_studies.utils.temporal import (
    arima_one_step_forecast,
    filtered_state_probs,
    fit_hmm_kmeans_init,
    refit_boundaries,
    sort_states_by_mean,
    walk_forward_feature,
    write_model_based,
)
from case_studies.utils.warning_policy import apply_notebook_warning_policy
from data import load_cme_futures
from utils.artifact_specs import load_setup_config, resolve_label_buffer
from utils.cv_splits import generate_cv_splits, load_evaluation_config, select_folds
from utils.paths import get_case_study_dir
from utils.reproducibility import set_global_seeds
from utils.style import COLORS, show_plotly_with_alt

apply_notebook_warning_policy()

# %% [markdown]
# ## Configuration
#
# Five settings, and each one decides something a reader would otherwise have to guess at.
# `MAX_PRODUCTS` and `MAX_FOLDS` exist so a smoke test can run a fraction of the work; both
# are zero here, which means the full universe and every evaluation period.
#
# What is *not* here is the estimation schedule. How much history each fitted model spends
# before its first estimate, and how often it is refreshed, are part of what the feature
# means rather than settings to trade runtime against, so they are read from `setup.yaml`
# below alongside the feature windows. `REFIT_EVERY_OVERRIDE` is the reduction lever for
# that: zero here, meaning "use what `setup.yaml` declares", and a positive value replaces
# both declared cadences for one run. It is named so nothing reading this file can mistake
# a reduction for the definition.

# %% tags=["parameters"]
CASE_STUDY_ID = "cme_futures"
SEED = 42
# Number of products to model. Zero means all thirty; a positive value takes that many
# from the front of the list and is only for a fast check that the code runs.
MAX_PRODUCTS = 0
# Number of walk-forward evaluation periods to resolve. Zero means all of them. No model
# is fitted per period any more, so this narrows what section F screens over and what the
# coverage tables report; the fits cost the same either way.
MAX_FOLDS = 0
# How many past sessions each Fourier transform reads: 252, one trading year. A cycle
# can only be measured if the window is long enough to contain it more than once, so a
# year-long window is the shortest one from which a half-year cycle is legible.
FFT_WINDOW = 252
# The two cycle lengths whose strength is reported as a feature, in trading sessions:
# 63 is a quarter and 126 is half a year. Agricultural and energy contracts have
# seasonal supply and demand at both.
FFT_TARGET_PERIODS = [63, 126]
# The share of false positives tolerated among the features section F declares
# significant, after correcting for how many were tested at once.
FDR_ALPHA = 0.05
# 0 keeps both declared refit cadences. A positive value replaces them, which is how a
# smoke run bounds the two walks without narrowing the universe: fewer estimates, the same
# rows and the same columns. The burn-ins are never overridden - a shorter one would move
# which sessions carry a value, and the coverage assertions below are about exactly that.
REFIT_EVERY_OVERRIDE = 0

# %% [markdown]
# Three more settings come from `config/setup.yaml`, the file that also configures
# [`03_financial_features`](03_financial_features.ipynb). The universe is the thirty
# products and the sectors they belong to. The two windows are the ones that stage uses
# to smooth carry and to express it as a z-score - the number of standard deviations
# carry sits from its own recent average - so that the series built in section C is the
# same series that stage writes under the name `carry_zscore_63d`, in a different shape.

# %%
CASE_DIR = get_case_study_dir(CASE_STUDY_ID)
FEATURES_DIR = CASE_DIR / "features"
LABELS_DIR = CASE_DIR / "labels"
STRATEGY_ID = CASE_STUDY_ID
set_global_seeds(SEED)

SETUP = load_setup_config(CASE_STUDY_ID)
PRODUCT_GROUPS = SETUP["universe"]["product_groups"]
ALL_PRODUCTS = [p for products in PRODUCT_GROUPS.values() for p in products]
assert len(ALL_PRODUCTS) == SETUP["universe"]["n_products"], (
    f"setup.yaml declares {SETUP['universe']['n_products']} products, "
    f"product_groups lists {len(ALL_PRODUCTS)}"
)

CARRY_SMOOTHING = int(SETUP["features"]["windows"]["carry_smoothing"])
CARRY_ZSCORE_WINDOW = int(SETUP["features"]["windows"]["carry_zscore"][0])

# The estimation schedule, read rather than typed, so the comments in `setup.yaml` that
# say what each count decides stay next to the value the notebook uses.
MODEL_BASED = SETUP["model_based"]
ARIMA_BURNIN = int(MODEL_BASED["arima"]["burnin"])
ARIMA_REFIT_FREQ = int(MODEL_BASED["arima"]["refit_every"])
ARIMA_ORDER = tuple(int(v) for v in MODEL_BASED["arima"]["order"])
HMM_BURNIN = int(MODEL_BASED["hmm"]["burnin"])
HMM_REFIT_EVERY = int(MODEL_BASED["hmm"]["refit_every"])
HMM_N_STATES = int(MODEL_BASED["hmm"]["n_states"])
if REFIT_EVERY_OVERRIDE:
    ARIMA_REFIT_FREQ = HMM_REFIT_EVERY = REFIT_EVERY_OVERRIDE
    print(f"Reduced run: both refit cadences replaced with {REFIT_EVERY_OVERRIDE}")

# Two sessions carry one clearing venue's settlement file and not the other's; `setup.yaml`
# says which and why. They are dropped here so no series is differenced across a date on
# which half the universe has no settlement price.
EXCLUDED_SESSIONS = [
    date.fromisoformat(str(d)) for d in SETUP["universe"].get("excluded_sessions", [])
]

if MAX_PRODUCTS > 0:
    ARIMA_PRODUCTS = ALL_PRODUCTS[:MAX_PRODUCTS]
else:
    ARIMA_PRODUCTS = ALL_PRODUCTS

print(f"Carry is smoothed over {CARRY_SMOOTHING} sessions before anything reads it.")
print(
    f"Its z-score is taken against the previous {CARRY_ZSCORE_WINDOW} sessions of that "
    f"smoothed series."
)
print(f"Modelling {len(ARIMA_PRODUCTS)} of the {len(ALL_PRODUCTS)} products in the universe.")
print("Estimation schedule, in sessions of each model's own series:")
print(f"  ARIMA        burn-in {ARIMA_BURNIN:>4}, refit every {ARIMA_REFIT_FREQ:>3}")
print(f"  carry regime burn-in {HMM_BURNIN:>4}, refit every {HMM_REFIT_EVERY:>3}")

# %% [markdown]
# ## The data these models read
#
# One row per product, expiry and session, carrying that contract's settlement price.
# **Product** is a futures contract's underlying - corn, gold, the S&P 500 index - and
# each product trades in several contracts at once that differ only in when they expire.
# Those are indexed by `position`: position 0 is the contract expiring soonest, the
# **front month**; position 1 is the one after it; position 2 the one after that.

# %%
df = load_cme_futures(products=sorted(ALL_PRODUCTS)).rename(
    {"session_date": "timestamp", "tenor": "position"}
)
df = df.filter(~pl.col("timestamp").is_in(EXCLUDED_SESSIONS))

if MAX_PRODUCTS > 0:
    df = df.filter(pl.col("product").is_in(ARIMA_PRODUCTS))

print(f"Loaded {len(df):,} rows, {df['product'].n_unique()} products")
print(f"Date range: {df['timestamp'].min()} to {df['timestamp'].max()}")

# %% [markdown]
# The thirty products are not thirty interchangeable series. They are seven groups of
# things that move for their own reasons, and the models below fit one product at a time,
# so a product's group is what the reader should carry forward about it.
#
# Each row is one sector: the products in it, the session its earliest product first
# quoted, the session by which all of them were quoting, and how many front-month
# product-sessions it contributes. Two things to read off it.
#
# The panel starts together. Every sector's two date columns hold the same session, with
# one exception, and that exception is why the equity-index row contributes fewer sessions
# than any other four-product sector. Comparing the two date columns is how to find it.
#
# The rest of the spread in the session counts is holiday calendars. These sectors do not
# close on the same days, so per product the agricultural and livestock contracts quote
# about a hundred fewer sessions across the panel than the financial ones. That is a small
# effect for a model fitted one product at a time, and a large one for section C.3, which
# has to average carry across all of them on every session and therefore has to decide
# what to do about the ones that did not settle.

# %%
_front = df.filter(pl.col("position") == 0)
_sector_of = {p: sector for sector, products in PRODUCT_GROUPS.items() for p in products}
universe_table = (
    _front.with_columns(
        pl.col("product").replace_strict(_sector_of, default="unclassified").alias("sector")
    )
    .group_by(["sector", "product"])
    .agg(pl.col("timestamp").min().alias("product_start"), pl.len().alias("sessions"))
    .group_by("sector")
    .agg(
        pl.col("product").sort().str.join(" ").alias("products"),
        pl.col("product").n_unique().alias("n_products"),
        pl.col("product_start").min().alias("first_product_quoting"),
        pl.col("product_start").max().alias("all_products_quoting"),
        pl.col("sessions").sum().alias("front_month_sessions"),
    )
    .sort("front_month_sessions", descending=True)
)
universe_table

# %% [markdown]
# ## A. Why a feature built from a fitted model is a different hazard
#
# The features in [`03_financial_features`](03_financial_features.ipynb) are formulas. A
# 63-session average of carry on 3 March reads carry on the 63 sessions up to 3 March and
# nothing else, so whether it could have been computed at the time is settled by looking
# at the formula.
#
# The features in this notebook are not formulas. Each one is the output of a model whose
# **parameters were estimated from data**, and those parameters are part of what the
# feature knows. Suppose the hidden Markov model in section C is estimated once on the
# whole price history and then asked which state the market was in on 3 March 2016. Its
# answer depends on the two state means and the transition probabilities it settled on,
# and those were computed from every session in the file - including 2023. Nothing in the
# formula for 3 March mentions 2023. The dependence runs through the parameters instead,
# and it is invisible at the point where the number is used.
#
# That failure is worth naming precisely because it does not announce itself. The
# notebook runs without error, the feature looks reasonable, and it correlates with future
# returns better than it should - because it was partly built from them. A model trained
# on such a feature reports a performance the same strategy could never have earned, and
# the gap only appears when someone tries to trade it.
#
# The rule that removes it is one sentence: **no parameter behind the value for a session
# may have seen that session or any later one.** It has two halves, and both are enforced
# below.
#
# **Bound where the parameters come from.** Both fitted models here do it the same way:
# **re-estimate as the walk proceeds.** Spend a burn-in, fit, use that fit for the next
# `refit_every` sessions, then refit on everything up to that point. No session is ever
# used to estimate the model that speaks for it, whether or not a period boundary happens
# to sit nearby.
#
# **A walk-forward period does not do this job, and until 2026-09-04 the hidden Markov
# model in C.3 relied on it to.** Estimating once per period on that period's whole
# training window and then filtering forward from the *start* of that window is causal for
# the evaluation sessions and is not causal for the training ones: the earliest training
# rows of an eight-year window carried parameters estimated from eight years of their own
# future, while every evaluation row carried parameters estimated only from its past. A
# model downstream was then fitted on one version of the column and scored on another.
# Nothing raised, because a period's rows are internally consistent and the artifact
# recorded no estimation window. C.1's ARIMA was already on a refit schedule; C.3 now is
# too, and the schedule is what bounds an estimate rather than the period - which is why
# the file this notebook writes carries no period column at all.
#
# **Run the fitted model forward, never backward.** Even a model estimated on training
# data can look ahead when it is *applied*. A hidden Markov model can be asked two
# different questions about 3 March: what is the most likely state given everything up to
# 3 March, or given the whole series. The second question is the one the standard library
# call answers by default, and its answer for 3 March changes when data from April
# arrives. Only the first is a quantity that existed on 3 March. Section C runs the first
# and demonstrates the difference by deleting the later observations and checking the
# number does not move.
#
# ## B. The periods, and what they are for
#
# The walk-forward boundaries are resolved here, and what they bound is the **screen** in
# section F, not the fits. A screen run over the sessions an estimate read reports how well
# a feature fits history rather than whether it predicts, so section F is cut to the
# evaluation windows. What bounds a fit is the schedule above.
#
# The one boundary that does bind every fit is where the holdout opens: past it neither
# model re-estimates, and each carries the last estimate it made before the boundary
# across the window frozen. A coefficient refitted on holdout sessions is a parameter
# estimated on the holdout however causal the forecast around it looks.
#
# They are derived from the forward-return file rather than from the price file. The two
# do not span the same dates: a forward return needs a window after it to resolve, so the
# label file stops earlier than the prices. The model notebooks downstream cut their
# periods from the label file, so cutting from the same frame here is what makes a period
# number in this artifact mean the same thing on both sides of the join.

# %% [markdown]
# Three things are resolved here and used everywhere below.
#
# **Which forward return the case study is built around.** It is read from the
# configuration rather than typed, because the same choice has to pick three things at
# once: the file [`02_labels`](02_labels.ipynb) wrote, the gap left between each training
# and evaluation window, and the correlation lag in section F.
#
# **The gap between training and evaluation.** A decision made on the last training
# session is only settled `LABEL_HORIZON_SESSIONS` sessions later. If evaluation began the
# next session, the model would be scored on days whose outcome overlaps days it was
# trained on. So the two windows are held that far apart. The practice is called
# **purging**, and the gap is what `LABEL_BUFFER` sizes.
#
# **Where the holdout begins.** The last stretch of history is held back and not read by
# anything in the research process, so that there is one period left at the end on which
# the finished strategy can be run as if for the first time. Its first session is
# `HOLDOUT_START`.
#
# Section F needs a stricter boundary than that. It scores features against a forward
# return, and a decision on date `t` is settled `LABEL_HORIZON_SESSIONS` sessions after
# `t`. For that outcome to be observable outside the holdout, `t` itself has to fall that
# many sessions earlier than the holdout does - so the last date section F may score is
# `LAST_SCORABLE_DECISION_DATE`, counted on the sessions the exchange actually traded.

# %%
PRIMARY_LABEL = SETUP["labels"]["primary"]
LABEL_BUFFER = resolve_label_buffer(CASE_STUDY_ID, PRIMARY_LABEL, SETUP)
assert LABEL_BUFFER, f"No label buffer configured for {PRIMARY_LABEL}"
LABEL_HORIZON_SESSIONS = int(re.match(r"^(\d+)", LABEL_BUFFER).group(1))

label_frame = pl.read_parquet(CASE_DIR / "labels" / f"{PRIMARY_LABEL}.parquet")
splits = generate_cv_splits(
    label_frame.select("timestamp").unique().sort("timestamp"),
    case_study_id=CASE_STUDY_ID,
    label_buffer=LABEL_BUFFER,
)
if MAX_FOLDS > 0:
    splits = select_folds(splits, range(MAX_FOLDS))


def _as_date(value) -> date:
    return pd.Timestamp(value).date()


_evaluation_config = load_evaluation_config(CASE_STUDY_ID)
HOLDOUT_START = _as_date(_evaluation_config["holdout_start"])
HOLDOUT_END = _as_date(_evaluation_config["holdout_end"])
_sessions = df.select("timestamp").unique().sort("timestamp")["timestamp"].to_list()
_pre_holdout = [d for d in _sessions if d < HOLDOUT_START]
LAST_SCORABLE_DECISION_DATE = _pre_holdout[-(LABEL_HORIZON_SESSIONS + 1)]

print(
    f"Predicting {PRIMARY_LABEL}, so training and evaluation are held "
    f"{LABEL_HORIZON_SESSIONS} sessions apart."
)
print(f"{len(splits)} walk-forward periods, most recent first:")
for s in splits:
    print(
        f"  Period {s['fold']}: train {s['train_start']} → {s['train_end']}, "
        f"evaluate {s['val_start']} → {s['val_end']}"
    )
print(
    f"The holdout opens {HOLDOUT_START}. Section F scores no decision after "
    f"{LAST_SCORABLE_DECISION_DATE}, so every outcome it reads is settled before that."
)

# %% [markdown]
# The figure draws the five evaluation windows and the gap in front of each. Read it as a
# picture of what section F screens over, not of what bounds a fit - nothing here bounds a
# fit any more. The gap between each pair of bars is the purge, sized to the label horizon
# so that no training session's outcome reaches into the window a model is scored on, and
# no bar crosses into the shaded holdout.
#
# The estimation schedule is drawn separately, once the series each model reads has been
# built: ARIMA at the end of C.1 and the regime chain at the end of C.3, each against its
# own series, because their burn-ins are paid on series that begin at different dates.
#
# **The file this notebook writes carries rows dated inside the holdout, and that is
# deliberate.** A holdout evaluation downstream needs a feature value on those sessions.
# What must not reach into the holdout is an estimate, and neither model makes one there:
# both stop re-estimating at the last session before the boundary and carry that estimate
# across frozen. Section E prints, column by column, what is populated where.

# %%
fig = go.Figure()
_span_style = {
    "Training window": COLORS["blue"],
    "Evaluation window": COLORS["amber"],
}
_seen: set[str] = set()
for split in splits:
    row = f"Period {split['fold']}"
    for kind, (start, end) in (
        ("Training window", (split["train_start"], split["train_end"])),
        ("Evaluation window", (split["val_start"], split["val_end"])),
    ):
        fig.add_trace(
            go.Scatter(
                x=[pd.Timestamp(start).isoformat(), pd.Timestamp(end).isoformat()],
                y=[row, row],
                mode="lines",
                line={"width": 18, "color": _span_style[kind]},
                name=kind,
                legendgroup=kind,
                showlegend=kind not in _seen,
            )
        )
        _seen.add(kind)

fig.add_vrect(
    x0=pd.Timestamp(HOLDOUT_START).isoformat(),
    x1=pd.Timestamp(df["timestamp"].max()).isoformat(),
    fillcolor=COLORS["neutral"],
    opacity=0.10,
    line_width=0,
    layer="below",
)
fig.add_vline(
    x=pd.Timestamp(HOLDOUT_START).isoformat(), line_dash="dash", line_color=COLORS["negative"]
)
fig.update_layout(
    title=(
        "Each period trains, waits out the label horizon, then evaluates"
        "<br><sup>The gap between the bars is the purge. The dashed rule is where the "
        "holdout opens; the shaded region is held out."
        "<br>These windows bound the screen in section F, not the fits.</sup>"
    ),
    xaxis_title="Session",
    yaxis_title="",
    height=360,
    margin={"l": 90, "t": 110},
)
show_plotly_with_alt(
    fig,
    "Horizontal timeline with one row per period, periods 0 to 4, running from about 2012 to "
    "2025. Each row shows a long dark training window followed, after a visible gap, by a "
    "shorter amber evaluation window; the gap between them is the purge that waits out the label "
    "horizon. The windows step forward period by period. A dashed vertical rule marks where the "
    "holdout opens and a shaded band covers everything after it; no evaluation window reaches "
    "into that band.",
)

# %% [markdown]
# ## The input all three models read: carry
#
# Carry is how far the front-month contract settles above the next one along, as a
# fraction of the front price, scaled by twelve:
#
# $$c_{p,t} = 12 \times \frac{F^{(0)}_{p,t} - F^{(1)}_{p,t}}{F^{(0)}_{p,t}}$$
#
# where the superscript is the contract position. A trader holding the front month has to
# replace it with the next contract before it expires, and that gap is what the
# replacement earns or costs. It is positive in **backwardation**, where the nearer
# contract is the dearer one, and negative in **contango**, where it is the cheaper one.
#
# **The twelve is a scale factor, not an annual rate.** It would turn a one-month spread
# into a yearly one, and the four energy curves in this universe do list a contract every
# month. The other twenty-six are on quarterly or irregular cycles, so for them twelve is
# the wrong multiple for an annual rate and the number is not one. It is the same constant
# for every product on every date, so it changes no ranking and no z-score; what it does
# not deliver is a level that means the same thing on a Treasury curve as on a crude one.
#
# Two derived series come out of it. **Smoothed carry** is carry averaged over
# `CARRY_SMOOTHING` sessions, which removes the daily settlement noise the models would
# otherwise fit. **The carry z-score** expresses that smoothed level as the number of
# standard deviations it sits from its own average over the previous
# `CARRY_ZSCORE_WINDOW` sessions, so that gold in dollars and corn in cents can be
# compared on one scale.
#
# [`03_financial_features`](03_financial_features.ipynb) writes the same z-score under the
# name `carry_zscore_63d`, but with one row per contract rather than one per product. The
# models here need one series per product, so the same definition is recomputed in that
# shape from the raw prices - both windows read from the same configuration - rather than
# reshaped out of that file. This is why the file this notebook writes records the raw
# prices as its input and no other feature file.


# %%
def compute_carry(data: pl.DataFrame) -> pl.DataFrame:
    """Compute carry percentage from front and deferred month prices."""
    # Raw (unadjusted) close: the term-structure spread must read contemporaneous
    # tenor levels, not the ratio-adjusted series whose levels encode roll history.
    front = (
        data.filter(pl.col("position") == 0)
        .select(["product", "timestamp", "raw_close"])
        .rename({"raw_close": "c0_price"})
    )
    second = (
        data.filter(pl.col("position") == 1)
        .select(["product", "timestamp", "raw_close"])
        .rename({"raw_close": "c1_price"})
    )

    carry_df = front.join(second, on=["product", "timestamp"], how="inner")
    carry_df = carry_df.with_columns(
        ((pl.col("c0_price") - pl.col("c1_price")) / pl.col("c0_price") * 12).alias("carry_pct")
    )

    # Smoothed carry and z-score, on the windows setup.yaml declares
    carry_df = carry_df.sort(["product", "timestamp"])
    carry_df = carry_df.with_columns(
        pl.col("carry_pct")
        .rolling_mean(window_size=CARRY_SMOOTHING)
        .over("product")
        .alias("carry_smoothed"),
    )
    carry_df = carry_df.with_columns(
        (
            (
                pl.col("carry_smoothed")
                - pl.col("carry_smoothed").rolling_mean(CARRY_ZSCORE_WINDOW).over("product")
            )
            / pl.col("carry_smoothed")
            .rolling_std(CARRY_ZSCORE_WINDOW)
            .over("product")
            .clip(lower_bound=1e-6)
        )
        .clip(lower_bound=-5.0, upper_bound=5.0)
        .alias("carry_zscore")
    )

    return carry_df.select(
        ["product", "timestamp", "carry_pct", "carry_smoothed", "carry_zscore"]
    ).drop_nulls()


carry = compute_carry(df)
print(f"Carry data: {len(carry):,} product-dates")

# %% [markdown]
# ---
#
# ## C. The three models
#
# Each subsection below states what the model infers, where its parameters are allowed to
# come from, and ends with an assertion that runs - not a comment claiming the
# window held, but a check that fails the notebook if it did not.
#
# ### C.1 ARIMA: what carry does next
#
# Term structure changes gradually, so today's carry z-score carries information about
# tomorrow's. ARIMA is the standard model for that kind of series: it writes the next
# value as a weighted sum of recent values and of recent forecast errors, and estimates
# the weights. Its three orders say how many of each go in - `p` past values, `q` past
# errors, and `d` differences taken first if the series drifts rather than reverting.
#
# The orders are declared in `setup.yaml` rather than searched for at each refit, and the
# cell below sets out the measurement behind that. Only the weights are re-estimated on the
# schedule; `p`, `d` and `q` stay put, so the forecast means the same thing in every block.
# Seasonal terms play no part here, because the seasonality this case study cares about is
# measured directly in C.2.
#
# **Why every value it emits is a forecast and not a fit.** One call walks each product's
# whole history a session at a time: at each step the model sees the history up to that
# session and predicts the next one. What is emitted for a session is therefore the
# prediction made before the session happened, and the weights behind it are re-estimated
# every `ARIMA_REFIT_FREQ` steps. There are no fitted-in-place values anywhere in the
# output. The walk cannot begin until there is enough history to estimate from, so the
# first `ARIMA_BURNIN` sessions of each product's history get no value.
#
# The two features are the forecast itself, `arima_carry_forecast`, and what it missed,
# `arima_carry_residual` - the realised z-score minus the forecast. A large residual says
# carry moved in a way its own recent path did not imply.

# %% [markdown]
# Every product with enough history goes into one call, as a long frame keyed by product
# and date. The library walks those series together, spread across cores. It is the same
# walk-forward routine as
# [`10_uncertainty_features`](../../09_model_based_features/10_uncertainty_features.ipynb),
# and the two settings that govern it - the burn-in and the refit cadence - are the ones
# `setup.yaml` declares under `model_based.arima`.

# %%
_carry_ts_dtype = carry.schema["timestamp"]


# %% [markdown]
# The carry frame carries `pl.Date`, so a bound written as a Python date is cast to the
# column's own dtype; that is what makes an inclusive upper bound cover the last session.


# %%
def _date_lit(value) -> pl.Expr:
    """Cast a Python date or timestamp to the carry frame's timestamp dtype."""
    return pl.lit(pd.Timestamp(value).date()).cast(_carry_ts_dtype)


# %% [markdown]
# **One walk per product over the whole history, not one per period.**
#
# The periods do not bound this model. It re-estimates every `ARIMA_REFIT_FREQ` sessions on
# everything up to that point, so a forecast for a session is made by weights fitted only on
# earlier ones, whether or not a period boundary happens to sit nearby. That is the schedule
# section A describes, and the hidden Markov model in C.3 is on one too - so the forecasts
# are no longer replicated onto anything. One value per product and session goes into the
# file.
#
# Cutting the walk per period bought nothing and cost two things. The cross-validation call
# it used took one `n_windows` for every series and validated it against the shortest, so one
# call per period meant one walk length per period, and RTY, listed 2017-07-10, was the
# shortest eligible series in all five. Every other product was truncated to RTY's length,
# the forecasts landed at the end of each window, and every period's ARIMA began in 2018
# whatever its window was - including the period that opens in 2015. The last period's
# training rows ended up 99.87% empty of a feature its fits declared they were using. The
# second cost is quieter: the burn-in year was paid once per period rather than once.
#
# Carry is not thin early. It is 99.3% non-null across all thirty products back to 2011; the
# gap was the call shape.
#
# One walk per product is also **cheaper than what it replaces**, which is not the usual
# direction for a correctness fix. The five periods overlap - the most recent spans 2015 to
# 2023, the oldest 2011 to 2019 - so the per-period design forecast the same product-dates up
# to five times and kept one. On this panel: 86,204 forecasts against 123,870.

# %% [markdown]
# ### The order is declared, and it used to be searched for
#
# Until this notebook was standardized the order was chosen automatically at every refit, by
# a stepwise search that picked whichever `(p, d, q)` scored best on the history available at
# that point. It is now read from `setup.yaml`, and since that is a change to the model
# rather than to the code around it, here is the measurement behind it.
#
# Sampling eight refit cutoffs on each of the 30 eligible products - 240 order searches - the
# automatic selection returned **34 distinct orders**, and **no product held a single order
# across its own walk**. The most common, `(2,0,1)`, took 11.7% of the searches and `(2,0,2)`
# 10.4%; the steadiest product spent 6 of its 8 cutoffs on one order and the rest moved more
# than that.
#
# An order that changes almost every month, on an expanding window of the same series, is the
# information criterion tracking the sample rather than structure being found. It also makes
# `arima_carry_forecast` a different quantity in every block, which is the property that
# breaks a comparison across chapters: two products' forecasts, or the same product's in two
# periods, were not made by the same model.
#
# `(2,0,1)` is the modal selection, and the more parsimonious of the two orders that are
# indistinguishable from each other. The differencing is not a search result at all: carry is
# already a rolling z-score clipped to plus or minus five, so it is stationary before this
# model reads it, and `d=0` follows from how the input is built. 199 of the 240 searches
# agreed; the 41 that differenced it were over-differencing a bounded series.


# %%
def _arima_fit(train: np.ndarray):
    """Estimate the coefficients on one block, at the order `setup.yaml` declares."""
    with warnings.catch_warnings():
        # Convergence chatter on a short block is expected and the walk's own burn-in is
        # what handles it; a fit that genuinely fails raises and stops the walk.
        warnings.simplefilter("ignore")
        return ARIMA(train[:, 0], order=ARIMA_ORDER).fit()


def _arima_apply(fitted, prefix: np.ndarray) -> np.ndarray:
    """One-step forecasts across a prefix, under the coefficients that block estimated."""
    return arima_one_step_forecast(fitted, prefix).reshape(-1, 1)


def _arima_one_product(payload: tuple[str, np.ndarray, np.ndarray, int]) -> pl.DataFrame:
    """Walk one product. Module level and picklable, because it runs in a worker process."""
    product, values, dates, frozen_after = payload
    forecast = walk_forward_feature(
        values.reshape(-1, 1),
        timestamps=dates,
        burnin=ARIMA_BURNIN,
        refit_every=ARIMA_REFIT_FREQ,
        fit=_arima_fit,
        apply=_arima_apply,
        n_features=1,
        freeze_after=frozen_after,
    )[:, 0]
    return pl.DataFrame(
        {
            "timestamp": dates,
            "product": [product] * len(dates),
            "arima_carry_forecast": forecast,
            "arima_carry_residual": values - forecast,
        }
    )


# %% [markdown]
# One walk, where there used to be two. The old shape cut its input at `HOLDOUT_START` and
# then ran a second walk across the holdout with refitting switched off, because the first
# emitted nothing inside the holdout and a holdout evaluation downstream needs a value on
# every one of those sessions.
#
# `freeze_after` is that distinction expressed once. The walk runs over the whole series,
# holdout sessions included, and past the last pre-holdout session it stops re-estimating and
# keeps applying what it last fitted. What must not reach into the holdout is an *estimate* -
# a coefficient refitted on holdout sessions is a parameter estimated on the holdout however
# causal the forecast around it looks - and freezing is what prevents that, rather than a cut
# on the input. Each forecast still conditions only on carry strictly before its own date,
# which is the property `arima_one_step_forecast` carries and section A states.


# %%
def _arima_walk() -> pl.DataFrame:
    """One-step walk-forward ARIMA forecasts per product, over each product's whole history."""
    full = (
        carry.filter(pl.col("product").is_in(ARIMA_PRODUCTS))
        .drop_nulls(subset=["carry_zscore"])
        .sort(["product", "timestamp"])
    )
    empty = pl.DataFrame(
        schema={
            "timestamp": pl.Date,
            "product": pl.String,
            "arima_carry_forecast": pl.Float64,
            "arima_carry_residual": pl.Float64,
        }
    )

    # Eligibility is measured on the pre-holdout history, because that is what the estimates
    # are allowed to come from: a product whose carry only starts inside the holdout has
    # nothing to fit on, whatever its total length.
    development = full.filter(pl.col("timestamp") < _date_lit(HOLDOUT_START))
    lengths = development.group_by("product").len().sort("len")
    required = ARIMA_BURNIN + 30
    eligible = lengths.filter(pl.col("len") >= required)
    excluded = lengths.filter(pl.col("len") < required)
    if excluded.height:
        # Named, not counted. A product missing from this feature changes what it covers, and
        # until 2026-08-23 the exclusion happened silently.
        listed = ", ".join(f"{row[0]} ({row[1]})" for row in excluded.iter_rows())
        print(f"  excluded, under {required} carry sessions before the holdout: {listed}")
    if not eligible.height:
        print("  no eligible products")
        return empty

    payloads = []
    for product in sorted(eligible["product"].to_list()):
        series = full.filter(pl.col("product") == product)
        dates = series["timestamp"].to_numpy()
        values = series["carry_zscore"].to_numpy()
        # `_date_lit` builds an expression, which against a Series yields another expression
        # rather than a mask; the count itself is what the walk freezes on.
        frozen_after = int(series.filter(pl.col("timestamp") < _date_lit(HOLDOUT_START)).height)
        payloads.append((product, values, dates, frozen_after))

    # One call per product, spread across processes. The fits are per product and independent,
    # and running them sequentially was still going after 49 minutes where the whole notebook
    # used to take 19. A fork context is named rather than left to the default, because Python
    # 3.14 defaults to forkserver, which re-imports the parent module and cannot reach a
    # function defined in a notebook kernel.
    workers = max(1, min(len(payloads), (os.cpu_count() or 2) - 1))
    print(f"  fitting {len(payloads)} products across {workers} processes", flush=True)
    with ProcessPoolExecutor(
        max_workers=workers, mp_context=multiprocessing.get_context("fork")
    ) as pool:
        frames = list(pool.map(_arima_one_product, payloads))

    walked = pl.concat(frames).sort(["product", "timestamp"])
    emitted = walked.drop_nulls(subset=["arima_carry_forecast"])
    print(
        f"  {len(payloads)} products fitted, {len(emitted):,} forecasts across "
        f"{len(walked):,} product-sessions",
        flush=True,
    )
    return walked


arima_t0 = time.time()
arima_pl = _arima_walk()
if arima_pl.height and arima_pl["timestamp"].dtype != pl.Date:
    arima_pl = arima_pl.with_columns(pl.col("timestamp").cast(pl.Date))
arima_elapsed = time.time() - arima_t0

# %%
if arima_pl.height:
    print(
        f"\nARIMA total: {len(arima_pl):,} rows across "
        f"{arima_pl['product'].n_unique()} products in {arima_elapsed:.0f}s"
    )
else:
    print("No ARIMA results generated")

# %% [markdown]
# **Check what the emitted rows are dated, and what each period actually receives.**
#
# The date checks are about the boundary rather than about periods. Rows dated before the
# holdout come from the walk; rows dated inside it come from the frozen tail and have to
# exist, because a holdout evaluation reads them - an empty holdout was the defect this
# section was rewritten to fix, and a check that only forbade holdout-dated values would
# have been satisfied perfectly by emitting none.
#
# What none of this establishes is where the weights came from, and no assertion over the
# output frame could, because the weights are not in the frame. What bounds them is the
# shape of the call - every refit reads a prefix ending before the sessions it goes on to
# forecast - together with the holdout cut being applied to the input in `_arima_walk`.
#
# Then the coverage, **in both windows of every period**. Until 2026-08-23 this reported the
# evaluation window only, and a model trains on the other one - so the number printed here
# was healthy on every run while the oldest period's training rows were 99.87% empty of the
# same feature. A coverage diagnostic that reads the window the model does not fit on is not
# evidence about the fit. `rules/notebook-standards.md` C16 now requires both, and this cell
# is what motivated the clause.

# %%
if len(arima_pl) > 0:
    _key = arima_pl.select(pl.struct("product", "timestamp").is_duplicated().sum()).item()
    assert _key == 0, f"{_key} duplicate (product, timestamp) rows in the ARIMA frame"

    _holdout_window = arima_pl.filter(pl.col("timestamp") >= _date_lit(HOLDOUT_START))
    assert _holdout_window.height > 0, (
        "nothing was emitted inside the holdout window, so a holdout retrain would fit "
        "this column on nulls"
    )
    assert _holdout_window["arima_carry_forecast"].null_count() == 0, (
        "a holdout-dated row carries no forecast"
    )
    assert _holdout_window["timestamp"].max() <= HOLDOUT_END, (
        "a row was emitted past the end of the holdout window"
    )
    _pre = arima_pl.filter(pl.col("timestamp") < _date_lit(HOLDOUT_START))
    print(
        f"One walk over {arima_pl['product'].n_unique()} products: "
        f"{_pre.drop_nulls('arima_carry_forecast').height:,} forecasts before the holdout "
        f"opens {HOLDOUT_START}, and {_holdout_window.height:,} inside it from "
        f"{_holdout_window['timestamp'].min()} to {_holdout_window['timestamp'].max()}, on "
        f"coefficients estimated before the boundary."
    )
    print("Product-sessions carrying an ARIMA value, per period, in each window:")
    for split in splits:
        for window, start_date, end_date in (
            ("train", split["train_start"], split["train_end"]),
            ("valid", split["val_start"], split["val_end"]),
        ):
            _in = (pl.col("timestamp") >= _as_date(start_date)) & (
                pl.col("timestamp") <= _as_date(end_date)
            )
            quoted = carry.filter(_in).height
            covered = arima_pl.filter(_in).drop_nulls("arima_carry_forecast").height
            print(
                f"  period {split['fold']} {window}: {covered:>7,} of {quoted:>7,} quoted "
                f"({100 * covered / max(quoted, 1):5.1f}%)"
            )

# %% [markdown]
# **The schedule this walk actually ran, drawn against the series it read.** The grey
# stretch is the burn-in each product spends before its first forecast; the blue stretch is
# where the order and weights are re-chosen every `ARIMA_REFIT_FREQ` sessions; the amber
# stretch past the rule is the holdout, over which the last pre-boundary estimate is carried
# frozen. Products list at different dates, so the burn-in ends at a different session for
# each one and the rows are not aligned - the row for the shortest series is the one to read
# against the evaluation windows in section B.

# %%
if arima_pl.height:
    _arima_schedule = (
        carry.filter(pl.col("product").is_in(arima_pl["product"].unique().to_list()))
        .filter(pl.col("timestamp") < _date_lit(HOLDOUT_START))
        .drop_nulls(subset=["carry_zscore"])
        .group_by("product")
        .agg(
            pl.len().alias("pre_holdout_sessions"),
            pl.col("timestamp").min().alias("first_session"),
        )
        .with_columns(
            (pl.col("pre_holdout_sessions") - ARIMA_BURNIN).alias("forecasts"),
            (
                (pl.col("pre_holdout_sessions") - ARIMA_BURNIN + ARIMA_REFIT_FREQ - 1)
                // ARIMA_REFIT_FREQ
            ).alias("estimates"),
        )
        .sort("pre_holdout_sessions")
    )
    print(
        f"ARIMA burn-in {ARIMA_BURNIN}, refit every {ARIMA_REFIT_FREQ}: "
        f"{_arima_schedule['estimates'].sum():,} estimates over "
        f"{_arima_schedule['forecasts'].sum():,} forecasts, "
        f"{_arima_schedule['estimates'].min()} to {_arima_schedule['estimates'].max()} "
        f"per product."
    )
    _arima_schedule.head(5)

# %% [markdown]
# ---
#
# ### C.2 A rolling Fourier transform: which cycles carry is running at
#
# Crops are harvested at the same time each year and heating demand peaks each winter, so
# the cost of holding a corn or a natural gas position is not the same in every month.
# Carry inherits that rhythm. A Fourier transform is the tool for finding it: it rewrites
# a stretch of a series as a sum of waves of different lengths and reports how much of the
# series' movement each wave accounts for. That amount is conventionally called the
# **power** at that wave's length.
#
# Five numbers per product per session come out of the transform of the previous
# `FFT_WINDOW` sessions:
#
# - `fft_dominant_period` - the length, in sessions, of the wave with the most power.
#   Near 252 it says the product is running on an annual cycle; near 21 it says the
#   movement is monthly and probably not seasonal at all.
# - `fft_energy_63d` and `fft_energy_126d` - the share of total power sitting at the two
#   cycle lengths declared in `FFT_TARGET_PERIODS`, quarterly and half-yearly.
# - `fft_spectral_entropy` - how spread the power is across wave lengths. Low entropy
#   means one cycle dominates and the series is close to periodic; high entropy means the
#   power is scattered and no cycle stands out, which is what noise looks like.
# - `fft_spectral_energy` - the total, which is a measure of how much the series moved at
#   all over the window and puts the three shares in context.


# %% [markdown]
# The transform of one window. The window's own average is subtracted first, because the
# transform reports the flat part of a series - the wave of infinite length - as the
# largest component of all, and that says only that carry is negative on average, which
# is not a cycle. That component is dropped from every summary for the same reason.


# %%
def _fft_window_features(segment: np.ndarray, target_periods: list[int]) -> dict[str, float]:
    centered = segment - segment.mean()
    fft_vals = np.fft.rfft(centered)
    power = np.abs(fft_vals) ** 2
    freqs = np.fft.rfftfreq(len(segment))
    total_power = np.sum(power[1:])

    output = {
        "total_power": float(total_power),
        "dominant_period": float("nan"),
        "spectral_entropy": float("nan"),
    }
    for period in target_periods:
        output[f"energy_{period}d"] = float("nan")

    if len(power) <= 1 or total_power <= 0:
        return output

    dom_idx = np.argmax(power[1:]) + 1
    if freqs[dom_idx] > 0:
        output["dominant_period"] = float(1.0 / freqs[dom_idx])

    p_norm = power[1:] / total_power
    p_norm = p_norm[p_norm > 0]
    output["spectral_entropy"] = float(-np.sum(p_norm * np.log(p_norm)))

    for period in target_periods:
        target_freq = 1.0 / period
        freq_idx = np.argmin(np.abs(freqs - target_freq))
        low_idx = max(1, freq_idx - 1)
        high_idx = min(len(power), freq_idx + 2)
        output[f"energy_{period}d"] = float(np.sum(power[low_idx:high_idx]) / total_power)
    return output


# %% [markdown]
# The window slides one session at a time and each result is written at the index the
# window ends *before*, so a value at `t` never reads the observation at `t`.


# %%
def rolling_fft_features(
    signal: np.ndarray,
    window: int = 252,
    target_periods: list[int] | None = None,
) -> dict[str, np.ndarray]:
    if target_periods is None:
        target_periods = [63, 126]

    n = len(signal)
    spectral_energy = np.full(n, np.nan)
    dominant_period = np.full(n, np.nan)
    spectral_entropy = np.full(n, np.nan)
    freq_energies = {p: np.full(n, np.nan) for p in target_periods}

    for t in range(window, n):
        window_stats = _fft_window_features(signal[t - window : t], target_periods)
        spectral_energy[t] = window_stats["total_power"]
        dominant_period[t] = window_stats["dominant_period"]
        spectral_entropy[t] = window_stats["spectral_entropy"]
        for period in target_periods:
            freq_energies[period][t] = window_stats[f"energy_{period}d"]

    result = {
        "fft_spectral_energy": spectral_energy,
        "fft_dominant_period": dominant_period,
        "fft_spectral_entropy": spectral_entropy,
    }
    for period, energy in freq_energies.items():
        result[f"fft_energy_{period}d"] = energy
    return result


# %% [markdown]
# One product at a time, over its whole history. This transform is the exception in
# section C: it estimates nothing. ARIMA fits weights and the model in C.3 fits state
# means and transition probabilities, so both are bound by a refit schedule; the transform
# of a window is a fixed calculation on the numbers in it, with no parameters at all. That
# makes it safe to run once over the full history, on the same footing as a rolling
# average, and the window is backward-looking, so no session's value reads a later one.

# %%
fft_results = []

for product in ARIMA_PRODUCTS:
    prod_carry = (
        carry.filter(pl.col("product") == product)
        .sort("timestamp")
        .drop_nulls(subset=["carry_pct"])
    )
    if len(prod_carry) < FFT_WINDOW + 50:
        continue

    signal = prod_carry["carry_pct"].to_numpy()
    dates = prod_carry["timestamp"].to_list()

    fft_out = rolling_fft_features(signal, window=FFT_WINDOW, target_periods=FFT_TARGET_PERIODS)

    prod_df = pl.DataFrame({"timestamp": dates, "product": product, **fft_out})
    fft_results.append(prod_df)
    valid_count = prod_df.drop_nulls(subset=["fft_spectral_energy"]).height
    print(f"  {product}: {valid_count} valid FFT observations")

# %% [markdown]
# One value per product and session, and that is the whole frame. Until 2026-09-04 these
# values were then copied once per period, because the period number was part of the key
# the models downstream joined on and every feature had to carry it. Nothing was estimated,
# so the copies were identical; they existed to satisfy a key that no longer has a period
# in it.

# %%
if fft_results:
    fft_pl = pl.concat(fft_results)
    fft_base = fft_pl
    print(f"\nSpectral features computed on {len(fft_pl):,} distinct product-sessions")
else:
    fft_pl = pl.DataFrame(
        schema={
            "timestamp": pl.Date,
            "product": pl.String,
            "fft_spectral_energy": pl.Float64,
            "fft_dominant_period": pl.Float64,
            "fft_spectral_entropy": pl.Float64,
            "fft_energy_63d": pl.Float64,
            "fft_energy_126d": pl.Float64,
        }
    )
    fft_base = fft_pl
    print("No FFT results generated")

# %% [markdown]
# **Check the window looks backward.** With no parameters, the only way this transform
# could read the future is through the window itself - an off-by-one in the slice would
# be enough. Recomputation is what settles it rather than re-reading the code: delete
# every observation after date `t`, transform what is left, and the value at `t` has to
# come back identical.

# %%
if fft_results:
    probe_product = fft_base["product"][0]
    probe_signal = (
        carry.filter(pl.col("product") == probe_product)
        .sort("timestamp")
        .drop_nulls(subset=["carry_pct"])["carry_pct"]
        .to_numpy()
    )
    probe_t = FFT_WINDOW + 100
    full_pass = rolling_fft_features(
        probe_signal, window=FFT_WINDOW, target_periods=FFT_TARGET_PERIODS
    )
    truncated = rolling_fft_features(
        probe_signal[: probe_t + 1], window=FFT_WINDOW, target_periods=FFT_TARGET_PERIODS
    )
    for key in full_pass:
        assert np.isclose(full_pass[key][probe_t], truncated[key][probe_t], equal_nan=True), key
    print(
        f"Recomputation agrees: for {probe_product} at session {probe_t}, deleting the "
        f"{len(probe_signal) - probe_t - 1} observations that come after it leaves every "
        f"one of its spectral values unchanged."
    )

# %% [markdown]
# ---
#
# ### C.3 A hidden Markov model: which of two states the book is in
#
# The first two models look at one product at a time. This one looks at the whole book:
# its input is a single number per session, carry averaged across the thirty products.
#
# A **hidden Markov model** assumes the series was generated by a system that is in one
# of a small number of states at any moment, that each state produces observations with
# its own average and spread, and that the system switches between states with fixed
# probabilities. The states are hidden because they are never observed directly - only
# the numbers they produce are - and fitting the model means estimating, from the
# observations alone, what those averages, spreads and switching probabilities are.
#
# Two states are used here, and they correspond to the two shapes the term structure
# takes. In one, the front contract settles above the next one, so rolling a long
# position forward earns the difference; that is **backwardation**, and carry is
# positive. In the other, the next contract is the dearer one, so the same roll pays the
# difference; that is **contango**, and carry is negative.
#
# Two features come out: `hmm_carry_regime_prob`, the probability the book is in the
# higher-carry state, and `hmm_regime_duration`, how many consecutive sessions the more
# probable state has held.
#
# **Two things have to be got right, and they are different things.** The parameters are
# re-estimated on a schedule, each estimate reading only sessions earlier than the ones it
# then speaks for. And the state probabilities are obtained by running the model *forward*
# - the answer for a session uses that session and every earlier one, and nothing later.
# The library's own `predict_proba` answers a different question, conditioning on the
# entire series, and its answer for a given session changes when data from months
# afterwards arrives. That quantity did not exist at the time and cannot be a feature.
# Both are checked by assertion below.
#
# Three pieces of machinery are shared with the other case studies that fit a hidden
# Markov model, in `case_studies/utils/temporal.py`: the fit that starts EM from a
# k-means partition, the ordering rule below, and the forward recursion. The recursion
# in particular reaches into a private part of `hmmlearn`, which is a thing to write
# once and document once rather than to copy into every notebook that needs it.

# %% [markdown]
# **Fitting the same numbers twice.** The estimation runs on one thread. The seed fixes
# which random draw is taken, not the order the arithmetic happens in: k-means adds up
# its distances in parallel, floating-point addition is not associative, so a
# multi-threaded fit lands on starting means that differ in their last bits, and EM
# carries that difference into the transition probabilities. Pinned to one thread, two
# runs of this notebook produce the same feature values - which is what the content
# fingerprint written in section E is a statement about.

# %% [markdown]
# **Giving the two states a stable identity.** EM returns them in whatever order it
# converged to, so without a rule the same fitted state can come back as state 0 for one
# estimate and state 1 for the next, and a feature named after one of them would mean
# different things along its own length. The rule has to be the quantity the feature name
# claims: `hmm_carry_regime_prob` is the probability of the *higher-carry* state, so the
# states are ordered on their estimated average carry, lower first. Section D draws the two
# averages across estimates, which is where that ordering can be checked.

# %% [markdown]
# #### Building the one number per session the model reads
#
# Carry averaged across the universe sounds simple and is not. Which products go into the
# average has to be the same from one session to the next, or the number moves when the
# set of contributors changes rather than when carry does. The sectors on this exchange
# keep different holiday calendars: a session that closes the metals pits leaves the
# grains settling as usual, and an average taken over whatever happened to settle jumps
# for a reason that has nothing to do with the term structure.
#
# So a product that does not settle keeps the carry of its last settlement for
# `HOLD_LAST_SETTLE_SESSIONS` sessions, carried forward only and never backward. A product
# absent for longer than that, or not yet trading at all, is left out of that session's
# average rather than represented by a stale number.
#
# The hold covers part of the problem and the cell below measures which part: how many
# absences there are, how many last the single closed session the holiday explanation
# predicts, how long the longest one runs, and what share of the missing product-sessions
# a two-session hold fills. What the hold does not reach shows up as a smaller set of
# contributors, and the per-session count of them is printed under it.
#
# That measurement is what sets `HOLD_LAST_SETTLE_SESSIONS`, so it is taken over
# pre-holdout sessions only. A constant chosen by looking at the holdout is a parameter
# estimated on the holdout, whatever the code that consumes it does afterwards. The
# observation series the models read stops at the same boundary, for the same reason.

# %%
HOLD_LAST_SETTLE_SESSIONS = 2  # sessions a last settlement stands in for

pre_holdout_carry = carry.filter(pl.col("timestamp") < _date_lit(HOLDOUT_START))
_carry_sessions = (
    pre_holdout_carry.select("timestamp").unique().sort("timestamp")["timestamp"].to_list()
)
_session_index = {d: i for i, d in enumerate(_carry_sessions)}
_absence_runs = []
for (_product,), _product_rows in pre_holdout_carry.group_by("product"):
    _seen = np.sort(np.array([_session_index[d] for d in _product_rows["timestamp"].to_list()]))
    _gaps = np.diff(_seen) - 1
    _absence_runs.extend(int(g) for g in _gaps[_gaps > 0])
_absence_runs = np.array(_absence_runs)
_missing_cells = int(_absence_runs.sum())
_held_cells = int(np.minimum(_absence_runs, HOLD_LAST_SETTLE_SESSIONS).sum())

print(
    f"A product goes missing mid-history {len(_absence_runs):,} times, over "
    f"{_missing_cells:,} product-sessions of "
    f"{len(_carry_sessions) * pre_holdout_carry['product'].n_unique():,}."
)
print(
    f"  gone for one session: {(_absence_runs == 1).sum():,}   "
    f"two: {(_absence_runs == 2).sum():,}   "
    f"longer: {(_absence_runs > 2).sum():,}   longest: {_absence_runs.max()} sessions"
)
print(
    f"Holding the last settlement for {HOLD_LAST_SETTLE_SESSIONS} sessions covers "
    f"{_held_cells:,} of the {_missing_cells:,} missing product-sessions "
    f"({100 * _held_cells / _missing_cells:.0f}%)."
)

# %% [markdown]
# The average itself: every product on every session, the hold applied forward, and the
# mean over whatever is present.

# %%
_basket_grid = (
    pre_holdout_carry.select("timestamp")
    .unique()
    .join(pre_holdout_carry.select("product").unique(), how="cross")
)
held_carry = (
    _basket_grid.join(pre_holdout_carry, on=["product", "timestamp"]