wfo_optimization.py

def month_starts(index: pd.DatetimeIndex, start: pd.Timestamp) -> list[pd.Timestamp]:
    """All month starts >= start that exist in the index."""
    months = pd.date_range(start=start, end=index.max(), freq="MS")
    # Keep only months that have data
    return [m for m in months if (index >= m).any()]


def optimize_params(train_feat: pd.DataFrame, train_regime: pd.Series) -> dict:
    """
    Grid search parameters on training set. Objective: maximize Sharpe (simple & common).
    """
    best = None
    best_sharpe = -1e9

    for z_thr in Z_THRESH_GRID:
        for range_size in RANGE_SIZE_GRID:
            for lowvol_size in LOWVOL_SIZE_GRID:
                for highvol_size in HIGHVOL_SIZE_GRID:
                    params = {
                        "z_thr": z_thr,
                        "range_size": range_size,
                        "lowvol_size": lowvol_size,
                        "highvol_size": highvol_size,
                    }
                    pos = make_positions(train_feat, train_regime.loc[train_feat.index], params)
                    bt = backtest(train_feat, pos)

                    # Sharpe on arithmetic returns of equity
                    eq = bt["equity"]
                    r = eq.pct_change().fillna(0.0)
                    ann_ret = r.mean() * 252
                    ann_vol = r.std(ddof=0) * math.sqrt(252)
                    sharpe = float(ann_ret / (ann_vol + 1e-12))

                    if sharpe > best_sharpe:
                        best_sharpe = sharpe
                        best = params

    return best


def walk_forward_oos(df_feat: pd.DataFrame, regime: pd.Series) -> tuple[pd.Series, pd.DataFrame]:
    """
    Monthly WFO:
    - For each month in OOS, optimize params on the last TRAIN_YEARS of data,
      then trade the next month using those params.
    Returns:
      - oos_equity (rebased to 1.0 at OOS_START)
      - chosen_params_df (month -> best params)
    """
    oos_start = pd.to_datetime(OOS_START)
    months = month_starts(df_feat.index, oos_start)

    # Collect monthly OOS returns (log-return proxy 'net') and chosen parameters
    oos_net = []
    param_rows = []

    for m in months:
        # test month: [m, next_month)
        next_m = (m + pd.offsets.MonthBegin(1)).normalize()
        test_idx = df_feat.index[(df_feat.index >= m) & (df_feat.index < next_m)]
        if len(test_idx) < 5:
            continue

        # training window: last TRAIN_YEARS up to month start
        train_start = (m - pd.DateOffset(years=TRAIN_YEARS)).normalize()
        train_idx = df_feat.index[(df_feat.index >= train_start) & (df_feat.index < m)]
        if len(train_idx) < 100:
            continue

        train_feat = df_feat.loc[train_idx]
        train_reg = regime.loc[train_idx]

        best_params = optimize_params(train_feat, train_reg)

        # Evaluate on test month
        # Add one prior day so the pos.shift(1) has a previous bar
        prev_day_idx = df_feat.index[df_feat.index < m]
        if len(prev_day_idx) > 0:
            start_for_test = prev_day_idx.max()
            test_feat_ext = df_feat.loc[(df_feat.index >= start_for_test) & (df_feat.index < next_m)]
        else:
            test_feat_ext = df_feat.loc[(df_feat.index >= m) & (df_feat.index < next_m)]

        test_reg_ext = regime.loc[test_feat_ext.index]
        pos = make_positions(test_feat_ext, test_reg_ext, best_params)
        bt = backtest(test_feat_ext, pos)

        # Keep only the test-month rows
        bt_test = bt.loc[test_idx]
        oos_net.append(bt_test["net"])

        param_rows.append({
            "month_start": m.date().isoformat(),
            **best_params,
        })

    if not oos_net:
        raise RuntimeError("No OOS months were generated. Check your dates and data.")

    oos_net = pd.concat(oos_net).sort_index()

    # Build equity (rebased)
    log_eq = oos_net.cumsum()
    equity = np.exp(log_eq)
    equity.name = "equity_oos"

    params_df = pd.DataFrame(param_rows)
    return equity, params_df