Signature Methods (Part 4 — Rolling Signatures in Volatility Forecasts)

# SPY settings
SPY_TICKER = "SPY"
SPY_START = pd.Timestamp("2000-01-01", tz="UTC")
SPY_END = pd.Timestamp("2023-12-31", tz="UTC")
SPY_BDAYS = pd.bdate_range(SPY_START, SPY_END, tz="UTC")
SPY_START_BDAY = SPY_BDAYS.min()
SPY_END_BDAY = SPY_BDAYS.max()

# Fixed split indices
SPY_IS_INDEX = 200  # In-sample start
SPY_OS_INDEX = 4000  # Out-of-sample start

# Feature configuration
N_LAGS = 0  # Only use current day's return
INIT_WINDOW = 500  # Initialization window for STES models

# XGBoost parameters
XGB_PARAMS = {
    "booster": "gblinear",
    "updater": "coord_descent",
    "max_depth": 5,
    "learning_rate": 0.01,
    "subsample": 1.0,
    "colsample_bytree": 1.0,
    "min_child_weight": 1.0,
    "reg_lambda": 1.0,
    "reg_alpha": 0.0,
    "verbosity": 0,
}
XGB_ROUNDS = 200
SCALE_FACTOR = 100.0  # Scale returns to volatility points

print(f"Data range: {SPY_START_BDAY.date()} to {SPY_END_BDAY.date()}")
print(f"In-sample starts at index: {SPY_IS_INDEX}")
print(f"Out-of-sample starts at index: {SPY_OS_INDEX}")

@dataclass(frozen=True)
class SignatureConfig:
    """Configuration for signature feature generation."""
    name: str
    lookback: int  # Number of days to look back
    level: int  # Signature truncation level
    augmentations: list[str]  # Augmentation pipeline
    sig_tfm: str  # 'signature' or 'logsignature'

# Example configuration
example_sig = SignatureConfig(
    name="L20_L2",
    lookback=20,
    level=2,
    augmentations=["cumsum", "basepoint", "addtime", "leadlag"],
    sig_tfm="signature"
)

print(f"Example: {example_sig.name}")
print(f"  Lookback: {example_sig.lookback} days")
print(f"  Level: {example_sig.level}")
print(f"  Augmentations: {' -> '.join(example_sig.augmentations)}")

# — from volatility_forecast/features/signature_features.py —

class SignatureFeaturesTemplate(_PriceTemplate):

    def transform(self, ctx, params, slice, state):
        lags      = int(params["lags"])          # rolling window length
        sig_level = int(params["sig_level"])      # signature truncation depth

        # 1) Build rolling path: one fixed-length window per timestamp
        logret = self._logret(panel, price_col)   # log-returns series
        path_df = pd.concat(
            {f"lag_{i}": logret.groupby(level="entity_id").shift(i)
             for i in range(lags, -1, -1)},
            axis=1,
        ).dropna()

        # 2) Resolve augmentation pipeline (order matters)
        augmentation_list = self._resolve_augmentation_pipeline(
            params.get("augmentation_list", "all")
        )

        # 3) Instantiate sktime SignatureTransformer
        sig_transformer = SignatureTransformer(
            augmentation_list=augmentation_list,   # e.g. ("cumsum","basepoint","addtime","leadlag")
            depth=sig_level,                       # truncation level
            window_name="global",                  # one window = one pre-built path
            rescaling="post",                      # factorial rescaling
            sig_tfm="signature",                   # full (not log-) signature
            backend="esig",
        )

        # 4) Reshape to sktime Panel format and compute
        paths   = path_df.to_numpy(dtype=float)    # (n_samples, n_timepoints)
        X_panel = paths[:, None, :]                 # (n_samples, 1_channel, n_timepoints)

        sig_df = sig_transformer.fit_transform(X_panel)
        sig_df.index = path_df.index                # realign (entity, time) index
        # → sig_df columns are the signature coefficients used as features

def _base_feature_requests(lags: int) -> tuple[FeatureRequest, ...]:
    """Define baseline features: log returns, absolute returns, squared returns."""
    return (
        FeatureRequest(
            template=LagLogReturnTemplate(),
            params={
                "lags": lags,
                "source": "tiingo",
                "table": "market.ohlcv",
                "price_col": "close",
            },
        ),
        FeatureRequest(
            template=LagAbsLogReturnTemplate(),
            params={
                "lags": lags,
                "source": "tiingo",
                "table": "market.ohlcv",
                "price_col": "close",
            },
        ),
        FeatureRequest(
            template=LagSquaredLogReturnTemplate(),
            params={
                "lags": lags,
                "source": "tiingo",
                "table": "market.ohlcv",
                "price_col": "close",
            },
        ),
    )

def build_spy_spec_baseline(lags: int = N_LAGS) -> VolDatasetSpec:
    """Baseline spec identical to Part 2 (no signature features)."""
    features = _base_feature_requests(lags)

    target = TargetRequest(
        template=NextDaySquaredReturnTarget(),
        params={
            "source": "tiingo",
            "table": "market.ohlcv",
            "price_col": "close",
            "scale": 1.0,
        },
        horizon=1,
        name="y",
    )

    return VolDatasetSpec(
        universe=UniverseSpec(entities=[SPY_TICKER]),
        time=TimeSpec(
            start=SPY_START_BDAY,
            end=SPY_END_BDAY,
            calendar="XNYS",
            grid="B",
            asof=None,
        ),
        features=features,
        target=target,
        join_policy=JoinPolicy(how="inner", sort_index=True),
        missingness=MissingnessPolicy(final_row_policy="drop_if_any_nan"),
    )

def build_spy_spec_with_signature(sig: SignatureConfig, lags: int = N_LAGS):
    """Spec that adds signature features on top of the baseline features."""
    base = list(_base_feature_requests(lags))

    # Add signature feature template
    sig_req = FeatureRequest(
        template=SignatureFeaturesTemplate(),
        params={
            "lags": int(sig.lookback),
            "sig_level": int(sig.level),
            "sig_tfm": str(sig.sig_tfm),
            "augmentation_list": (
                "none" if not sig.augmentations else "->".join(sig.augmentations)
            ),
            "source": "tiingo",
            "table": "market.ohlcv",
            "price_col": "close",
        },
    )

    features = tuple(base + [sig_req])

    target = TargetRequest(
        template=NextDaySquaredReturnTarget(),
        params={
            "source": "tiingo",
            "table": "market.ohlcv",
            "price_col": "close",
            "scale": 1.0,
        },
        horizon=1,
        name="y",
    )

    return VolDatasetSpec(
        universe=UniverseSpec(entities=[SPY_TICKER]),
        time=TimeSpec(
            start=SPY_START_BDAY,
            end=SPY_END_BDAY,
            calendar="XNYS",
            grid="B",
            asof=None,
        ),
        features=features,
        target=target,
        join_policy=JoinPolicy(how="inner", sort_index=True),
        missingness=MissingnessPolicy(final_row_policy="drop_if_any_nan"),
    )

def build_spy_dataset_with_signature(
    ctx,
    *,
    sig: SignatureConfig,
    baseline_X: pd.DataFrame,
    baseline_index: pd.Index,
    baseline_columns: list[str],
    lags: int = N_LAGS,
) -> pd.DataFrame:
    """Build signature dataset and align to baseline index.
    
    Key steps:
    1. Build dataset with signature features (drops early rows)
    2. Reindex to baseline index (restores dropped rows)
    3. Restore baseline columns from baseline dataset
    4. Fill signature columns with 0.0 for warmup rows
    """
    spec = build_spy_spec_with_signature(sig, lags)
    X, _y, _returns, _catalog = build_vol_dataset(ctx, spec, persist=False)

    X1 = _split_entity_frame(X, SPY_TICKER)
    if isinstance(X1, pd.Series):
        X1 = X1.to_frame()

    if "const" not in X1.columns:
        X1["const"] = 1.0

    # Align to baseline index
    X_aligned = X1.reindex(baseline_index)

    # Restore baseline columns
    X_aligned.loc[baseline_index, baseline_columns] = baseline_X.loc[
        baseline_index, baseline_columns
    ].to_numpy()

    # Fill signature features for warmup rows
    sig_cols = [c for c in X_aligned.columns if c not in set(baseline_columns)]
    if sig_cols:
        X_aligned[sig_cols] = X_aligned[sig_cols].fillna(0.0)

    return X_aligned

subset = sig_cols[:50]
corr = X_sig_only[subset].corr().replace([np.inf, -np.inf], np.nan).fillna(0.0)

dist = 1.0 - np.abs(corr.values)
np.fill_diagonal(dist, 0.0)
dist[np.isclose(dist, 0.0)] = 0.0

Z = linkage(squareform(dist, checks=False), method="average")
order = leaves_list(Z)
corr_reordered = corr.iloc[order, order]

plt.figure(figsize=(10, 9))
sns.heatmap(
    corr_reordered,
    cmap="RdBu_r",
    center=0,
    square=True,
    cbar_kws={"shrink": 0.8},
)
plt.title("Clustered correlation matrix (ordered by hierarchical clustering)")
plt.tight_layout()
plt.show()

def select_features_by_cluster(
    X: pd.DataFrame,
    y: pd.Series,
    feature_cols: list[str],
    train_slice: slice,
    corr_threshold: float = 0.95,
    max_features: int | None = None,
) -> tuple[list[str], dict]:
    """Select one feature per correlation cluster, fit on training data only.

    Parameters
    ----------
    X : DataFrame with all features
    y : target series (aligned to X index)
    feature_cols : columns to cluster and select from
    train_slice : training indices (causal — clustering uses train only)
    corr_threshold : features with |corr| >= this are grouped together.
        This is the *starting* threshold — if ``max_features`` is set and
        the initial clustering yields too many clusters, the threshold is
        progressively lowered (i.e. clusters are merged more aggressively)
        until the feature count is within budget.
    max_features : int or None
        Hard cap on the number of selected features.  When set, the
        correlation threshold is lowered in steps of 0.05 until
        ``n_clusters <= max_features``.

    Returns
    -------
    selected : list of str
        Selected feature column names (one per cluster).
    cluster_map : dict
        Mapping from cluster id to list of member feature names.
    """
    X_train = X.iloc[train_slice][feature_cols]
    y_train = y.iloc[train_slice]

    # 0. Drop zero-variance features (e.g. warmup zeros that never vary in training)
    variances = X_train.var()
    live_cols = [c for c in feature_cols if variances[c] > 0]
    dropped_cols = [c for c in feature_cols if variances[c] <= 0]
    if dropped_cols:
        print(f"  Dropped {len(dropped_cols)} zero-variance features in training window")

    X_train = X_train[live_cols]

    # 1. Correlation matrix on training data
    #    NaN correlations (e.g. near-constant columns) treated as perfectly correlated
    #    so they get grouped together and only one survives
    corr_mat = X_train.corr().replace([np.inf, -np.inf], np.nan)
    corr_mat = corr_mat.fillna(1.0)

    # 2. Distance = 1 - |corr|
    dist = 1.0 - np.abs(corr_mat.values)
    np.fill_diagonal(dist, 0.0)
    dist = np.clip(dist, 0.0, None)  # numerical safety

    # 3. Hierarchical clustering — adapt threshold to meet max_features budget
    Z = linkage(squareform(dist, checks=False), method="average")

    threshold = corr_threshold
    while True:
        cut_distance = 1.0 - threshold
        labels = fcluster(Z, t=cut_distance, criterion="distance")
        n_clusters = len(set(labels))
        if max_features is None or n_clusters <= max_features:
            break
        # Lower threshold (merge more aggressively) and retry
        threshold -= 0.05
        if threshold < 0.30:
            # Safety floor — don't merge features with |corr| < 0.30
            break

    if threshold < corr_threshold:
        print(f"  Lowered corr_threshold {corr_threshold:.2f} → {threshold:.2f} "
              f"to meet max_features={max_features} (got {n_clusters} clusters)")

    # 4. For each cluster, pick feature with highest |corr| with target
    target_corr = X_train.corrwith(y_train).abs().fillna(0.0)
    selected = []
    cluster_map = {}  # cluster_id -> list of features
    for col, cid in zip(live_cols, labels):
        cluster_map.setdefault(cid, []).append(col)

    for cid, members in sorted(cluster_map.items()):
        best = max(members, key=lambda c: target_corr.get(c, 0.0))
        selected.append(best)

    return selected, cluster_map

train_sl = slice(SPY_IS_INDEX, SPY_OS_INDEX)

selected_sig_cols, cluster_map = select_features_by_cluster(
    X_sig, y_base, sig_cols, train_slice=train_sl, corr_threshold=0.95,
)

print(f"Original signature features: {len(sig_cols)}")
print(f"Clusters found: {len(cluster_map)}")
print(f"Selected features (one per cluster): {len(selected_sig_cols)}")
print()

for cid, members in sorted(cluster_map.items(), key=lambda x: -len(x[1])):
    chosen = [m for m in members if m in selected_sig_cols][0]
    if len(members) > 1:
        dropped = [m for m in members if m != chosen]
        print(f"  Cluster {cid} ({len(members)} features): keep {chosen}, drop {dropped}")
    else:
        print(f"  Cluster {cid} (1 feature): keep {chosen}")

fig, axes = plt.subplots(1, 2, figsize=(18, 8))

# Left: all signature features
corr_all = X_sig_only.corr().replace([np.inf, -np.inf], np.nan).fillna(0.0)
dist_all = 1.0 - np.abs(corr_all.values)
np.fill_diagonal(dist_all, 0.0)
dist_all = np.clip(dist_all, 0.0, None)
Z_all = linkage(squareform(dist_all, checks=False), method="average")
order_all = leaves_list(Z_all)
corr_all_reordered = corr_all.iloc[order_all, order_all]

sns.heatmap(corr_all_reordered, cmap="RdBu_r", center=0, square=True,
            cbar_kws={"shrink": 0.7}, ax=axes[0],
            xticklabels=False, yticklabels=False)
axes[0].set_title(f"All signature features ({len(sig_cols)})")

# Right: selected features only
X_selected = X_sig_only[selected_sig_cols]
corr_sel = X_selected.corr().replace([np.inf, -np.inf], np.nan).fillna(0.0)
if len(selected_sig_cols) > 2:
    dist_sel = 1.0 - np.abs(corr_sel.values)
    np.fill_diagonal(dist_sel, 0.0)
    dist_sel = np.clip(dist_sel, 0.0, None)
    Z_sel = linkage(squareform(dist_sel, checks=False), method="average")
    order_sel = leaves_list(Z_sel)
    corr_sel_reordered = corr_sel.iloc[order_sel, order_sel]
else:
    corr_sel_reordered = corr_sel

sns.heatmap(corr_sel_reordered, cmap="RdBu_r", center=0, square=True,
            cbar_kws={"shrink": 0.7}, ax=axes[1],
            xticklabels=False, yticklabels=False)
axes[1].set_title(f"After cluster selection ({len(selected_sig_cols)} features)")

plt.suptitle("Correlation structure: before vs after cluster-based feature selection", fontsize=13)
plt.tight_layout()
plt.show()

"""Taken from volatility_forecast/example/volatility_forecast_2.py"""

class ModelName(str, Enum):
    ES = "ES"
    STES_EAESE = "STES_EAESE"
    XGBSTES_E2E = "XGBSTES_E2E"

def _make_xgb_stes_model(
    *,
    seed: int | None,
    fit_method: str = "end_to_end",
    loss: str = "mse",
    xgb_params: dict | None = None,
    monotonic_constraints: dict[str, int] | None = None,
) -> XGBoostSTESModel:
    """Create XGBoost STES model with specified configuration."""
    params = dict(XGB_PARAMS)
    if xgb_params:
        params.update(xgb_params)
    
    # Strip gblinear-specific params when using tree booster
    if params.get("booster", "gbtree") != "gblinear":
        params.pop("updater", None)
    
    e2e_kwargs = {}
    if fit_method == "end_to_end":
        e2e_kwargs = {
            "e2e_grad_hess_scale": 1.0,
            "e2e_debug": True,
            "e2e_debug_print_once": True,
        }

    return XGBoostSTESModel(
        xgb_params=params,
        num_boost_round=XGB_ROUNDS,
        init_window=INIT_WINDOW,
        fit_method=fit_method,
        loss=loss,
        random_state=seed,
        monotonic_constraints=monotonic_constraints,
        **e2e_kwargs,
    )

def _infer_monotone_constraints(cols: list[str]) -> dict[str, int]:
    """Enforce -1 monotonicity on shock-magnitude features."""
    out: dict[str, int] = {}
    for c in cols:
        cl = c.lower()
        if any(k in cl for k in ["abs", "squared", "sq", "r2", "vol", "rv"]):
            out[c] = -1
    return out

def _xgb_variant_overrides(variant: str, cols: list[str]) -> dict:
    """Translate variant names into model configuration."""
    v = variant.upper()
    fit_method = "end_to_end" if "E2E" in v or "END_TO_END" in v else "alternating"
    loss = "pseudohuber" if "HUBER" in v else "mse"
    mono = _infer_monotone_constraints(cols) if "MONO" in v else None
    return {"fit_method": fit_method, "loss": loss, "monotonic_constraints": mono}

def _is_signature_col(col: str) -> bool:
    return col.startswith("sig.") or col.startswith("mcsig.")

def _scale_train_test(
    X: pd.DataFrame, train_slice: slice, test_slice: slice
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Standardize all features (including signatures) on train, apply to test.

    Matches volatility_forecast_2.py: every column except 'const' is scaled.
    """
    X_tr = X.iloc[train_slice].copy()
    X_te = X.iloc[test_slice].copy()

    cols_to_scale = [c for c in X.columns if c != "const"]

    if not cols_to_scale:
        return X_tr, X_te

    scaler = StandardScaler().fit(X_tr[cols_to_scale])
    X_tr.loc[:, cols_to_scale] = pd.DataFrame(
        scaler.transform(X_tr[cols_to_scale]),
        index=X_tr.index,
        columns=cols_to_scale,
    )
    X_te.loc[:, cols_to_scale] = pd.DataFrame(
        scaler.transform(X_te[cols_to_scale]),
        index=X_te.index,
        columns=cols_to_scale,
    )

    return X_tr, X_te

"""Taken from volatility_forecast/example/volatility_forecast_2.py"""
def _fit_predict_oos(
    *,
    variant: str,
    X: pd.DataFrame,
    y: pd.Series,
    r: pd.Series,
    train_slice: slice,
    test_slice: slice,
    seed: int | None = None,
    xgb_params: dict | None = None,
    extra_cols: list[str] | None = None,
    return_alpha: bool = False,
    cv_splits: int = 5,
) -> tuple[pd.Index, np.ndarray] | tuple[pd.Index, np.ndarray, pd.Series]:
    """Fit model on train slice and predict on test slice with warm-start.
    
    For XGBoost-STES variants, supports three fitting modes via suffix:
      - No suffix or _COLD: cold-start (default)
      - _WS: warm-start — initialize from baseline z-scores via base_margin
      - _2S: two-stage — freeze baseline, train only on signature features
    
    Tree booster mode: include _TREE in the variant name to use gbtree
    instead of gblinear.  Warm-start / two-stage are not supported with
    gbtree (the tree builder is greedy, so does not suffer from gblinear's
    non-convexity issues).
    """
    from scipy.special import expit as _expit
    
    variant_name = variant.value if isinstance(variant, Enum) else str(variant)
    
    # Detect warm-start / two-stage mode from variant suffix
    fit_mode = "cold"
    clean_variant = variant_name
    if variant_name.endswith("_WS"):
        fit_mode = "warm_start"
        clean_variant = variant_name[:-3]  # strip _WS
    elif variant_name.endswith("_2S"):
        fit_mode = "two_stage"
        clean_variant = variant_name[:-3]  # strip _2S
    
     # Detect tree booster from variant name
    is_tree = "_TREE" in clean_variant.upper()
    
    if is_tree and fit_mode != "cold":
        logger.warning(
            f"[{variant_name}] Warm-start/two-stage not supported for "
            f"gbtree — falling back to cold-start"
        )
        fit_mode = "cold"
    
    # For tree variants, inject gbtree-specific XGBoost parameters
    if is_tree:
        tree_defaults = {
            "booster": "gbtree",
            "max_depth": 2,
            "eta": 0.05,
            "subsample": 1.0,       # keep 1.0 for temporal STES
            "colsample_bytree": 1.0,
            "min_child_weight": 5.0,
            "reg_lambda": 1.0,
            "reg_alpha": 0.0,
            "verbosity": 0,
        }
        # User's xgb_params override tree defaults
        if xgb_params:
            tree_defaults.update(xgb_params)
        xgb_params = tree_defaults
    
    # Prepare data
    if clean_variant.startswith("XGBSTES_"):
        # Drop 'const' — gblinear already has an unpenalized bias term;
        # keeping const creates a redundant, partly-penalized intercept.
        if extra_cols:
            # Use baseline (non-sig) columns + ONLY the specified extra columns
            base_only = [c for c in X.columns if not _is_signature_col(c)]
            cols = base_only + [c for c in extra_cols if c in X.columns and c not in base_only]
        else:
            cols = [c for c in X.columns if not _is_signature_col(c)]
        cols = [c for c in cols if c != "const"]
        X_sel = X[cols]
        X_tr, X_te = _scale_train_test(X_sel, train_slice, test_slice)
        y_tr, r_tr = y.iloc[train_slice], r.iloc[train_slice]
        r_te = r.iloc[test_slice]
        
        over = _xgb_variant_overrides(clean_variant, cols)
        model = _make_xgb_stes_model(seed=None, xgb_params=xgb_params, **over)
    else:
        cols = select_variant_columns(X, clean_variant) or ["const"]
        cols = [c for c in cols if not _is_signature_col(c)]
        if extra_cols:
            cols = cols + [c for c in extra_cols if c not in set(cols)]
        
        X_sel = X[cols]
        X_tr, X_te = _scale_train_test(X_sel, train_slice, test_slice)
        y_tr, r_tr = y.iloc[train_slice], r.iloc[train_slice]
        r_te = r.iloc[test_slice]
        
        if clean_variant == ModelName.ES.value:
            model = ESModel(random_state=seed) if seed is not None else ESModel()
        else:
            model = STESModel(random_state=seed) if seed is not None else STESModel()
    
    # Fit model
    if clean_variant.startswith("XGBSTES_"):
        r_tr_scaled = r_tr * SCALE_FACTOR
        y_tr_scaled = y_tr * (SCALE_FACTOR**2)
        
        # CV grid: tree models need different hyperparameters
        if is_tree:
            grid_e2e = [
                {"max_depth": 2, "min_child_weight": 5.0, "eta": 0.05},
                {"max_depth": 3, "min_child_weight": 10.0, "eta": 0.03},
                {"max_depth": 2, "min_child_weight": 20.0, "eta": 0.1},
            ]
        else:
            grid_e2e = [
                {"min_child_weight": 5.0, "learning_rate": 0.05, "max_depth": 3},
                {"min_child_weight": 20.0, "learning_rate": 0.1, "max_depth": 3},
                {"min_child_weight": 50.0, "learning_rate": 0.05, "max_depth": 2},
            ]
        
        # ── WARM-START / TWO-STAGE: fit baseline first ──────────────────
        base_margin_train = None
        base_margin_all = None
        X_tr_fit = X_tr      # features used for fitting
        X_all_fit = None      # will be set later for prediction
        
        if fit_mode in ("warm_start", "two_stage") and extra_cols:
            # 1. Fit baseline model (without signature features)
            base_cols = [c for c in cols if not _is_signature_col(c)]
            X_tr_base_ws, X_te_base_ws = _scale_train_test(X[base_cols], train_slice, test_slice)
            
            model_base = _make_xgb_stes_model(seed=None, xgb_params=xgb_params, **over)
            model_base.fit(
                X_tr_base_ws, y_tr_scaled, returns=r_tr_scaled,
                perform_cv=True, cv_grid=grid_e2e, cv_splits=cv_splits,
                start_index=0, end_index=len(X_tr_base_ws),
            )
            
            # 2. Extract baseline z-scores: z = bias + Σ w_j * x_j
            dump_str = model_base.model_.get_dump()[0]
            b_bl, w_bl = _parse_gblinear_dump(dump_str)
            
            # Map baseline weights to the full (augmented) feature space
            w_embedded = np.zeros(len(cols))
            for j, c in enumerate(cols):
                if c in base_cols:
                    base_j = base_cols.index(c)
                    w_embedded[j] = w_bl[base_j]
            
            # Compute z_baseline on scaled augmented features
            X_tr_np = X_tr.values.astype(float)
            z_baseline_train = b_bl + X_tr_np @ w_embedded
            
            # Also for test (needed at prediction time)
            X_te_np = X_te.values.astype(float)
            z_baseline_test = b_bl + X_te_np @ w_embedded
            
            base_margin_train = z_baseline_train
            base_margin_all = np.concatenate([z_baseline_train, z_baseline_test])
            
            if fit_mode == "two_stage":
                # Only train on signature features; baseline is frozen in base_margin
                sig_fit_cols = [c for c in cols if _is_signature_col(c)]
                X_tr_fit = X_tr[sig_fit_cols]
                # Create separate model for sig-only features
                model = _make_xgb_stes_model(seed=None, xgb_params=xgb_params, **over)
            
            logger.info(
                f"[{fit_mode}] Baseline fitted (bias={b_bl:.4f}). "
                f"base_margin range: [{z_baseline_train.min():.3f}, {z_baseline_train.max():.3f}]"
            )
        
        model.fit(
            X_tr_fit, y_tr_scaled, returns=r_tr_scaled,
            perform_cv=True, cv_grid=grid_e2e, cv_splits=cv_splits,
            start_index=0, end_index=len(X_tr_fit),
            base_margin=base_margin_train,
        )
    elif extra_cols and isinstance(model, STESModel):
        # Group L2: penalize only signature features, leave baseline free.
        # Uniform L2 shrinks useful baseline weights when trying to suppress
        # noisy signature weights — group-specific penalties avoid this.
        from sklearn.model_selection import TimeSeriesSplit
        from scipy.optimize import least_squares as _lsq

        X_tr_np = X_tr.values.astype(float)
        y_tr_np = y_tr.values.astype(float)
        r_tr_np = r_tr.values.astype(float)
        p = X_tr_np.shape[1]
        col_list = list(X_tr.columns)

        # Mask: 1.0 for signature features, 0.0 for baseline / const
        sig_mask = np.array(
            [1.0 if _is_signature_col(c) else 0.0 for c in col_list]
        )

        l2_sig_grid = [0.0, 0.01, 0.1, 1.0, 10.0]
        tscv = TimeSeriesSplit(n_splits=cv_splits)
        best_l2, best_mse = 0.0, np.inf

        for l2_cand in l2_sig_grid:
            fold_mses = []
            for tr_idx, va_idx in tscv.split(X_tr_np):
                Xf_tr, Xf_va = X_tr_np[tr_idx], X_tr_np[va_idx]
                yf_tr, yf_va = y_tr_np[tr_idx], y_tr_np[va_idx]
                rf_tr, rf_va = r_tr_np[tr_idx], r_tr_np[va_idx]

                scale_tr = np.linalg.norm(yf_tr)
                p_vec = (
                    np.sqrt(l2_cand) * scale_tr * sig_mask
                    if l2_cand > 0
                    else None
                )

                rng = np.random.default_rng(seed)
                x0 = rng.normal(0, 1, size=p)
                res = _lsq(
                    STESModel._objective,
                    x0=x0,
                    jac=STESModel._jacobian,
                    args=(rf_tr, Xf_tr, yf_tr, 0, len(Xf_tr), p_vec),
                )
                beta = res.x

                alphas_va = _expit(Xf_va @ beta)
                r2_va = rf_va ** 2
                v = (
                    float(np.mean(rf_tr[-20:] ** 2))
                    if len(rf_tr) >= 20
                    else float(r2_va[0])
                )
                yhat_va = np.empty(len(yf_va))
                for t in range(len(yf_va)):
                    v_next = alphas_va[t] * r2_va[t] + (1 - alphas_va[t]) * v
                    yhat_va[t] = v_next
                    v = v_next
                fold_mses.append(float(np.mean((yf_va - yhat_va) ** 2)))

            avg_mse = float(np.mean(fold_mses))
            if avg_mse < best_mse:
                best_mse, best_l2 = avg_mse, l2_cand

        logger.info(
            f"Group-L2 CV selected λ_sig={best_l2:.4f} (MSE={best_mse:.2e})"
        )

        # Final fit with best group penalty
        scale = np.linalg.norm(y_tr_np)
        p_vec_final = (
            np.sqrt(best_l2) * scale * sig_mask if best_l2 > 0 else None
        )
        rng = np.random.default_rng(seed)
        x0 = rng.normal(0, 1, size=p)
        res = _lsq(
            STESModel._objective,
            x0=x0,
            jac=STESModel._jacobian,
            args=(r_tr_np, X_tr_np, y_tr_np, 0, len(X_tr_np), p_vec_final),
        )
        model.params = res.x
        model._set_schema(col_list, p)
    else:
        model.fit(X_tr, y_tr, returns=r_tr, start_index=0, end_index=len(X_tr))
    
    # Warm-start prediction: concatenate train+test to avoid cold start
    n_tr = len(X_tr)
    X_all = pd.concat([X_tr, X_te], axis=0)
    r_all = pd.concat([r_tr, r_te], axis=0)
    
    if clean_variant.startswith("XGBSTES_"):
        r_all_scaled = r_all * SCALE_FACTOR
        
        # For two-stage prediction, use sig-only features
        if fit_mode == "two_stage" and extra_cols:
            sig_fit_cols = [c for c in cols if _is_signature_col(c)]
            X_all_pred = pd.concat([X_tr[sig_fit_cols], X_te[sig_fit_cols]], axis=0)
        else:
            X_all_pred = X_all
        
        y_hat_all_scaled = model.predict(
            X_all_pred, returns=r_all_scaled,
            base_margin=base_margin_all,
        )
        y_hat_all = np.asarray(y_hat_all_scaled, dtype=float) / (SCALE_FACTOR**2)
        alpha_all = model.get_alphas(X_all_pred, base_margin=base_margin_all)
    else:
        y_hat_all = model.predict(X_all, returns=r_all)
        if isinstance(model, STESModel):
            _yhat_vals, alpha_vals = model.predict_with_alpha(X_all, returns=r_all)
            alpha_all = pd.Series(alpha_vals, index=X_all.index, name="alpha")
        else:
            alpha_val = getattr(model, "alpha_", np.nan)
            alpha_all = pd.Series(alpha_val, index=X_all.index, name="alpha")
    
    # Extract test predictions
    y_hat_arr = np.asarray(y_hat_all, dtype=float)[n_tr:]
    keep = np.isfinite(y_hat_arr)
    
    if return_alpha:
        return X_te.index[keep], y_hat_arr[keep], alpha_all
    return X_te.index[keep], y_hat_arr[keep]

"""Taken from volatility_forecast/example/volatility_forecast_2.py"""
def evaluate_fixed_split(
    X: pd.DataFrame,
    y: pd.Series,
    r: pd.Series,
    *,
    variants: list[str],
    seeds: list[int],
    is_index: int = SPY_IS_INDEX,
    os_index: int = SPY_OS_INDEX,
    extra_cols_by_variant: dict[str, list[str]] | None = None,
    cv_splits: int = 5,
) -> pd.DataFrame:
    """Evaluate variants across multiple random seeds."""
    train_sl = slice(is_index, os_index)
    test_sl = slice(os_index, len(y))

    rows = []
    for variant in variants:
        rmses, maes, medaes = [], [], []
        extra_cols = (extra_cols_by_variant or {}).get(variant)

        for seed in seeds:
            try:
                idx_te, y_hat = _fit_predict_oos(
                    variant=variant,
                    X=X,
                    y=y,
                    r=r,
                    train_slice=train_sl,
                    test_slice=test_sl,
                    seed=seed,
                    extra_cols=extra_cols,
                    cv_splits=cv_splits,
                )
                y_true = y.loc[idx_te].values
                rmses.append(rmse(y_true, y_hat))
                maes.append(mae(y_true, y_hat))
                medaes.append(medae(y_true, y_hat))
            except Exception as e:
                logger.exception(
                    f"Fixed-split eval failed: variant={variant}, seed={seed}: {e}"
                )

        rows.append(
            {
                "variant": variant,
                "rmse": float(np.mean(rmses)) if rmses else np.nan,
                "mae": float(np.mean(maes)) if maes else np.nan,
                "medae": float(np.mean(medaes)) if medaes else np.nan,
            }
        )

    return pd.DataFrame(rows).sort_values("rmse")

train_slice = slice(SPY_IS_INDEX, SPY_OS_INDEX)
test_slice  = slice(SPY_OS_INDEX, len(y_base))

alpha_variants = {
    "STES (no sig)":   dict(variant="STES_EAESE",   extra_cols=None),
    "STES + sig":      dict(variant=stes_sig_name,   extra_cols=selected_sig_cols),
    "XGB-STES (no sig)": dict(variant="XGBSTES_E2E", extra_cols=None),
    "XGB-STES + sig":    dict(variant=xgb_sig_name,   extra_cols=selected_sig_cols),
}

alphas = {}
for label, kw in alpha_variants.items():
    _, _, alpha = _fit_predict_oos(
        X=X_sig_selected, y=y_base, r=r_base,
        train_slice=train_slice, test_slice=test_slice,
        seed=0, return_alpha=True, **kw,
    )
    alphas[label] = alpha.iloc[SPY_OS_INDEX - SPY_IS_INDEX :]

fig, axes = plt.subplots(2, 1, figsize=(14, 7), sharex=True)

for label in ["STES (no sig)", "STES + sig"]:
    axes[0].plot(alphas[label].index, alphas[label].values, label=label, linewidth=0.7)
axes[0].set_ylabel(r"$\alpha_t$")
axes[0].set_title("STES — predicted smoothing parameter (OOS)")
axes[0].legend()
axes[0].grid(True, alpha=0.3)

for label in ["XGB-STES (no sig)", "XGB-STES + sig"]:
    axes[1].plot(alphas[label].index, alphas[label].values, label=label, linewidth=0.7)
axes[1].set_ylabel(r"$\alpha_t$")
axes[1].set_title("XGBoost-STES — predicted smoothing parameter (OOS)")
axes[1].legend()
axes[1].grid(True, alpha=0.3)

fig.tight_layout()
plt.show()

# ── Quantify the alpha level shift visible in the chart ──────────────────
print("OOS SMOOTHING PARAMETER (α) — LEVEL SHIFT ANALYSIS")
print("=" * 72)
print(f"\n{'Model':>25s}  {'mean(α)':>10s}  {'std(α)':>10s}  {'median':>10s}  {'IQR':>12s}")
print("-" * 72)
for label in ["STES (no sig)", "STES + sig", "XGB-STES (no sig)", "XGB-STES + sig"]:
    a = alphas[label].values
    q25, q75 = np.percentile(a, [25, 75])
    print(f"{label:>25s}  {a.mean():>10.4f}  {a.std():>10.4f}  {np.median(a):>10.4f}  [{q25:.4f}, {q75:.4f}]")

# ── Level shift magnitudes ──────────────────────────────────────────────
a_stes_base = alphas["STES (no sig)"].values
a_stes_sig  = alphas["STES + sig"].values
a_xgb_base  = alphas["XGB-STES (no sig)"].values
a_xgb_sig   = alphas["XGB-STES + sig"].values

print(f"\nLevel shifts:")
print(f"  STES:      mean α jumps from {a_stes_base.mean():.4f} → {a_stes_sig.mean():.4f}  (Δ = {a_stes_sig.mean() - a_stes_base.mean():+.4f}, {(a_stes_sig.mean()/a_stes_base.mean()-1)*100:+.1f}%)")
print(f"  XGB-STES:  mean α jumps from {a_xgb_base.mean():.4f} → {a_xgb_sig.mean():.4f}  (Δ = {a_xgb_sig.mean() - a_xgb_base.mean():+.4f}, {(a_xgb_sig.mean()/a_xgb_base.mean()-1)*100:+.1f}%)")

# ── Correlation between baseline and sig alpha paths ─────────────────────
corr_stes = np.corrcoef(a_stes_base, a_stes_sig)[0, 1]
corr_xgb = np.corrcoef(a_xgb_base, a_xgb_sig)[0, 1]
print(f"\nCorrelation between baseline α and sig α:")
print(f"  STES:     ρ = {corr_stes:.4f}")
print(f"  XGB-STES: ρ = {corr_xgb:.4f}")

# ── What does the level shift imply for the ES recursion? ────────────────
# α controls the weight on new shocks:  σ²_{t+1} = α·r²_t + (1-α)·σ²_t
# Higher α → faster reaction to shocks, less smoothing, more volatile forecasts
# Lower α → more smoothing, slower reaction, stabler forecasts
mean_shift_stes = a_stes_sig.mean() - a_stes_base.mean()
mean_shift_xgb = a_xgb_sig.mean() - a_xgb_base.mean()

# Implied half-life of a shock:  h = log(0.5) / log(1-α)
hl_stes_base = np.log(0.5) / np.log(1 - a_stes_base.mean())
hl_stes_sig = np.log(0.5) / np.log(1 - a_stes_sig.mean())
hl_xgb_base = np.log(0.5) / np.log(1 - a_xgb_base.mean())
hl_xgb_sig = np.log(0.5) / np.log(1 - a_xgb_sig.mean())

print(f"\nImplied shock half-life (days):")
print(f"  STES:     {hl_stes_base:.1f} days → {hl_stes_sig:.1f} days")
print(f"  XGB-STES: {hl_xgb_base:.1f} days → {hl_xgb_sig:.1f} days")
print(f"\n  The signature model reacts faster to shocks (shorter half-life).")
print(f"  This systematically changes the forecast character.")

from sktime.transformations.panel.signature_based import SignatureTransformer


def build_multichannel_rolling_signatures(
    returns: pd.Series,
    lookback: int = 20,
    level: int = 2,
    augmentations: tuple[str, ...] = ("basepoint", "addtime", "leadlag"),
    sig_tfm: str = "signature",
    rescaling: str = "post",
) -> pd.DataFrame:
    """Build signature features from a 4-channel return path over rolling windows.

    Channels:
        0: cumsum(r_t)        — log-price path (rebased per window)
        1: |r_t|              — absolute return
        2: r_t^2              — squared return
        3: r_t * 1_{r_t < 0}  — negative return (leverage proxy)

    Channel 0 is computed *within each window* so the path is rebased to 0.
    This removes drift/level effects that would otherwise leak non-stationarity
    into the signature features.

    Parameters
    ----------
    returns : pd.Series with DatetimeIndex
    lookback : rolling window length
    level : signature truncation depth
    augmentations : augmentation pipeline (no cumsum — handled manually)
    sig_tfm : 'signature' or 'logsignature'
    rescaling : 'post', 'pre', or None

    Returns
    -------
    pd.DataFrame with signature features, indexed by date.
    First valid date is returns.index[lookback - 1].
    """
    r = returns.values.astype(float)
    n = len(r)

    # Roll into panel: one instance per window
    n_windows = n - lookback + 1
    # sktime numpy3D convention: (N, C, L) = (instances, channels, time)
    panel = np.zeros((n_windows, 4, lookback))
    for i in range(n_windows):
        window = r[i : i + lookback]
        ch_logprice = np.cumsum(window)  # rebased log-price path
        ch_abs_r = np.abs(window)
        ch_r2 = window ** 2
        ch_r_neg = window * (window < 0).astype(float)
        panel[i] = np.vstack([ch_logprice, ch_abs_r, ch_r2, ch_r_neg])

    # Compute signatures
    sig_transformer = SignatureTransformer(
        augmentation_list=augmentations,
        depth=level,
        window_name="global",
        rescaling=rescaling,
        sig_tfm=sig_tfm,
        backend="esig",
    )
    sig_features = sig_transformer.fit_transform(panel)

    # Align to dates: window [i : i+lookback] corresponds to date index[i+lookback-1]
    valid_dates = returns.index[lookback - 1 :]
    assert len(valid_dates) == n_windows, (
        f"Date alignment mismatch: {len(valid_dates)} dates vs {n_windows} windows"
    )
    sig_features = pd.DataFrame(
        sig_features.values if hasattr(sig_features, "values") else sig_features,
        index=valid_dates,
    )
    sig_features.columns = [f"mcsig.{i}" for i in range(sig_features.shape[1])]

    return sig_features


def build_multichannel_dataset(
    baseline_X: pd.DataFrame,
    baseline_index: pd.Index,
    baseline_columns: list[str],
    returns: pd.Series,
    lookback: int = 20,
    level: int = 2,
) -> tuple[pd.DataFrame, list[str]]:
    """Build aligned dataset with baseline features + multi-channel signatures.

    Returns (X_mc, mc_sig_cols).
    """
    sig_df = build_multichannel_rolling_signatures(
        returns, lookback=lookback, level=level,
    )

    # Reindex to baseline, fill warmup with 0
    sig_df = sig_df.reindex(baseline_index).fillna(0.0)
    mc_sig_cols = list(sig_df.columns)

    X_mc = baseline_X.copy()
    X_mc[mc_sig_cols] = sig_df[mc_sig_cols].values

    return X_mc, mc_sig_cols

selected_mc_cols, mc_cluster_map = select_features_by_cluster(
    X_mc, y_base, mc_sig_cols, train_slice=train_sl,
    corr_threshold=0.95, max_features=20,
)

print(f"Multi-channel signature features: {len(mc_sig_cols)}")
print(f"Clusters found: {len(mc_cluster_map)}")
print(f"Selected features (one per cluster): {len(selected_mc_cols)}")
print()

for cid, members in sorted(mc_cluster_map.items(), key=lambda x: -len(x[1])):
    chosen = [m for m in members if m in selected_mc_cols][0]
    if len(members) > 1:
        dropped = [m for m in members if m != chosen]
        print(f"  Cluster {cid} ({len(members)} features): keep {chosen}, drop {dropped}")
    else:
        print(f"  Cluster {cid} (1 feature): keep {chosen}")

fig, axes = plt.subplots(1, 2, figsize=(18, 8))

sns.heatmap(corr_mc_reordered, cmap="RdBu_r", center=0, square=True,
            cbar_kws={"shrink": 0.7}, ax=axes[0],
            xticklabels=False, yticklabels=False)
axes[0].set_title(f"All multi-channel sig features ({len(mc_sig_cols)})")

X_mc_sel = X_mc_only[selected_mc_cols]
corr_mc_sel = X_mc_sel.corr().replace([np.inf, -np.inf], np.nan).fillna(0.0)
if len(selected_mc_cols) > 2:
    dist_mc_sel = 1.0 - np.abs(corr_mc_sel.values)
    np.fill_diagonal(dist_mc_sel, 0.0)
    dist_mc_sel = np.clip(dist_mc_sel, 0.0, None)
    Z_mc_sel = linkage(squareform(dist_mc_sel, checks=False), method="average")
    order_mc_sel = leaves_list(Z_mc_sel)
    corr_mc_sel_reordered = corr_mc_sel.iloc[order_mc_sel, order_mc_sel]
else:
    corr_mc_sel_reordered = corr_mc_sel

sns.heatmap(corr_mc_sel_reordered, cmap="RdBu_r", center=0, square=True,
            cbar_kws={"shrink": 0.7}, ax=axes[1])
axes[1].set_title(f"After cluster selection ({len(selected_mc_cols)} features)")

plt.suptitle("Multi-channel signatures: before vs after feature selection", fontsize=13)
plt.tight_layout()
plt.show()

upper_mc = corr_mc_sel.where(np.triu(np.ones(corr_mc_sel.shape, dtype=bool), k=1))
pairs_mc = upper_mc.stack().rename("corr").reset_index()
pairs_mc.columns = ["f1", "f2", "corr"]
high_mc = pairs_mc[pairs_mc["corr"].abs() >= 0.90].sort_values("corr", key=abs, ascending=False)
print(f"Selected features: {len(selected_mc_cols)}")
print(f"Pairs with |corr| >= 0.90: {len(high_mc)}")
if len(high_mc):
    display(high_mc)

train_slice = slice(SPY_IS_INDEX, SPY_OS_INDEX)
test_slice  = slice(SPY_OS_INDEX, len(y_base))

mc_alpha_variants = {
    "STES (no sig)":       dict(variant="STES_EAESE",  extra_cols=None),
    "STES + mc-sig":       dict(variant=mc_stes_name,   extra_cols=selected_mc_cols),
    "XGB-STES (no sig)":   dict(variant="XGBSTES_E2E", extra_cols=None),
    "XGB-STES + mc-sig":   dict(variant=mc_xgb_name,    extra_cols=selected_mc_cols),
}

mc_alphas = {}
for label, kw in mc_alpha_variants.items():
    _, _, alpha = _fit_predict_oos(
        X=X_mc_selected, y=y_base, r=r_base,
        train_slice=train_slice, test_slice=test_slice,
        seed=0, return_alpha=True, **kw,
    )
    mc_alphas[label] = alpha.iloc[SPY_OS_INDEX - SPY_IS_INDEX :]

fig, axes = plt.subplots(2, 1, figsize=(14, 7), sharex=True)

for label in ["STES (no sig)", "STES + mc-sig"]:
    axes[0].plot(mc_alphas[label].index, mc_alphas[label].values, label=label, linewidth=0.7)
axes[0].set_ylabel(r"$\alpha_t$")
axes[0].set_title("STES — predicted smoothing parameter (OOS, multi-channel)")
axes[0].legend()
axes[0].grid(True, alpha=0.3)

for label in ["XGB-STES (no sig)", "XGB-STES + mc-sig"]:
    axes[1].plot(mc_alphas[label].index, mc_alphas[label].values, label=label, linewidth=0.7)
axes[1].set_ylabel(r"$\alpha_t$")
axes[1].set_title("XGBoost-STES — predicted smoothing parameter (OOS, multi-channel)")
axes[1].legend()
axes[1].grid(True, alpha=0.3)

fig.tight_layout()
plt.show()

baseline_feat_cols = [c for c in baseline_columns if c != "const"]
train_sl_diag = slice(SPY_IS_INDEX, SPY_OS_INDEX)

# X_sig has both baseline + sig columns
all_cols = baseline_feat_cols + selected_sig_cols
cross_corr = X_sig[all_cols].iloc[train_sl_diag].corr()
cross_block = cross_corr.loc[selected_sig_cols, baseline_feat_cols]

# Short labels for display
short_base = {c: c.split(".")[1][:12] for c in baseline_feat_cols}
short_sig = {c: c.replace("sig.2.", "S_")[:14] for c in selected_sig_cols}

fig, ax = plt.subplots(figsize=(8, 5))
sns.heatmap(
    cross_block.rename(index=short_sig, columns=short_base),
    annot=True, fmt=".2f", cmap="RdBu_r", center=0,
    ax=ax, cbar_kws={"shrink": 0.8}, vmin=-1, vmax=1,
)
ax.set_title("Cross-correlation: selected signature features vs baseline features\n(training set)")
plt.tight_layout()
plt.show()

# Flag notable correlations
print("\nPairs with |corr| > 0.3:")
for sig_col in selected_sig_cols:
    for base_col in baseline_feat_cols:
        c = cross_block.loc[sig_col, base_col]
        if abs(c) > 0.3:
            print(f"  {sig_col} <-> {base_col}: {c:.3f}")

max_abs = cross_block.abs().max().max()
print(f"\nMax |cross-corr|: {max_abs:.3f}")

train_slice = slice(SPY_IS_INDEX, SPY_OS_INDEX)
test_slice = slice(SPY_OS_INDEX, len(y_base))

# ---- Baseline STES ----
cols_stes_base = select_variant_columns(X_base, "STES_EAESE") or ["const"]
X_sel_stes_base = X_base[cols_stes_base]
X_tr_stes_base, X_te_stes_base = _scale_train_test(X_sel_stes_base, train_slice, test_slice)
y_tr = y_base.iloc[train_slice]
r_tr = r_base.iloc[train_slice]

model_stes_base = STESModel(random_state=0)
model_stes_base.fit(X_tr_stes_base, y_tr, returns=r_tr, start_index=0, end_index=len(X_tr_stes_base))

# ---- STES + single-channel sig (with uniform L2 CV) ----
cols_stes_sig = [c for c in select_variant_columns(X_sig, "STES_EAESE") if not _is_signature_col(c)]
cols_stes_sig = cols_stes_sig + [c for c in selected_sig_cols if c not in set(cols_stes_sig)]
X_sel_stes_sig = X_sig[cols_stes_sig]
X_tr_stes_sig, X_te_stes_sig = _scale_train_test(X_sel_stes_sig, train_slice, test_slice)

l2_grid = [{"l2_reg": 0.0}, {"l2_reg": 0.01}, {"l2_reg": 0.1}, {"l2_reg": 1.0}, {"l2_reg": 10.0}]
model_stes_sig = STESModel(random_state=0)
model_stes_sig.fit(
    X_tr_stes_sig, y_tr, returns=r_tr,
    start_index=0, end_index=len(X_tr_stes_sig),
    perform_cv=True, cv_grid=l2_grid, cv_splits=5,
 )

# ---- Compact comparison table ----
base_df = pd.DataFrame({"feature": cols_stes_base, "beta_base": model_stes_base.params})
aug_df = pd.DataFrame({"feature": cols_stes_sig, "beta_aug": model_stes_sig.params})
cmp = base_df.merge(aug_df, on="feature", how="outer")
cmp["group"] = np.where(
    cmp["feature"].apply(_is_signature_col),
    "sig",
    np.where(cmp["feature"] == "const", "const", "base"),
 )
cmp["delta"] = cmp["beta_aug"] - cmp["beta_base"]
cmp["delta_pct"] = np.where(
    cmp["beta_base"].abs() > 1e-10,
    cmp["delta"] / cmp["beta_base"].abs() * 100.0,
    np.nan,
 )
_grp_ord = {"const": 0, "base": 1, "sig": 2}
cmp["_ord"] = cmp["group"].map(_grp_ord).fillna(9)
cmp = cmp.sort_values(["_ord", "feature"]).drop(columns=["_ord"])
display(cmp)

# Evaluate OOS RMSE for both
y_te = y_base.iloc[test_slice]
r_te = r_base.iloc[test_slice]

X_all_base = pd.concat([X_tr_stes_base, X_te_stes_base])
r_all = pd.concat([r_tr, r_te])
yhat_base_all = model_stes_base.predict(X_all_base, returns=r_all)
yhat_base_os = np.asarray(yhat_base_all)[len(X_tr_stes_base):]

X_all_sig = pd.concat([X_tr_stes_sig, X_te_stes_sig])
yhat_sig_all = model_stes_sig.predict(X_all_sig, returns=r_all)
yhat_sig_os = np.asarray(yhat_sig_all)[len(X_tr_stes_sig):]

rmse_base = calc_rmse(y_te.values, yhat_base_os)
rmse_sig = calc_rmse(y_te.values, yhat_sig_os)
print(f"\nOut-of-sample RMSE")
print(f"  Baseline STES:    {rmse_base:.6f}")
print(f"  STES + Sig (L2):  {rmse_sig:.6f}")
print(f"  Δ RMSE:           {rmse_sig - rmse_base:+.6f} ({(rmse_sig/rmse_base - 1)*100:+.2f}%)")

l2_values = [0.0, 0.001, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0, 10.0, 50.0, 100.0]

path_results = []
path_params = {c: [] for c in cols_stes_sig}

for l2 in l2_values:
    model = STESModel(random_state=0, l2_reg=l2)
    model.fit(X_tr_stes_sig, y_tr, returns=r_tr, start_index=0, end_index=len(X_tr_stes_sig))

    # OOS prediction
    yhat_all = model.predict(X_all_sig, returns=r_all)
    yhat_os = np.asarray(yhat_all)[len(X_tr_stes_sig):]
    rmse_os = calc_rmse(y_te.values, yhat_os)

    # IS prediction
    yhat_is = np.asarray(yhat_all)[:len(X_tr_stes_sig)]
    rmse_is = calc_rmse(y_tr.values, yhat_is)

    # Parameter norms
    sig_params = [p for c, p in zip(cols_stes_sig, model.params) if _is_signature_col(c)]
    base_params = [p for c, p in zip(cols_stes_sig, model.params) if not _is_signature_col(c) and c != "const"]
    const_param = [p for c, p in zip(cols_stes_sig, model.params) if c == "const"]

    for c, p in zip(cols_stes_sig, model.params):
        path_params[c].append(p)

    path_results.append({
        "l2": l2,
        "rmse_os": rmse_os,
        "rmse_is": rmse_is,
        "norm_sig": np.sqrt(np.sum(np.array(sig_params)**2)),
        "norm_base": np.sqrt(np.sum(np.array(base_params)**2)),
        "const": const_param[0] if const_param else np.nan,
    })

path_df = pd.DataFrame(path_results)

fig, axes = plt.subplots(1, 3, figsize=(16, 5))

# Panel 1: RMSE vs L2
ax = axes[0]
ax.semilogx(path_df["l2"].replace(0, 1e-4), path_df["rmse_os"], "o-", label="OS RMSE (sig)", color="C1")
ax.semilogx(path_df["l2"].replace(0, 1e-4), path_df["rmse_is"], "s--", label="IS RMSE (sig)", color="C1", alpha=0.5)
ax.axhline(rmse_base, color="C0", ls="--", lw=1.5, label=f"Baseline OS RMSE = {rmse_base:.6f}")
ax.set_xlabel("L2 penalty (λ)")
ax.set_ylabel("RMSE")
ax.set_title("RMSE vs L2 penalty")
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

# Panel 2: Parameter norms
ax = axes[1]
ax.semilogx(path_df["l2"].replace(0, 1e-4), path_df["norm_sig"], "o-", label="‖β_sig‖", color="C3")
ax.semilogx(path_df["l2"].replace(0, 1e-4), path_df["norm_base"], "s-", label="‖β_base‖", color="C0")
ax.set_xlabel("L2 penalty (λ)")
ax.set_ylabel("Parameter L2 norm")
ax.set_title("Parameter norms vs L2 penalty")
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

# Panel 3: Individual parameter paths
ax = axes[2]
l2_plot = [max(v, 1e-4) for v in l2_values]
for c in cols_stes_sig:
    is_sig = _is_signature_col(c)
    style = "--" if is_sig else "-"
    lw = 1 if is_sig else 2
    parts = c.split(".")
    short = parts[-1][:8] if is_sig else (parts[1][:10] if len(parts) > 1 else c[:10])
    ax.semilogx(l2_plot, path_params[c], style, lw=lw, label=short)
ax.set_xlabel("L2 penalty (λ)")
ax.set_ylabel("β value")
ax.set_title("Individual parameter paths")
ax.legend(fontsize=6, ncol=2)
ax.grid(True, alpha=0.3)

plt.suptitle("L2 regularization path: STES + single-channel signatures", fontsize=13)
plt.tight_layout()
plt.show()

print(f"\n{'L2':>8s}  {'OS RMSE':>10s}  {'IS RMSE':>10s}  {'‖β_sig‖':>10s}  {'‖β_base‖':>10s}  {'const':>10s}")
for _, row in path_df.iterrows():
    marker = " ← CV" if row["l2"] == model_stes_sig.l2_reg else ""
    print(f"{row['l2']:>8.3f}  {row['rmse_os']:>10.6f}  {row['rmse_is']:>10.6f}  {row['norm_sig']:>10.6f}  {row['norm_base']:>10.6f}  {row['const']:>10.4f}{marker}")
print(f"\nBaseline OS RMSE (no sig): {rmse_base:.6f}")
print(f"Best augmented OS RMSE:    {path_df['rmse_os'].min():.6f} at λ={path_df.loc[path_df['rmse_os'].idxmin(), 'l2']:.3f}")

def fit_stes_group_l2(X_tr, y_tr, r_tr, cols, l2_sig, l2_base=0.0, seed=0):
    """Fit STES with group-specific L2 penalties."""
    X_np = X_tr.values.astype(float)
    y_np = y_tr.values.astype(float)
    r_np = r_tr.values.astype(float)
    n, p = X_np.shape

    # Build group penalty vector
    mask = np.zeros(p)
    for i, c in enumerate(cols):
        if _is_signature_col(c):
            mask[i] = l2_sig
        elif c != "const":
            mask[i] = l2_base

    scale = np.linalg.norm(y_np)
    penalty_vec = np.sqrt(mask) * scale if np.any(mask > 0) else None

    rng = np.random.default_rng(seed)
    x0 = rng.normal(0, 1, size=p)

    result = least_squares(
        STESModel._objective, x0=x0, jac=STESModel._jacobian,
        args=(r_np, X_np, y_np, 0, len(X_np), penalty_vec),
    )
    return result.x

# Sweep l2_sig with l2_base = 0
l2_sig_values = [0.0, 0.001, 0.01, 0.05, 0.1, 0.5, 1.0, 5.0, 10.0, 50.0, 100.0]

group_results = []
for l2_sig in l2_sig_values:
    params = fit_stes_group_l2(X_tr_stes_sig, y_tr, r_tr, cols_stes_sig, l2_sig=l2_sig, l2_base=0.0)

    # Predict
    alphas = expit(X_all_sig.values @ params)
    r_all_np = r_all.values
    returns2 = r_all_np**2
    n = len(r_all_np)
    v = np.empty(n + 1)
    v[0] = returns2[0]
    yhat = np.empty(n)
    for t in range(n):
        yhat[t] = alphas[t] * returns2[t] + (1 - alphas[t]) * v[t]
        v[t + 1] = yhat[t]

    yhat_os = yhat[len(X_tr_stes_sig):]
    rmse_os = calc_rmse(y_te.values, yhat_os)

    sig_params = [p for c, p in zip(cols_stes_sig, params) if _is_signature_col(c)]
    base_params = [p for c, p in zip(cols_stes_sig, params) if not _is_signature_col(c) and c != "const"]

    group_results.append({
        "l2_sig": l2_sig,
        "rmse_os": rmse_os,
        "norm_sig": np.sqrt(np.sum(np.array(sig_params)**2)),
        "norm_base": np.sqrt(np.sum(np.array(base_params)**2)),
    })

group_df = pd.DataFrame(group_results)

fig, axes = plt.subplots(1, 2, figsize=(12, 5))

ax = axes[0]
ax.semilogx(group_df["l2_sig"].replace(0, 1e-4), group_df["rmse_os"], "o-", label="OS RMSE (group L2)", color="C2")
ax.semilogx(path_df["l2"].replace(0, 1e-4), path_df["rmse_os"], "s--", label="OS RMSE (uniform L2)", color="C1", alpha=0.6)
ax.axhline(rmse_base, color="C0", ls="--", lw=1.5, label=f"Baseline OS RMSE = {rmse_base:.6f}")
ax.set_xlabel("L2 penalty on signature features (λ_sig)")
ax.set_ylabel("RMSE")
ax.set_title("Group L2 vs Uniform L2")
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

ax = axes[1]
ax.semilogx(group_df["l2_sig"].replace(0, 1e-4), group_df["norm_sig"], "o-", label="‖β_sig‖", color="C3")
ax.semilogx(group_df["l2_sig"].replace(0, 1e-4), group_df["norm_base"], "s-", label="‖β_base‖ (unpenalized)", color="C0")
ax.set_xlabel("L2 penalty on signature features (λ_sig)")
ax.set_ylabel("Parameter L2 norm")
ax.set_title("Group L2: parameter norms")
ax.legend(fontsize=8)
ax.grid(True, alpha=0.3)

plt.suptitle("Group-penalized STES: penalize only signature features (λ_base = 0)", fontsize=13)
plt.tight_layout()
plt.show()

print(f"\n{'λ_sig':>8s}  {'OS RMSE':>10s}  {'‖β_sig‖':>10s}  {'‖β_base‖':>10s}")
for _, row in group_df.iterrows():
    print(f"{row['l2_sig']:>8.3f}  {row['rmse_os']:>10.6f}  {row['norm_sig']:>10.6f}  {row['norm_base']:>10.6f}")
print(f"\nBaseline OS RMSE: {rmse_base:.6f}")
print(f"Best group-L2 OS RMSE: {group_df['rmse_os'].min():.6f} at λ_sig={group_df.loc[group_df['rmse_os'].idxmin(), 'l2_sig']:.3f}")

grid_e2e = [
    {"min_child_weight": 5.0, "learning_rate": 0.05, "max_depth": 3},
    {"min_child_weight": 20.0, "learning_rate": 0.1, "max_depth": 3},
    {"min_child_weight": 50.0, "learning_rate": 0.05, "max_depth": 2},
]

r_tr_scaled = r_tr * SCALE_FACTOR
y_tr_scaled = y_tr * (SCALE_FACTOR**2)

def _parse_gblinear_dump(dump_str):
    """Parse gblinear dump → (bias, {feature_index: weight})."""
    raw = dump_str.strip()
    bias_sec, weight_sec = raw.split("weight:")
    bias = float(bias_sec.replace("bias:", "").strip())
    weights = [float(v) for v in weight_sec.strip().split("\n")]
    return bias, weights

# ── 1) Baseline-only model ─────────────────────────────────────────
# Drop 'const' — gblinear has its own unpenalized bias term
cols_base_xgb = [c for c in baseline_columns if c != "const"]
X_sel_base = X_base[cols_base_xgb]
X_tr_base, X_te_base = _scale_train_test(X_sel_base, train_slice, test_slice)

over = _xgb_variant_overrides("XGBSTES_E2E", cols_base_xgb)
model_xgb_base = _make_xgb_stes_model(seed=None, **over)
model_xgb_base.fit(
    X_tr_base, y_tr_scaled, returns=r_tr_scaled,
    perform_cv=True, cv_grid=grid_e2e, cv_splits=5,
    start_index=0, end_index=len(X_tr_base),
)
bias_base, weights_base = _parse_gblinear_dump(model_xgb_base.model_.get_dump()[0])
assert len(weights_base) == len(cols_base_xgb)
w_base = dict(zip(cols_base_xgb, weights_base))

cols_xgb_sig = [c for c in baseline_columns + selected_sig_cols if c != "const"]
X_sel_xgb_sig = X_sig[cols_xgb_sig]
X_tr_xgb_sig, X_te_xgb_sig = _scale_train_test(X_sel_xgb_sig, train_slice, test_slice)

over_sig = _xgb_variant_overrides("XGBSTES_E2E", cols_xgb_sig)
model_xgb_sig = _make_xgb_stes_model(seed=None, **over_sig)
model_xgb_sig.fit(
    X_tr_xgb_sig, y_tr_scaled, returns=r_tr_scaled,
    perform_cv=True, cv_grid=grid_e2e, cv_splits=5,
    start_index=0, end_index=len(X_tr_xgb_sig),
)
bias_sig, weights_sig = _parse_gblinear_dump(model_xgb_sig.model_.get_dump()[0])
assert len(weights_sig) == len(cols_xgb_sig)
w_sig = dict(zip(cols_xgb_sig, weights_sig))

# ── 3) Comparison table ───────────────────────────────────────────
print(f"{'Feature':>42s}  {'Baseline':>10s}  {'+ Sig':>10s}  {'Δw':>10s}  {'Δ%':>8s}")
print("-" * 86)
for c in cols_base_xgb:
    wb = w_base[c]
    ws = w_sig[c]
    dw = ws - wb
    dp = 100.0 * dw / wb if abs(wb) > 1e-12 else float("nan")
    print(f"  {c:>40s}  {wb:>10.6f}  {ws:>10.6f}  {dw:>+10.6f}  {dp:>+7.1f}%")
print(f"  {'bias':>40s}  {bias_base:>10.6f}  {bias_sig:>10.6f}  {bias_sig - bias_base:>+10.6f}")
print()
for c in cols_xgb_sig:
    if _is_signature_col(c):
        print(f"  {c:>40s}  {'—':>10s}  {w_sig[c]:>10.6f}   ← SIG")

# ── 4) Side-by-side bar chart ─────────────────────────────────────
fig, axes = plt.subplots(1, 2, figsize=(14, 5), sharey=False)

# Left panel: baseline-only weights
ax = axes[0]
pairs_b = sorted(w_base.items(), key=lambda x: abs(x[1]))
names_b = [p[0].split(".")[-1][:12] for p in pairs_b]
vals_b = [p[1] for p in pairs_b]
ax.barh(range(len(names_b)), vals_b, color="#1f77b4")
ax.set_yticks(range(len(names_b)))
ax.set_yticklabels(names_b, fontsize=9)
ax.set_xlabel("gblinear weight")
ax.set_title("Baseline only")
ax.axvline(0, color="k", lw=0.5)

# Right panel: baseline + signature weights
ax = axes[1]
pairs_s = sorted(w_sig.items(), key=lambda x: abs(x[1]))
names_s = [p[0].split(".")[-1][:12] for p in pairs_s]
vals_s = [p[1] for p in pairs_s]
colors = ["#d62728" if _is_signature_col(p[0]) else "#1f77b4" for p in pairs_s]
ax.barh(range(len(names_s)), vals_s, color=colors)
ax.set_yticks(range(len(names_s)))
ax.set_yticklabels(names_s, fontsize=9)
ax.set_xlabel("gblinear weight")
ax.set_title("Baseline + Signature")
ax.axvline(0, color="k", lw=0.5)

fig.suptitle("XGBoost gblinear feature weights (blue = baseline, red = signature)", fontsize=12)
plt.tight_layout()
plt.show()

print(f"\nreg_lambda: {XGB_PARAMS['reg_lambda']}")

from scipy.special import expit as _expit

# ── 4) Mean linear score and resulting alpha ──────────────────────────────
X_tr_base_np = X_tr_base.to_numpy()
X_tr_sig_np = X_tr_xgb_sig.to_numpy()

z_base = bias_base + X_tr_base_np @ np.array(weights_base)
z_sig = bias_sig + X_tr_sig_np @ np.array(weights_sig)

alpha_base = _expit(z_base)
alpha_sig = _expit(z_sig)

print("4. MEAN LINEAR SCORE z AND RESULTING α = σ(z)")
print(f"   {'':>20s}  {'mean(z)':>12s}  {'std(z)':>12s}  {'mean(α)':>12s}  {'std(α)':>12s}")
print("   " + "-" * 64)
print(f"   {'Baseline':>20s}  {z_base.mean():>12.6f}  {z_base.std():>12.6f}  {alpha_base.mean():>12.6f}  {alpha_base.std():>12.6f}")
print(f"   {'+ Signature':>20s}  {z_sig.mean():>12.6f}  {z_sig.std():>12.6f}  {alpha_sig.mean():>12.6f}  {alpha_sig.std():>12.6f}")
print(f"   {'Δ':>20s}  {z_sig.mean() - z_base.mean():>+12.6f}  {z_sig.std() - z_base.std():>+12.6f}  {alpha_sig.mean() - alpha_base.mean():>+12.6f}  {alpha_sig.std() - alpha_base.std():>+12.6f}")

# ── 5) Decompose mean(z) into bias + feature contributions ───────────────
feat_contrib_base = X_tr_base_np @ np.array(weights_base)
feat_contrib_sig = X_tr_sig_np @ np.array(weights_sig)

# Per-feature contribution to mean(z): w_j * mean(x_j)
print(f"\n5. DECOMPOSITION OF mean(z) = bias + Σ w_j·mean(x_j)")
print(f"   {'':>20s}  {'bias':>12s}  {'Σw·x̄':>12s}  {'mean(z)':>12s}")
print("   " + "-" * 52)
print(f"   {'Baseline':>20s}  {bias_base:>12.6f}  {feat_contrib_base.mean():>12.6f}  {z_base.mean():>12.6f}")
print(f"   {'+ Signature':>20s}  {bias_sig:>12.6f}  {feat_contrib_sig.mean():>12.6f}  {z_sig.mean():>12.6f}")

print(f"\n   Since training features are centered, Σw·x̄ ≈ 0 and mean(z) ≈ bias.")
print(f"   The bias determines the long-run alpha level: σ({bias_base:.3f}) = {_expit(bias_base):.4f}")
print(f"   vs σ({bias_sig:.3f}) = {_expit(bias_sig):.4f}")
print(f"   Δα (from bias alone) = {_expit(bias_sig) - _expit(bias_base):+.4f}")

# ── 6) IS-sample / OOS RMSE comparison ───────────────────────────────────
r_tr_np = r_tr_scaled.to_numpy().reshape(-1)
y_tr_np = y_tr_scaled.to_numpy().reshape(-1)

yhat_base_all = model_xgb_base.predict(
    pd.concat([X_tr_base, X_te_base]),
    returns=pd.concat([r_tr_scaled, r_te * SCALE_FACTOR]),
)
yhat_sig_all = model_xgb_sig.predict(
    pd.concat([X_tr_xgb_sig, X_te_xgb_sig]),
    returns=pd.concat([r_tr_scaled, r_te * SCALE_FACTOR]),
)
yhat_base_is = np.asarray(yhat_base_all)[:len(X_tr_base)]
yhat_sig_is = np.asarray(yhat_sig_all)[:len(X_tr_xgb_sig)]
yhat_base_os = np.asarray(yhat_base_all)[len(X_tr_base):]
yhat_sig_os = np.asarray(yhat_sig_all)[len(X_tr_xgb_sig):]
y_te_np = y_te.to_numpy().reshape(-1) * (SCALE_FACTOR**2)

rmse_is_base = np.sqrt(np.mean((y_tr_np - yhat_base_is) ** 2))
rmse_is_sig = np.sqrt(np.mean((y_tr_np - yhat_sig_is) ** 2))
rmse_os_base = np.sqrt(np.mean((y_te_np - yhat_base_os) ** 2))
rmse_os_sig = np.sqrt(np.mean((y_te_np - yhat_sig_os) ** 2))

print(f"\n6. IN-SAMPLE vs OUT-OF-SAMPLE RMSE (scaled units)")
print(f"   {'':>20s}  {'IS RMSE':>12s}  {'OOS RMSE':>12s}  {'OOS/IS':>8s}")
print("   " + "-" * 56)
print(f"   {'Baseline':>20s}  {rmse_is_base:>12.6f}  {rmse_os_base:>12.6f}  {rmse_os_base/rmse_is_base:>8.4f}")
print(f"   {'+ Signature':>20s}  {rmse_is_sig:>12.6f}  {rmse_os_sig:>12.6f}  {rmse_os_sig/rmse_is_sig:>8.4f}")
print(f"   {'Δ (sig-base)':>20s}  {rmse_is_sig - rmse_is_base:>+12.6f}  {rmse_os_sig - rmse_os_base:>+12.6f}")

if rmse_is_sig > rmse_is_base:
    print(f"\n   ⚠ Signature model has HIGHER in-sample loss → possible convergence issue")
else:
    print(f"\n   ✓ Signature model has lower IS loss (as expected with more features)")
    print(f"     But higher OOS loss → overfitting, not a convergence problem")

from scipy.special import expit as _expit

X_tr_base_np = X_tr_base.to_numpy()
X_tr_sig_np = X_tr_xgb_sig.to_numpy()
z_base = bias_base + X_tr_base_np @ np.array(weights_base)
z_sig = bias_sig + X_tr_sig_np @ np.array(weights_sig)

# ── 8) COUNTERFACTUAL: sig model weights but baseline bias ───────────────
print("8. COUNTERFACTUAL: SIG MODEL WITH BASELINE BIAS")
z_sig_cf_tr = bias_base + X_tr_sig_np @ np.array(weights_sig)
z_sig_cf_te = bias_base + X_te_xgb_sig.to_numpy() @ np.array(weights_sig)
z_sig_te = bias_sig + X_te_xgb_sig.to_numpy() @ np.array(weights_sig)

print(f"   α with actual sig bias ({bias_sig:.3f}):     train={_expit(z_sig).mean():.4f}  test={_expit(z_sig_te).mean():.4f}")
print(f"   α with baseline bias   ({bias_base:.3f}):     train={_expit(z_sig_cf_tr).mean():.4f}  test={_expit(z_sig_cf_te).mean():.4f}")
print(f"   Δα from bias shift:                  train={_expit(z_sig).mean() - _expit(z_sig_cf_tr).mean():+.4f}  test={_expit(z_sig_te).mean() - _expit(z_sig_cf_te).mean():+.4f}")

# ── 9) Feature contribution to z dynamics ────────────────────────────────
print(f"\n9. |w·std(x)| — FEATURE CONTRIBUTION TO α DYNAMICS")

# Short labels for compact table
def _short_label(c):
    if _is_signature_col(c):
        return c.split(".")[-1][:10]
    return c.split(".")[-2][:10] if "." in c else c[:10]

# Augmented model contributions (sorted by |w·std|)
scores_sig = [(c, w_sig[c], X_tr_xgb_sig[c].std(), abs(w_sig[c] * X_tr_xgb_sig[c].std()))
              for c in cols_xgb_sig]
total_sig = sum(s[3] for s in scores_sig)

print(f"\n   {'Feature':>14s}  {'w':>9s}  {'|w·std|':>9s}  {'%':>6s}  {'type':>5s}")
print("   " + "-" * 50)
for c, w, s, ws in sorted(scores_sig, key=lambda x: -x[3]):
    tag = "SIG" if _is_signature_col(c) else "base"
    print(f"   {_short_label(c):>14s}  {w:>+9.4f}  {ws:>9.4f}  {100*ws/total_sig:>5.1f}%  {tag:>5s}")

sig_pct = sum(s[3] for s in scores_sig if _is_signature_col(s[0])) / total_sig * 100
base_pct = 100 - sig_pct

print(f"\n   Baseline features: {base_pct:.1f}% of z dynamics")
print(f"   Signature features: {sig_pct:.1f}% of z dynamics")
print(f"   std(z): baseline-only = {z_base.std():.4f}, +sig = {z_sig.std():.4f} ({z_sig.std()/z_base.std():.2f}x)")
print(f"\n   INTERPRETATION:")
print(f"   The bias is NOT suboptimal — it is correctly compensating for the")
print(f"   increased z variance that signatures introduce. The sig features")
print(f"   contribute {sig_pct:.0f}% of α dynamics, increasing std(z) by {z_sig.std()/z_base.std():.2f}x.")
print(f"   The bias shifts up by {bias_sig - bias_base:+.3f} to keep the mean α at a")
print(f"   sensible level ({_expit(bias_base):.2%} → {_expit(bias_sig):.2%}), because the wider")
print(f"   z distribution pushes more mass into the sigmoid tails.")