from pathlib import Path
import pandas as pd
import requests
from io import StringIO
import matplotlib.pyplot as plt
import numpy as np

DATA_DIR = Path("data")
DATA_DIR.mkdir(exist_ok=True)

RAW_CSV_PATH = DATA_DIR / "feds200628.csv"   # local cached file
DAILY_ZC_CSV_PATH = DATA_DIR / "gsw_zero_coupon_daily.csv"
MONTHLY_ZC_CSV_PATH = DATA_DIR / "gsw_zero_coupon_monthly.csv"

GSW_URL = "https://www.federalreserve.gov/data/yield-curve-tables/feds200628.csv"

headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                  "AppleWebKit/537.36 (KHTML, like Gecko) "
                  "Chrome/127.0.0.1 Safari/537.36"
}

if not RAW_CSV_PATH.exists():
    print("Local GSW file not found. Downloading from Fed...")
    resp = requests.get(GSW_URL, headers=headers)
    resp.raise_for_status()

    # Some versions require skipping first 9 rows, but we write the raw text first
    with open(RAW_CSV_PATH, "w", encoding="utf-8") as f:
        f.write(resp.text)

    print(f"Saved raw GSW CSV to {RAW_CSV_PATH.resolve()}")
else:
    print("Local GSW file already exists. Skipping download.")

Local GSW file already exists. Skipping download.

print("Loading local GSW CSV...")

df_raw = pd.read_csv(RAW_CSV_PATH, skiprows=9)

print("Raw shape:", df_raw.shape)
df_raw.head()

Loading local GSW CSV...
Raw shape: (16853, 100)

df_raw.columns

Index(['Date', 'BETA0', 'BETA1', 'BETA2', 'BETA3', 'SVEN1F01', 'SVEN1F04',
       'SVEN1F09', 'SVENF01', 'SVENF02', 'SVENF03', 'SVENF04', 'SVENF05',
       'SVENF06', 'SVENF07', 'SVENF08', 'SVENF09', 'SVENF10', 'SVENF11',
       'SVENF12', 'SVENF13', 'SVENF14', 'SVENF15', 'SVENF16', 'SVENF17',
       'SVENF18', 'SVENF19', 'SVENF20', 'SVENF21', 'SVENF22', 'SVENF23',
       'SVENF24', 'SVENF25', 'SVENF26', 'SVENF27', 'SVENF28', 'SVENF29',
       'SVENF30', 'SVENPY01', 'SVENPY02', 'SVENPY03', 'SVENPY04', 'SVENPY05',
       'SVENPY06', 'SVENPY07', 'SVENPY08', 'SVENPY09', 'SVENPY10', 'SVENPY11',
       'SVENPY12', 'SVENPY13', 'SVENPY14', 'SVENPY15', 'SVENPY16', 'SVENPY17',
       'SVENPY18', 'SVENPY19', 'SVENPY20', 'SVENPY21', 'SVENPY22', 'SVENPY23',
       'SVENPY24', 'SVENPY25', 'SVENPY26', 'SVENPY27', 'SVENPY28', 'SVENPY29',
       'SVENPY30', 'SVENY01', 'SVENY02', 'SVENY03', 'SVENY04', 'SVENY05',
       'SVENY06', 'SVENY07', 'SVENY08', 'SVENY09', 'SVENY10', 'SVENY11',
       'SVENY12', 'SVENY13', 'SVENY14', 'SVENY15', 'SVENY16', 'SVENY17',
       'SVENY18', 'SVENY19', 'SVENY20', 'SVENY21', 'SVENY22', 'SVENY23',
       'SVENY24', 'SVENY25', 'SVENY26', 'SVENY27', 'SVENY28', 'SVENY29',
       'SVENY30', 'TAU1', 'TAU2'],
      dtype='str')

# 1. Normalize column names (just in case)
df = df_raw.copy()

# Make sure the date column is correctly named
date_col_candidates = ["Date", "date", "DATE"]
date_col = None
for c in date_col_candidates:
    if c in df.columns:
        date_col = c
        break

if date_col is None:
    raise ValueError(f"Could not find a date column in raw data. Columns: {df.columns}")

df[date_col] = pd.to_datetime(df[date_col])
df = df.sort_values(by=date_col)

# 2. Select zero-coupon columns (SVENYxx)
zc_cols = [c for c in df.columns if c.startswith("SVENY")]
print("Zero-coupon columns:", zc_cols)

if not zc_cols:
    raise ValueError("No zero-coupon (SVENYxx) columns found. Check raw CSV format.")

# 3. Keep Date + zero-coupon columns
df_zc = df[[date_col] + zc_cols].copy()
df_zc = df_zc.rename(columns={date_col: "date"})
df_zc.set_index("date", inplace=True)
df_zc.sort_index(inplace=True)

print("Zero-coupon daily data (raw units):")
df_zc.head()

Zero-coupon columns: ['SVENY01', 'SVENY02', 'SVENY03', 'SVENY04', 'SVENY05', 'SVENY06', 'SVENY07', 'SVENY08', 'SVENY09', 'SVENY10', 'SVENY11', 'SVENY12', 'SVENY13', 'SVENY14', 'SVENY15', 'SVENY16', 'SVENY17', 'SVENY18', 'SVENY19', 'SVENY20', 'SVENY21', 'SVENY22', 'SVENY23', 'SVENY24', 'SVENY25', 'SVENY26', 'SVENY27', 'SVENY28', 'SVENY29', 'SVENY30']
Zero-coupon daily data (raw units):

# The SVENYxx convention:
# SVENY01 -> 1-year zero-coupon, SVENY02 -> 2-year, ..., typically up to 30.
# We'll map 'SVENY01' -> 1.0, 'SVENY02' -> 2.0, etc., and rename columns to "1.0","2.0",...

maturities_years = []
for c in zc_cols:
    # x = last two chars -> '01', '02', etc.
    # Some files may have 3 digits if > 99, but for Treasuries we expect <= 30.
    suffix = c.replace("SVENY", "")
    try:
        mat = int(suffix)
    except ValueError:
        raise ValueError(f"Unexpected SVENY column name format: {c}")
    maturities_years.append(mat)

# New column names as string years, e.g. "1.0", "2.0", "3.0", ...
new_cols = [f"{mat:.1f}" for mat in maturities_years]

zc_renaming = dict(zip(zc_cols, new_cols))
df_zc = df_zc.rename(columns=zc_renaming)

# Convert from percent to decimals
df_zc = df_zc.astype(float) / 100.0

print("Zero-coupon daily yields in decimals:")
df_zc.head()

Zero-coupon daily yields in decimals:

df_zc.to_csv(DAILY_ZC_CSV_PATH, index=True)

# Resample to monthly (end-of-month yields)
df_zc_monthly = df_zc.resample("ME").last()

# Forward-fill any gaps (holidays etc.)
df_zc_monthly = df_zc_monthly.ffill()

# Drop rows that are completely NaN (if any)
df_zc_monthly = df_zc_monthly.dropna(how="all")

print("Monthly zero-coupon yields (decimals):")
df_zc_monthly.head()

Monthly zero-coupon yields (decimals):

df_zc_monthly.to_csv(MONTHLY_ZC_CSV_PATH, index=True)
# print(f"Saved monthly zero-coupon panel to: {MONTHLY_ZC_CSV_PATH.resolve()}")
print(f"Saved monthly zero-coupon panel locally.")

Saved monthly zero-coupon panel locally.

# Plot a few maturities to check series look reasonable
sample_mats = ["1.0", "5.0", "10.0", "30.0"]
sample_mats = [m for m in sample_mats if m in df_zc_monthly.columns]

df_zc_monthly[sample_mats].plot(figsize=(10, 5))
plt.title("GSW Zero-Coupon Yields (Monthly)")
plt.xlabel("Date")
plt.ylabel("Yield (decimal)")
plt.grid(True)
plt.show()

from statsmodels.tsa.api import VAR

DATA_DIR = Path("data")

df_yields = pd.read_csv(DATA_DIR / "gsw_zero_coupon_monthly.csv",
                        index_col=0, parse_dates=True)

# only get data from 1990 onwards
df_yields = df_yields["1990-01-01":]

print(df_yields.shape)
df_yields.head()

(433, 30)

def dl_loadings(maturities: np.ndarray, lam: float) -> np.ndarray:
    tau = maturities
    lam_tau = lam * tau

    with np.errstate(divide="ignore", invalid="ignore"):
        f1 = np.ones_like(tau)
        f2 = (1 - np.exp(-lam_tau)) / lam_tau
        f3 = f2 - np.exp(-lam_tau)

    f2 = np.where(tau == 0, 1.0, f2)
    f3 = np.where(tau == 0, 0.0, f3)

    return np.column_stack([f1, f2, f3])

def estimate_diebold_li_factors(df_yields: pd.DataFrame, lam: float = 0.0609):
    maturities = np.array([float(c) for c in df_yields.columns])
    sort_idx = np.argsort(maturities)
    mats_sorted = maturities[sort_idx]

    df_sorted = df_yields.iloc[:, sort_idx]

    X = dl_loadings(mats_sorted, lam)

    betas = []
    dates = []

    for date, row in df_sorted.iterrows():
        y = row.values.astype(float)
        mask = ~np.isnan(y)

        X_m = X[mask]
        y_m = y[mask]

        if y_m.shape[0] < 3:
            continue

        beta_hat = np.linalg.inv(X_m.T @ X_m) @ X_m.T @ y_m
        betas.append(beta_hat)
        dates.append(date)

    return pd.DataFrame(betas, index=dates, columns=["level","slope","curvature"]).sort_index()

lam = 0.0609  # canonical Diebold–Li value
factors_ols = estimate_diebold_li_factors(df_yields, lam)
factors_ols.to_csv(DATA_DIR / "dl_factors_ols.csv")

print(factors_ols.head())

               level     slope  curvature
1990-01-31  0.076733  0.004404   0.016872
1990-02-28  0.075251  0.006550   0.021459
1990-03-31  0.079159  0.005193   0.012616
1990-04-30  0.080233  0.006865   0.018714
1990-05-31  0.072812  0.009628   0.024918

var_model = VAR(factors_ols)
var_res = var_model.fit(maxlags=1)

A = var_res.coefs[0]       # transition matrix
Q = var_res.sigma_u        # state noise covariance

print("Transition matrix A:")
print(A)
print("State noise covariance Q:")
print(Q)

Transition matrix A:
[[ 0.91940097  0.08694316 -0.04380703]
 [ 0.07225808  0.88825823  0.04739625]
 [ 0.1174264  -0.12157507  1.05306669]]
State noise covariance Q:
              level     slope  curvature
level      0.000129 -0.000121  -0.000251
slope     -0.000121  0.000121   0.000233
curvature -0.000251  0.000233   0.000525

C:\ProgramData\anaconda3\Lib\site-packages\statsmodels\tsa\base\tsa_model.py:473: ValueWarning: No frequency information was provided, so inferred frequency ME will be used.
  self._init_dates(dates, freq)

def estimate_measurement_cov(df_yields, factors_ols, lam):
    # Align by date
    df_y, df_b = df_yields.align(factors_ols, join="inner", axis=0)

    # Coerce to numeric and drop any rows with NaNs
    df_y = df_y.apply(pd.to_numeric, errors="coerce")
    df_b = df_b.apply(pd.to_numeric, errors="coerce")

    row_mask = df_y.notna().all(axis=1) & df_b.notna().all(axis=1)
    df_y = df_y.loc[row_mask]
    df_b = df_b.loc[row_mask]

    print("Measurement cov – using", df_y.shape[0], "dates")

    maturities = np.array([float(c) for c in df_y.columns])
    H = dl_loadings(maturities, lam)  # (n, 3)

    Y = df_y.values   # (T, n)
    B = df_b.values   # (T, 3)

    residuals = []
    for t in range(Y.shape[0]):
        y_t = Y[t, :]        # (n,)
        beta_t = B[t, :]     # (3,)
        y_hat_t = H @ beta_t # (n,)
        e_t = y_t - y_hat_t
        residuals.append(e_t)

    E = np.vstack(residuals)   # (T, n)
    R_full = np.cov(E, rowvar=False)  # cov across maturities

    print("Any NaNs in R_full?", np.isnan(R_full).any())

    sigma2 = float(np.nanmean(np.diag(R_full)))
    n = df_y.shape[1]
    R = sigma2 * np.eye(n)

    return R, H, df_y, df_b

R, H, df_y_clean, df_b_clean = estimate_measurement_cov(df_yields, factors_ols, lam)
print("Shapes: Y", df_y_clean.shape, "H", H.shape, "R", R.shape)

Measurement cov – using 433 dates
Any NaNs in R_full? False
Shapes: Y (433, 30) H (30, 3) R (30, 30)

def kalman_filter_smoother(Y, H, A, Q, R, beta0=None, P0=None):
    Y = np.asarray(Y)
    H = np.asarray(H)
    A = np.asarray(A)
    Q = np.asarray(Q)
    R = np.asarray(R)

    T, n_mats = Y.shape
    n_states = A.shape[0]

    assert H.shape == (n_mats, n_states), f"H shape {H.shape} != ({n_mats},{n_states})"
    assert A.shape == (n_states, n_states)
    assert Q.shape == (n_states, n_states)
    assert R.shape == (n_mats, n_mats)

    beta_pred = np.zeros((T, n_states))
    P_pred = np.zeros((T, n_states, n_states))
    beta_filt = np.zeros((T, n_states))
    P_filt = np.zeros((T, n_states, n_states))

    I = np.eye(n_states)

    if beta0 is None:
        beta0 = np.zeros(n_states)
    if P0 is None:
        P0 = 10.0 * np.eye(n_states)

    beta_prev = beta0
    P_prev = P0

    for t in range(T):
        # Prediction
        beta_t_pred = A @ beta_prev
        P_t_pred = A @ P_prev @ A.T + Q

        # Update
        y_t = Y[t, :]  # (n_mats,)
        S_t = H @ P_t_pred @ H.T + R   # (n_mats, n_mats)
        K_t = P_t_pred @ H.T @ np.linalg.inv(S_t)  # (n_states, n_mats)

        y_hat_t = H @ beta_t_pred     # (n_mats,)
        innov = y_t - y_hat_t         # (n_mats,)

        beta_t_filt = beta_t_pred + K_t @ innov
        P_t_filt = (I - K_t @ H) @ P_t_pred

        beta_pred[t] = beta_t_pred
        P_pred[t] = P_t_pred
        beta_filt[t] = beta_t_filt
        P_filt[t] = P_t_filt

        beta_prev = beta_t_filt
        P_prev = P_t_filt

    # RTS smoother
    beta_smooth = np.zeros_like(beta_filt)
    P_smooth = np.zeros_like(P_filt)

    beta_smooth[-1] = beta_filt[-1]
    P_smooth[-1] = P_filt[-1]

    for t in range(T - 2, -1, -1):
        P_f = P_filt[t]
        P_p_next = P_pred[t + 1]

        C_t = P_f @ A.T @ np.linalg.inv(P_p_next)  # (n_states, n_states)

        beta_smooth[t] = beta_filt[t] + C_t @ (beta_smooth[t + 1] - beta_pred[t + 1])
        P_smooth[t] = P_f + C_t @ (P_smooth[t + 1] - P_p_next) @ C_t.T

    return beta_filt, beta_smooth

# 1) OLS factors
factors_ols = estimate_diebold_li_factors(df_yields, lam)

# 2) A, Q from VAR on OLS factors
var_model = VAR(factors_ols)
var_res = var_model.fit(maxlags=1)
A = var_res.coefs[0]
Q = var_res.sigma_u

# 3) R, H, and cleaned yields/factors
R, H, df_y_clean, df_b_clean = estimate_measurement_cov(df_yields, factors_ols, lam)

print("Shapes before Kalman:")
print("Y:", df_y_clean.shape)
print("H:", H.shape)
print("A:", A.shape)
print("Q:", Q.shape)
print("R:", R.shape)

# 4) Run Kalman + smoother
Y = df_y_clean.values  # (T, n)
beta0 = df_b_clean.iloc[0].values  # first OLS beta as init
P0 = np.eye(3)

beta_filt, beta_smooth = kalman_filter_smoother(Y, H, A, Q, R, beta0, P0)

idx = df_b_clean.index
cols = ["level", "slope", "curvature"]

factors_filt = pd.DataFrame(beta_filt, index=idx, columns=cols)
factors_smooth = pd.DataFrame(beta_smooth, index=idx, columns=cols)

Measurement cov – using 433 dates
Any NaNs in R_full? False
Shapes before Kalman:
Y: (433, 30)
H: (30, 3)
A: (3, 3)
Q: (3, 3)
R: (30, 30)

C:\ProgramData\anaconda3\Lib\site-packages\statsmodels\tsa\base\tsa_model.py:473: ValueWarning: No frequency information was provided, so inferred frequency ME will be used.
  self._init_dates(dates, freq)

# Initial state: use first OLS estimate as beta0
beta0 = df_b_clean.iloc[0].values.astype(float)  # shape (3,)
P0 = np.eye(3) * 1.0  # initial covariance; you can tweak the scale

# Y is the observation matrix (T, n_mats)
Y = df_y_clean.values.astype(float)

beta_filt, beta_smooth = kalman_filter_smoother(
    Y, H, A, Q, R,
    beta0=beta0,
    P0=P0
)

beta_filt.shape, beta_smooth.shape

((433, 3), (433, 3))

idx = df_b_clean.index
cols = ["level", "slope", "curvature"]

factors_filt = pd.DataFrame(beta_filt, index=idx, columns=cols)
factors_smooth = pd.DataFrame(beta_smooth, index=idx, columns=cols)

factors_filt.to_csv(DATA_DIR / "dl_factors_kalman_filtered_sample.csv")
factors_smooth.to_csv(DATA_DIR / "dl_factors_kalman_smoothed_sample.csv")

factors_smooth.head()

plt.figure(figsize=(10,4))
plt.plot(df_b_clean["level"], label="OLS", alpha=0.6)
plt.plot(factors_smooth["level"], label="Kalman smooth", alpha=0.8)
plt.title("Level factor – OLS vs Kalman")
plt.legend()
plt.grid(True)
plt.show()

plt.figure(figsize=(10,4))
plt.plot(df_b_clean["slope"], label="OLS", alpha=0.6)
plt.plot(factors_smooth["slope"], label="Kalman smooth", alpha=0.8)
plt.title("Slope factor – OLS vs Kalman")
plt.legend()
plt.grid(True)
plt.show()

plt.figure(figsize=(10,4))
plt.plot(df_b_clean["curvature"], label="OLS", alpha=0.6)
plt.plot(factors_smooth["curvature"], label="Kalman smooth", alpha=0.8)
plt.title("Curvature factor – OLS vs Kalman")
plt.legend()
plt.grid(True)
plt.show()

def unpack_params(theta):
    """
    Map parameter vector theta (length 16) -> (A, Q, R_scalar).
    A: 3x3
    Q: 3x3 (PSD via Cholesky L L')
    R: scalar variance (we'll build R = R * I outside)
    """
    theta = np.asarray(theta)
    assert theta.size == 16

    # A entries
    A_flat = theta[0:9]
    A = A_flat.reshape(3, 3)

    # Cholesky L parameters for Q
    l11, l21, l22, l31, l32, l33 = theta[9:15]

    L = np.array([
        [np.exp(l11),     0.0,          0.0],
        [l21,         np.exp(l22),      0.0],
        [l31,             l32,      np.exp(l33)]
    ])
    Q = L @ L.T

    # Measurement variance
    log_sigma = theta[15]
    sigma = np.exp(log_sigma)
    R_scalar = sigma**2

    return A, Q, R_scalar

def kalman_loglik(theta, Y, H, beta0=None, P0=None):
    """
    Negative log-likelihood for given theta, using Kalman filter.

    theta: parameter vector (length 16)
    Y: (T, n) array of yields
    H: (n, 3) loadings matrix
    beta0: initial state mean (3,)
    P0: initial state covariance (3,3)

    Returns: negative log-likelihood (float)
    """
    Y = np.asarray(Y)
    T, n = Y.shape

    A, Q, R_scalar = unpack_params(theta)
    R = R_scalar * np.eye(n)

    k = 3
    if beta0 is None:
        beta0 = np.zeros(k)
    if P0 is None:
        P0 = 10.0 * np.eye(k)

    beta_prev = beta0
    P_prev = P0

    I_k = np.eye(k)

    loglik = 0.0
    const = n * np.log(2 * np.pi)

    for t in range(T):
        # Prediction
        beta_pred = A @ beta_prev
        P_pred = A @ P_prev @ A.T + Q

        y_t = Y[t, :]  # (n,)

        # Innovation
        S_t = H @ P_pred @ H.T + R  # (n,n)
        try:
            S_inv = np.linalg.inv(S_t)
            sign, logdet = np.linalg.slogdet(S_t)
            if sign <= 0:
                # Penalize non-PD S_t
                return 1e6
        except np.linalg.LinAlgError:
            return 1e6

        y_hat = H @ beta_pred
        innov = y_t - y_hat  # (n,)

        # Contribution to log-likelihood
        quad = innov.T @ S_inv @ innov
        loglik_t = -0.5 * (const + logdet + quad)
        loglik += loglik_t

        # Kalman update (for next step)
        K_t = P_pred @ H.T @ S_inv  # (3,n)
        beta_filt = beta_pred + K_t @ innov
        P_filt = (I_k - K_t @ H) @ P_pred

        beta_prev, P_prev = beta_filt, P_filt

    # We return negative log-likelihood for minimization
    return -float(loglik)

def cholesky_param_from_Q(Q):
    """
    Take a 3x3 PSD Q and get Cholesky parameter vector
    (l11, l21, l22, l31, l32, l33)
    such that L L' = Q and L has exp(diag) structure.
    """
    L0 = np.linalg.cholesky(Q)
    # enforce positive diag via exp parameterization
    l11 = np.log(L0[0,0])
    l22 = np.log(L0[1,1])
    l33 = np.log(L0[2,2])

    # off-diagonals stay as is
    l21 = L0[1,0]
    l31 = L0[2,0]
    l32 = L0[2,1]

    return np.array([l11, l21, l22, l31, l32, l33])

def initial_theta_from_2step(A_init, Q_init, R_sigma2_init):
    A_flat = A_init.flatten()
    l_params = cholesky_param_from_Q(Q_init)
    log_sigma0 = 0.5 * np.log(R_sigma2_init)
    theta0 = np.concatenate([A_flat, l_params, np.array([log_sigma0])])
    return theta0

# 1. Initial A, Q, R
A_init = np.array([
    [0.98,  0.01,  0.00],
    [0.00,  0.90,  0.05],
    [0.00, -0.05,  0.80]
])

Q_init = np.array([
    [0.0005, 0.0,     0.0],
    [0.0,    0.0010,  0.0],
    [0.0,    0.0,     0.0010]
])

R_sigma2_init = 1e-4

# 2. Turn Q into Cholesky parameters
l_params = cholesky_param_from_Q(Q_init)

# 3. Flatten A and make log-sigma
A_flat = A_init.flatten()
log_sigma0 = 0.5 * np.log(R_sigma2_init)

# 4. Full initial vector (length 16)
theta0 = np.concatenate([
    A_flat,
    l_params,
    np.array([log_sigma0])
])

print(theta0)

[ 0.98        0.01        0.          0.          0.9         0.05
  0.         -0.05        0.8        -3.80045123  0.         -3.45387764
  0.          0.         -3.45387764 -4.60517019]

from scipy.optimize import minimize

# Y, H from cleaned dataset and loadings
Y = df_y_clean.values.astype(float)
maturities = np.array([float(c) for c in df_y_clean.columns])
H = dl_loadings(maturities, lam)  # (n,3)

beta0 = df_b_clean.iloc[0].values  # or zeros
P0 = np.eye(3) * 1.0

def objective(theta):
    return kalman_loglik(theta, Y, H, beta0=beta0, P0=P0)

res = minimize(
    objective,
    theta0,
    method="L-BFGS-B",
    options={"maxiter": 200, "disp": True}
)

print("Converged:", res.success)
print("Final neg loglik:", res.fun)

Converged: True
Final neg loglik: -62125.423142875195

theta_hat = res.x
A_hat, Q_hat, R_scalar_hat = unpack_params(theta_hat)
R_hat = R_scalar_hat * np.eye(Y.shape[1])

beta_filt_mle, beta_smooth_mle = kalman_filter_smoother(
    Y, H, A_hat, Q_hat, R_hat,
    beta0=beta0,
    P0=P0
)

factors_smooth_mle = pd.DataFrame(
    beta_smooth_mle,
    index=df_y_clean.index,
    columns=["level","slope","curvature"]
)

def mle_given_lambda(lam, Y, maturities, theta0):
    H = dl_loadings(maturities, lam)

    def objective(theta):
        return kalman_loglik(theta, Y, H, beta0=beta0, P0=P0)

    res = minimize(objective, theta0, method="L-BFGS-B",
                   options={"maxiter": 200})
    return res.fun, res.x  # neg loglik, theta_hat

Y = df_y_clean.values
maturities = np.array([float(c) for c in df_y_clean.columns])

lambda_grid = np.linspace(0.01, 1.2, 30)  # adjust as you like

best_val = np.inf
best_lam = None
best_theta = None

for lam_try in lambda_grid:
    neg_ll, theta_hat = mle_given_lambda(lam_try, Y, maturities, theta0)
    print(lam_try, neg_ll)
    if neg_ll < best_val:
        best_val = neg_ll
        best_lam = lam_try
        best_theta = theta_hat

print("Best lambda:", best_lam, "neg loglik:", best_val)

# Summary of all ML estimates we will carry forward
A_hat, Q_hat, R_scalar_hat = unpack_params(best_theta)
R_hat = R_scalar_hat * np.eye(Y.shape[1])
H_hat = dl_loadings(maturities, best_lam)

0.01 -59609.83792847459
0.05103448275862069 -64114.66202427518
0.09206896551724138 -69908.23831762242
0.13310344827586207 -69668.23213746391
0.17413793103448277 -68954.21045002024
0.21517241379310348 -68232.16319568738
0.25620689655172413 -67366.85873680399
0.29724137931034483 -64653.09975090846
0.33827586206896554 -67662.39134411809
0.37931034482758624 -67668.60987332094
0.42034482758620695 -67720.91342463228
0.4613793103448276 -67503.25814309693
0.5024137931034482 -67624.32878743885
0.543448275862069 -67190.63888140296
0.5844827586206897 -67601.19367420413
0.6255172413793104 -67484.00442137267
0.6665517241379311 -67280.09686429493

C:\Users\thoma\AppData\Local\Temp\ipykernel_24368\3121696981.py:20: RuntimeWarning: overflow encountered in exp
  [l21,         np.exp(l22),      0.0],
C:\Users\thoma\AppData\Local\Temp\ipykernel_24368\3121696981.py:23: RuntimeWarning: invalid value encountered in matmul
  Q = L @ L.T
C:\ProgramData\anaconda3\Lib\site-packages\numpy\linalg\linalg.py:2120: RuntimeWarning: invalid value encountered in slogdet
  sign, logdet = _umath_linalg.slogdet(a, signature=signature)
C:\Users\thoma\AppData\Local\Temp\ipykernel_24368\609847046.py:60: RuntimeWarning: invalid value encountered in matmul
  K_t = P_pred @ H.T @ S_inv  # (3,n)

0.7075862068965517 nan
0.7486206896551725 -66660.78780290262
0.7896551724137931 -67227.74348501614
0.8306896551724139 -66442.50663102427
0.8717241379310345 -63184.87517358641
0.9127586206896552 -65779.09265378924
0.953793103448276 -65849.7416485883
0.9948275862068966 nan
1.0358620689655174 -63675.38250013805
1.076896551724138 -64866.1332875518
1.1179310344827587 -64639.29479314282
1.1589655172413793 -63772.95992611035
1.2 -55899.66809548512
Best lambda: 0.09206896551724138 neg loglik: -69908.23831762242

from scipy.linalg import logm

def compute_K_from_A(A, delta_t=1/12):
    """
    Given discrete-time A (3x3) and time step delta_t in years (monthly = 1/12),
    approximate continuous-time K via matrix logarithm.
    """
    A = np.asarray(A)
    K = - (1.0 / delta_t) * logm(A)  # <-- minus sign here
    K = np.real_if_close(K)
    return K

def afns_AB_grid(K, Q, delta0, delta1, theta, tau_grid, n_steps=200):
    """
    Compute A(tau), B(tau) on a grid of maturities tau_grid (in years)
    for an AFNS-like model with:
      dX_t = K (theta - X_t) dt + noise
      r_t = delta0 + delta1' X_t

    Uses simple Euler integration of the ODEs:
      dB/dtau = -K' B - delta1
      dA/dtau = -delta0 - (K theta)' B + 0.5 B' Q B

    K: (3,3)
    Q: (3,3) continuous-time state covariance
    delta0: scalar
    delta1: (3,) vector
    theta: (3,) long-run mean
    tau_grid: array of maturities in years
    n_steps: steps per year for numerical integration
    """
    K = np.asarray(K)
    Q = np.asarray(Q)
    delta1 = np.asarray(delta1).reshape(3,)
    theta = np.asarray(theta).reshape(3,)

    tau_grid = np.asarray(tau_grid)
    taus_sorted = np.sort(tau_grid)
    max_tau = taus_sorted[-1]

    dtau = 1.0 / n_steps  # step in years
    n_iter = int(max_tau / dtau) + 1

    B = np.zeros((3,))  # B(0)
    A = 0.0             # A(0)

    A_vals = {}
    B_vals = {}

    current_tau = 0.0
    idx_tau = 0

    K_T = K.T
    Ktheta = K @ theta

    for i in range(n_iter):
        # store values when we cross a tau in tau_grid
        while idx_tau < len(taus_sorted) and current_tau >= taus_sorted[idx_tau] - 1e-8:
            tau_val = taus_sorted[idx_tau]
            A_vals[tau_val] = A
            B_vals[tau_val] = B.copy()
            idx_tau += 1
            if idx_tau >= len(taus_sorted):
                break

        if idx_tau >= len(taus_sorted):
            break

        # ODEs:
        dB = -(K_T @ B) - delta1
        dA = -delta0 - (Ktheta @ B) + 0.5 * (B @ Q @ B)

        B = B + dB * dtau
        A = A + dA * dtau

        current_tau += dtau

    # Convert dicts to arrays aligned with original tau_grid order
    A_array = np.array([A_vals[tau] for tau in tau_grid])
    B_array = np.vstack([B_vals[tau] for tau in tau_grid])  # (n_tau, 3)

    return A_array, B_array

class AFNSFromDL:
    def __init__(self, A_P, Q_P, factors_df, maturities, delta0=None, delta1=None, delta_t=1/12):
        """
        A_P, Q_P: discrete-time DL MLE dynamics (3x3 each)
        factors_df: DataFrame with columns ['level','slope','curvature']
        maturities: array-like of maturities in years (e.g. [1.0, 2.0, ..., 30.0])
        delta0: scalar for short-rate intercept (if None, set to 0)
        delta1: length-3 array (if None, default [1,1,0])
        delta_t: time step in years for A_P (monthly = 1/12)
        """
        self.A_P = np.asarray(A_P)
        self.Q_P = np.asarray(Q_P)
        self.factors = factors_df
        self.maturities = np.asarray(maturities)
        self.delta_t = delta_t

        self.K = compute_K_from_A(self.A_P, delta_t=delta_t)
        # crude continuous-time Q: scale discrete Q by 1/delta_t
        self.Q_ct = self.Q_P * delta_t

        if delta1 is None:
            self.delta1 = np.array([1.0, 1.0, 0.0])
        else:
            self.delta1 = np.asarray(delta1).reshape(3,)

        if delta0 is None:
            self.delta0 = 0.0
        else:
            self.delta0 = float(delta0)

        # long-run mean theta: sample mean of factors
        self.theta = self.factors[["level","slope","curvature"]].mean().values

        # precompute A(tau), B(tau) and build linear mapping
        self.A_tau, self.B_tau = afns_AB_grid(
            self.K, self.Q_ct, self.delta0, self.delta1, self.theta,
            tau_grid=self.maturities
        )
        # Mapping: y_t(tau_i) = a_i + M_i dot X_t
        # with a_i = -A(tau_i)/tau_i, M_i = -B(tau_i)/tau_i
        self.a_vec = -self.A_tau / self.maturities
        self.M_mat = -self.B_tau / self.maturities[:, None]  # (n_tau, 3)

    def yields_from_factors(self, X_t):
        """
        Given X_t = [level, slope, curvature], return AFNS zero-coupon yields
        at all self.maturities.
        """
        X_t = np.asarray(X_t).reshape(3,)
        return self.a_vec + self.M_mat @ X_t

    def yields_from_path(self):
        """
        Apply AFNS mapping to the whole factor path.
        Returns DataFrame: index like factors_df, columns = maturities as strings.
        """
        X = self.factors[["level","slope","curvature"]].values  # (T,3)
        Y = X @ self.M_mat.T + self.a_vec  # (T, n_tau)
        cols = [f"{m:.1f}" for m in self.maturities]
        return pd.DataFrame(Y, index=self.factors.index, columns=cols)

# Suppose you have:
# A_hat, Q_hat from DL MLE
# factors_smooth_mle: DataFrame with level/slope/curvature
# maturities: e.g. np.array([1.0, 2.0, 3.0, 5.0, 7.0, 10.0, 20.0, 30.0])

afns_model = AFNSFromDL(
    A_P=A_hat,
    Q_P=Q_hat,
    factors_df=factors_smooth_mle,
    maturities=np.array([1.0, 2.0, 3.0, 5.0, 7.0, 10.0, 20.0, 30.0])
)

afns_yields = afns_model.yields_from_path()
afns_yields.head()

def ns_loadings(tau, lam):
    """
    Nelson–Siegel factor loadings:
    B1(tau) = 1
    B2(tau) = (1 - exp(-lam*tau)) / (lam*tau)
    B3(tau) = B2(tau) - exp(-lam*tau)
    """
    tau = np.asarray(tau, dtype=float)
    eps = 1e-8
    x = lam * np.maximum(tau, eps)
    exp_term = np.exp(-x)

    B1 = np.ones_like(tau)
    B2 = (1.0 - exp_term) / x
    B3 = B2 - exp_term
    return B1, B2, B3

def afns_yield_adjustment(tau, lam, sig1, sig2, sig3):
    """
    Independent-factor AFNS yield-adjustment term:
    C(t,T)/(T-t) as function of tau, lam, sigma1..3.
    """
    tau = np.asarray(tau, dtype=float)
    lam = float(lam)

    s1, s2, s3 = float(sig1), float(sig2), float(sig3)

    eps = 1e-8
    tau_safe = np.maximum(tau, eps)
    x = lam * tau_safe
    e1 = np.exp(-x)
    e2 = np.exp(-2 * x)

    # Level factor contribution
    I1 = (s1**2 / 6.0) * tau_safe**2

    # Slope factor contribution
    I2 = s2**2 * (
        1.0 / (2.0 * lam**2)
        - (1.0 / lam**3) * (1.0 - e1) / tau_safe
        + (1.0 / (4.0 * lam**3)) * (1.0 - e2) / tau_safe
    )

    # Curvature factor contribution
    I3 = s3**2 * (
        1.0 / (2.0 * lam**2)
        + (1.0 / lam**2) * e1
        - (1.0 / (4.0 * lam)) * tau_safe * e2
        - (3.0 / (4.0 * lam**2)) * e2
        - (2.0 / lam**3) * (1.0 - e1) / tau_safe
        + (5.0 / (8.0 * lam**3)) * (1.0 - e2) / tau_safe
    )

    return I1 + I2 + I3

def build_measurement_matrices(taus, lam, sig1, sig2, sig3):
    """
    Build:
    - H: N x 3 factor loading matrix
    - a: N-dimensional intercept vector (AFNS adj term)
    used in y_t = a + H X_t + eps_t
    """
    taus = np.asarray(taus, dtype=float)

    B1, B2, B3 = ns_loadings(taus, lam)
    H = np.column_stack([B1, B2, B3])

    C_adj = afns_yield_adjustment(taus, lam, sig1, sig2, sig3)
    a = -C_adj

    return H, a

def dns_yields_from_factors(X_t, taus, lam):
    """
    Produce DNS/Nelson–Siegel yields from factor vector X_t=[L,S,C]
    """
    X_t = np.asarray(X_t, dtype=float)
    L_t, S_t, C_t = X_t
    B1, B2, B3 = ns_loadings(taus, lam)
    return L_t * B1 + S_t * B2 + C_t * B3


def afns_yields_from_factors(X_t, taus, lam, sig1, sig2, sig3):
    """
    AFNS yields = DNS yields - no-arbitrage adjustment
    """
    dns = dns_yields_from_factors(X_t, taus, lam)
    adj = afns_yield_adjustment(taus, lam, sig1, sig2, sig3)
    return dns - adj

def unpack_theta(theta):
    """
    Unpack the 14-parameter AFNS reduced-form vector.
    """
    theta = np.asarray(theta, dtype=float)

    phi_L, phi_S, phi_C = theta[0:3]
    mu_L,  mu_S,  mu_C  = theta[3:6]
    log_qL, log_qS, log_qC = theta[6:9]
    log_lam = theta[9]
    log_sig1, log_sig2, log_sig3 = theta[10:13]
    log_r = theta[13]

    Phi = np.diag([phi_L, phi_S, phi_C])
    mu  = np.array([mu_L, mu_S, mu_C])
    Q   = np.diag([np.exp(log_qL)**2,
                   np.exp(log_qS)**2,
                   np.exp(log_qC)**2])

    lam = np.exp(log_lam)
    sig1, sig2, sig3 = np.exp(log_sig1), np.exp(log_sig2), np.exp(log_sig3)
    r = np.exp(log_r)

    return Phi, mu, Q, lam, sig1, sig2, sig3, r

class CurveModelAFNS:
    """
    DNS/AFNS term structure model with independent AR(1) P-dynamics
    and AFNS no-arbitrage adjustment in measurement eq.
    """

    def __init__(self, theta_hat, taus):
        """
        theta_hat: estimated 14-parameter vector
        taus: maturities in years (array-like)
        """
        self.theta = np.asarray(theta_hat, dtype=float)
        self.taus = np.asarray(taus, dtype=float)

        (
            self.Phi,
            self.mu,
            self.Q,
            self.lam,
            self.sig1,
            self.sig2,
            self.sig3,
            self.r,
        ) = unpack_theta(self.theta)

        # Build AFNS measurement structures
        self.H, self.a = build_measurement_matrices(
            self.taus, self.lam, self.sig1, self.sig2, self.sig3
        )
        self.R = (self.r**2) * np.eye(len(self.taus))

    # ---------- simulators ----------

    def simulate_factors(self, T, x0=None, rng=None):
        """
        Simulate factor path X_t under P-dynamics (AR(1)) for t=0..T-1.
        Returns array (T, 3).
        """
        if rng is None:
            rng = np.random.default_rng()

        if x0 is None:
            x0 = self.mu.copy()

        X = np.zeros((T, 3))
        X[0] = x0

        for t in range(1, T):
            eps = rng.multivariate_normal(mean=np.zeros(3), cov=self.Q)
            X[t] = self.mu + self.Phi @ (X[t-1] - self.mu) + eps

        return X

    def simulate_yields(self, X, model="afns"):
        """
        Given factor path X (T,3), return yields (T, N) under DNS or AFNS.
        """
        X = np.asarray(X, dtype=float)
        T = X.shape[0]
        N = len(self.taus)
        Y = np.zeros((T, N))

        for t in range(T):
            if model == "dns":
                Y[t] = dns_yields_from_factors(X[t], self.taus, self.lam)
            elif model == "afns":
                Y[t] = afns_yields_from_factors(
                    X[t], self.taus, self.lam, self.sig1, self.sig2, self.sig3
                )
            else:
                raise ValueError("model must be 'dns' or 'afns'")

        return Y

def kalman_loglik_afns(theta, Y, taus, beta0=None, P0=None):
    """
    AFNS Kalman log-likelihood.

    Model:
      X_t - mu = Phi (X_{t-1} - mu) + eta_t,  eta_t ~ N(0, Q)
      y_t = a + H X_t + eps_t,                 eps_t ~ N(0, R)

    where (Phi, mu, Q, lam, sig1..3, r) = unpack_theta(theta)
    and H, a are built via AFNS (no-arbitrage adjustment).

    Parameters
    ----------
    theta : array-like, shape (14,)
        Parameter vector as defined in unpack_theta.
    Y : array-like, shape (T, N)
        Observed yields (T time points, N maturities).
    taus : array-like, shape (N,)
        Maturities in years corresponding to columns of Y.
    beta0 : array-like, shape (3,), optional
        Initial state mean; if None, we use mu.
    P0 : array-like, shape (3, 3), optional
        Initial state covariance; if None, we use 0.1 * I.

    Returns
    -------
    neg_loglik : float
        Negative log-likelihood (for minimization).
    """
    Y = np.asarray(Y, dtype=float)
    T, N = Y.shape
    taus = np.asarray(taus, dtype=float)

    # Unpack parameters
    Phi, mu, Q, lam, sig1, sig2, sig3, r = unpack_theta(theta)

    # Measurement matrices (AFNS)
    H, a = build_measurement_matrices(taus, lam, sig1, sig2, sig3)
    R = (r**2) * np.eye(N)

    # Adjust observations to absorb intercept: y'_t = y_t - a
    Y_adj = Y - a[None, :]

    k = 3  # number of factors
    if beta0 is None:
        beta0 = mu.copy()
    if P0 is None:
        P0 = 0.1 * np.eye(k)

    beta_prev = beta0
    P_prev = P0
    I_k = np.eye(k)

    loglik = 0.0
    const = N * np.log(2 * np.pi)

    for t in range(T):
        # Prediction step: X_t|t-1
        beta_pred = mu + Phi @ (beta_prev - mu)
        P_pred = Phi @ P_prev @ Phi.T + Q

        # Observation for this time
        y_t = Y_adj[t, :]          # (N,)

        # Innovation covariance
        S_t = H @ P_pred @ H.T + R  # (N, N)

        try:
            S_inv = np.linalg.inv(S_t)
            sign, logdet = np.linalg.slogdet(S_t)
            if sign <= 0:
                # non-PD covariance → penalize
                return 1e6
        except np.linalg.LinAlgError:
            return 1e6

        # Innovation
        y_hat = H @ beta_pred      # (N,)
        innov = y_t - y_hat        # (N,)

        quad = innov.T @ S_inv @ innov
        loglik_t = -0.5 * (const + logdet + quad)
        loglik += loglik_t

        # Update step
        K_t = P_pred @ H.T @ S_inv      # (3, N)
        beta_filt = beta_pred + K_t @ innov
        P_filt = (I_k - K_t @ H) @ P_pred

        beta_prev, P_prev = beta_filt, P_filt

    # Return negative log-likelihood for minimization
    return -float(loglik)

def make_initial_theta(Y, taus):
    """
    Construct a rough initial guess for the 14-parameter AFNS vector.
    Y: (T, N) yields
    taus: (N,) maturities

    Returns
    -------
    theta0 : np.ndarray, shape (14,)
    """
    Y = np.asarray(Y, dtype=float)
    T, N = Y.shape

    # crude guesses
    # factor means: use average of shortest maturity for level, 0 for slope/curvature
    mu_L0 = float(Y[:, 0].mean())
    mu_S0 = 0.0
    mu_C0 = 0.0

    # AR coefficients: persistent level, less for slope/curvature
    phi_L0, phi_S0, phi_C0 = 0.98, 0.90, 0.80

    # state noise std devs (log scale)
    log_qL0 = np.log(0.01)
    log_qS0 = np.log(0.02)
    log_qC0 = np.log(0.02)

    # lambda around typical NS values (e.g. DL ~ 0.06–0.1 for monthly)
    log_lam0 = np.log(0.06)

    # AFNS vol parameters (continuous-time vols for level, slope, curvature)
    log_sig1_0 = np.log(0.01)
    log_sig2_0 = np.log(0.02)
    log_sig3_0 = np.log(0.02)

    # measurement noise std dev
    log_r0 = np.log(0.001)

    theta0 = np.array([
        phi_L0, phi_S0, phi_C0,
        mu_L0,  mu_S0,  mu_C0,
        log_qL0, log_qS0, log_qC0,
        log_lam0,
        log_sig1_0, log_sig2_0, log_sig3_0,
        log_r0
    ], dtype=float)

    return theta0


def fit_afns_mle(Y, taus, theta0=None, maxiter=300):
    """
    Estimate AFNS parameters by MLE via Kalman filter.

    Parameters
    ----------
    Y : (T, N) array
        Yield panel (time x maturities).
    taus : (N,) array
        Maturities in years (aligned with columns of Y).
    theta0 : array-like, optional
        Initial guess for parameters; if None, we use make_initial_theta().
    maxiter : int
        Maximum number of optimizer iterations.

    Returns
    -------
    theta_hat : np.ndarray
        Estimated parameter vector.
    res : OptimizeResult
        Full scipy.optimize result object.
    """
    Y = np.asarray(Y, dtype=float)
    taus = np.asarray(taus, dtype=float)

    if theta0 is None:
        theta0 = make_initial_theta(Y, taus)

    def objective(theta):
        return kalman_loglik_afns(theta, Y, taus)

    res = minimize(
        objective,
        theta0,
        method="L-BFGS-B",
        options={"maxiter": maxiter, "disp": True}
    )

    theta_hat = res.x
    return theta_hat, res

# Choose maturities (columns must exist in df_yields)
cols = ["1.0", "2.0", "3.0", "5.0", "7.0", "10.0", "20.0", "30.0"]
taus = np.array([float(c) for c in cols])

Y = df_yields[cols].values  # shape (T, N)

theta0 = make_initial_theta(Y, taus)
theta_hat, res = fit_afns_mle(Y, taus, theta0=theta0, maxiter=300)

print("Converged:", res.success)
print("Final negative log-likelihood:", res.fun)
print("Estimated parameters:", theta_hat)

Converged: True
Final negative log-likelihood: -17606.123533753547
Estimated parameters: [ 0.99102606  0.98948413  0.95145194  0.05888597 -0.04058459 -0.02338965
 -6.27675244 -5.83087401 -5.1210678  -1.25079721 -5.17113235 -3.66618834
 -4.41114251 -6.91906477]

cm = CurveModelAFNS(theta_hat, taus)

# simulate 10 years monthly
T = 10 * 12
X_sim = cm.simulate_factors(T)
Y_dns = cm.simulate_yields(X_sim, model="dns")
Y_afns = cm.simulate_yields(X_sim, model="afns")

plt.plot(Y_afns)

[<matplotlib.lines.Line2D at 0x202c7046cf0>,
 <matplotlib.lines.Line2D at 0x202c6f31d90>,
 <matplotlib.lines.Line2D at 0x202bfd58b90>,
 <matplotlib.lines.Line2D at 0x202c6f7bf20>,
 <matplotlib.lines.Line2D at 0x202c7045f70>,
 <matplotlib.lines.Line2D at 0x202c7046ed0>,
 <matplotlib.lines.Line2D at 0x202c7046fc0>,
 <matplotlib.lines.Line2D at 0x202c70470b0>]

def kalman_filter_afns(theta, Y, taus, beta0=None, P0=None):
    """
    Run the Kalman filter for the AFNS Option-B model.

    Model:
      X_t - mu = Phi (X_{t-1} - mu) + eta_t,  eta_t ~ N(0, Q)
      y_t = a + H X_t + eps_t,                eps_t ~ N(0, R)

    We absorb the intercept a into the observations:
      y'_t = y_t - a = H X_t + eps_t.

    Parameters
    ----------
    theta : array-like, shape (14,)
        Parameter vector as in unpack_theta.
    Y : (T, N) array
        Yield panel.
    taus : (N,) array
        Maturities in years.
    beta0 : (3,) array, optional
        Initial state mean; default = mu.
    P0 : (3,3) array, optional
        Initial state covariance; default = 0.1 * I.

    Returns
    -------
    filt_means : (T, 3)
        Filtered state means E[X_t | Y_1..t].
    filt_covs  : (T, 3, 3)
        Filtered covariance matrices.
    pred_means : (T, 3)
        One-step-ahead predicted means E[X_t | Y_1..t-1].
    pred_covs  : (T, 3, 3)
        One-step-ahead predicted covariances.
    loglik : float
        Total log-likelihood (same as in MLE, for reference).
    extra : dict
        Dict with (Phi, mu, Q, H, a, R) for reuse in smoother, plotting, etc.
    """
    Y = np.asarray(Y, dtype=float)
    taus = np.asarray(taus, dtype=float)
    T, N = Y.shape

    Phi, mu, Q, lam, sig1, sig2, sig3, r = unpack_theta(theta)
    H, a = build_measurement_matrices(taus, lam, sig1, sig2, sig3)
    R = (r**2) * np.eye(N)

    # absorb intercept
    Y_adj = Y - a[None, :]

    k = 3
    if beta0 is None:
        beta0 = mu.copy()
    if P0 is None:
        P0 = 0.1 * np.eye(k)

    filt_means = np.zeros((T, k))
    filt_covs  = np.zeros((T, k, k))
    pred_means = np.zeros((T, k))
    pred_covs  = np.zeros((T, k, k))

    beta_prev = beta0
    P_prev = P0
    I_k = np.eye(k)

    loglik = 0.0
    const = N * np.log(2 * np.pi)

    for t in range(T):
        # prediction
        beta_pred = mu + Phi @ (beta_prev - mu)
        P_pred = Phi @ P_prev @ Phi.T + Q

        pred_means[t] = beta_pred
        pred_covs[t] = P_pred

        y_t = Y_adj[t, :]          # (N,)
        S_t = H @ P_pred @ H.T + R # (N,N)

        # innovation covariance must be PD
        try:
            S_inv = np.linalg.inv(S_t)
            sign, logdet = np.linalg.slogdet(S_t)
            if sign <= 0:
                raise np.linalg.LinAlgError("Non-PD innovation covariance")
        except np.linalg.LinAlgError:
            # error handling
            return None, None, None, None, -np.inf, {}

        y_hat = H @ beta_pred      # (N,)
        innov = y_t - y_hat        # (N,)

        quad = innov.T @ S_inv @ innov
        loglik_t = -0.5 * (const + logdet + quad)
        loglik += loglik_t

        # update
        K_t = P_pred @ H.T @ S_inv     # (3,N)
        beta_filt = beta_pred + K_t @ innov
        P_filt = (I_k - K_t @ H) @ P_pred

        filt_means[t] = beta_filt
        filt_covs[t] = P_filt

        beta_prev, P_prev = beta_filt, P_filt

    extra = {
        "Phi": Phi,
        "mu": mu,
        "Q": Q,
        "H": H,
        "a": a,
        "R": R,
    }

    return filt_means, filt_covs, pred_means, pred_covs, float(loglik), extra

def rts_smoother_afns(filt_means, filt_covs, pred_means, pred_covs, Phi):
    """
    Rauch–Tung–Striebel smoother for AFNS model.

    Parameters
    ----------
    filt_means : (T, 3)
        Filtered means from Kalman filter.
    filt_covs  : (T, 3, 3)
        Filtered covariances.
    pred_means : (T, 3)
        One-step-ahead predicted means.
    pred_covs  : (T, 3, 3)
        One-step-ahead predicted covariances.
    Phi : (3,3)
        State transition matrix (constant over time in this model).

    Returns
    -------
    smooth_means : (T, 3)
        Smoothed state means E[X_t | Y_1..T].
    smooth_covs  : (T, 3, 3)
        Smoothed state covariances.
    """
    filt_means = np.asarray(filt_means, dtype=float)
    filt_covs  = np.asarray(filt_covs, dtype=float)
    pred_means = np.asarray(pred_means, dtype=float)
    pred_covs  = np.asarray(pred_covs, dtype=float)

    T, k = filt_means.shape
    smooth_means = np.zeros_like(filt_means)
    smooth_covs  = np.zeros_like(filt_covs)

    # initialize at T-1
    smooth_means[-1] = filt_means[-1]
    smooth_covs[-1]  = filt_covs[-1]

    Phi_T = Phi.T

    for t in range(T - 2, -1, -1):
        P_filt_t = filt_covs[t]
        P_pred_next = pred_covs[t + 1]

        # smoother gain
        J_t = P_filt_t @ Phi_T @ np.linalg.inv(P_pred_next)

        # smoothed mean
        smooth_means[t] = (
            filt_means[t]
            + J_t @ (smooth_means[t + 1] - pred_means[t + 1])
        )

        # smoothed covariance
        smooth_covs[t] = (
            P_filt_t
            + J_t @ (smooth_covs[t + 1] - P_pred_next) @ J_t.T
        )

    return smooth_means, smooth_covs

# theta_hat from MLE, Y and taus from GSW data
filt_means, filt_covs, pred_means, pred_covs, loglik, extra = kalman_filter_afns(
    theta_hat, Y, taus
)

Phi = extra["Phi"]

smooth_means, smooth_covs = rts_smoother_afns(
    filt_means, filt_covs, pred_means, pred_covs, Phi
)

# smooth_means is (T,3): AFNS factors L_t, S_t, C_t
L = smooth_means[:, 0]
S = smooth_means[:, 1]
C = smooth_means[:, 2]

plt.plot(filt_means)
plt.legend(["Level", "Slope", "Curvature"])

<matplotlib.legend.Legend at 0x202c7047ec0>

plt.plot(smooth_means)
plt.legend(["Level", "Slope", "Curvature"])

<matplotlib.legend.Legend at 0x202bfd71d30>

# ============================================================
# 0) USER INPUTS
# ============================================================

# Required:
# - X_smooth: array (T, 3) of Kalman-smoothed AFNS factors [L,S,C]
X_smooth = smooth_means
# - cm: calibrated CurveModelAFNS (needs lam, sig1..sig3 and AFNS yield function)

# Banking book + frequency
dt = 1.0 / 12.0         # monthly in years
tau_A = 3.0             # assets repricing maturity (years)
tau_L = 1.0             # liabilities repricing maturity (years)

# ---- Banking-book repricing gap -------------------------------
# Liability-sensitive book: rate-sensitive liabilities reprice faster/larger than
# rate-sensitive assets, so a parallel rate rise compresses NII (nonzero LEVEL
# sensitivity).
A_NOTIONAL = 60.0      # rate-sensitive assets, reprice at tau_A
L_NOTIONAL = 100.0     # rate-sensitive liabilities, reprice at tau_L
U_SCALE    = 25.0      # action scale: Delta h = u_raw * U_SCALE, u_raw in [-1,1]
A_notional = A_NOTIONAL
L_notional = L_NOTIONAL

# LQ weights (tune later)
alpha_nii = 1.0          # strength of "penalize NII" term
lambda_u = 1e-3          # trading penalty (smaller => more aggressive hedging)
lambda_h = 1e-6          # hedge inventory penalty (optional)

# Stress scenario
shock_bps = 200          # +200 bps
shock = shock_bps / 10000.0

# ============================================================
# 1) AFNS yield + forward-rate utilities
# ============================================================

def afns_yield_single_from_cm(cm, X, tau):
    """
    Compute AFNS yield y(tau;X) using cm parameters.
    Depends on 
    existing afns_yields_from_factors implementation.
    """
    taus = np.array([tau], dtype=float)
    y = afns_yields_from_factors(X, taus, cm.lam, cm.sig1, cm.sig2, cm.sig3)
    return float(y[0])


def forward_rate_cc_from_cm(cm, X, tau1, tau2):
    """
    Continuous-compounded forward rate f(tau1,tau2):
      f = (tau2*y(tau2) - tau1*y(tau1)) / (tau2 - tau1)
    """
    y1 = afns_yield_single_from_cm(cm, X, tau1)
    y2 = afns_yield_single_from_cm(cm, X, tau2)
    return (tau2 * y2 - tau1 * y1) / (tau2 - tau1)


# ============================================================
# 2) NII definition with FRA hedge
# ============================================================

def compute_unhedged_nii_path(cm, X_path, A_notional, L_notional, tau_A, tau_L, dt):
    """
    Unhedged NII_{t+1} = A*y_t(tauA)*dt - L*y_t(tauL)*dt
    Returns array length T-1.
    """
    T = X_path.shape[0]
    NII0 = np.zeros(T-1)
    for t in range(T-1):
        yA = afns_yield_single_from_cm(cm, X_path[t], tau_A)
        yL = afns_yield_single_from_cm(cm, X_path[t], tau_L)
        NII0[t] = A_notional * yA * dt - L_notional * yL * dt
    return NII0


def compute_hedged_nii_path_FRA(cm, X_path, h_path, A_notional, L_notional, tau_A, tau_L, dt):
    """
    Hedged NII_{t+1} = A*y_t(tauA)*dt - L*y_t(tauL)*dt + h_t*(K_t - y_{t+1}(tauL))*dt
    with K_t = forward(tauL, tauL+dt).
    """
    T = X_path.shape[0]
    NIIh = np.zeros(T-1)
    for t in range(T-1):
        yA = afns_yield_single_from_cm(cm, X_path[t], tau_A)
        yL = afns_yield_single_from_cm(cm, X_path[t], tau_L)

        # Issue 10: pay-fixed swaplet on the current reference rate y_t(tau_L),
        # h_t*(y_t(tau_L) - K_STRIKE)*dt, offsetting the contemporaneous liability
        # rate exposure (effective at sensible notional ~ H_VM).
        NIIh[t] = (A_notional * yA * dt
                   - L_notional * yL * dt
                   + h_path[t] * (yL - K_STRIKE) * dt)
    return NIIh


# ============================================================
# 3) Estimate factor dynamics from smoothed factors (AR(1))
#    X_{t+1} = c + Phi X_t + eps
#    We'll convert to mean-reverting form with mu if desired.
# ============================================================

def fit_var1(X):
    """
    Fit VAR(1): X_{t+1} = c + Phi X_t + eps, via OLS.
    Returns c (3,), Phi (3,3), Sigma (3,3).
    """
    X = np.asarray(X, dtype=float)
    Y = X[1:]             # (T-1,3)
    Z = X[:-1]            # (T-1,3)

    # add intercept
    Z1 = np.column_stack([np.ones(Z.shape[0]), Z])  # (T-1, 1+3)

    # OLS for each equation
    B = np.linalg.lstsq(Z1, Y, rcond=None)[0]        # (1+3, 3)
    c = B[0]                                         # (3,)
    Phi = B[1:].T                                    # (3,3) because (3x3)

    resid = Y - Z1 @ B                               # (T-1,3)
    Sigma = (resid.T @ resid) / (resid.shape[0] - (1 + 3))  # sample cov

    return c, Phi, Sigma


# ============================================================
# 4) Build LQ problem from scratch
#    State x_t = [L,S,C,h]
#    Dynamics:  X_{t+1} = c + Phi X_t + eps,  h_{t+1} = h_t + u_t
#
#    Objective: penalize (approx NII)^2 + lambda_h h^2 + lambda_u u^2
#
#    Key step: build H_x (sensitivity of NII to state) numerically
# ============================================================

def numerical_grad_y(cm, X_ref, tau, eps=1e-5):
    """
    Numerical gradient of y(tau;X) wrt X=(L,S,C) using central differences.
    Returns grad (3,).
    """
    grad = np.zeros(3)
    for i in range(3):
        d = np.zeros(3)
        d[i] = eps
        yp = afns_yield_single_from_cm(cm, X_ref + d, tau)
        ym = afns_yield_single_from_cm(cm, X_ref - d, tau)
        grad[i] = (yp - ym) / (2 * eps)
    return grad


def build_Hx_QR_from_nii(cm, X_ref,
                         A_notional, L_notional, tau_A, tau_L, dt,
                         alpha_nii, lambda_h, lambda_u):
    """
    Build H_x, Q_s, R for LQ:
      approx NII(x_t) ≈ H_x' [X_t; h_t]
      => (NII)^2 ≈ x' (alpha * H_x H_x') x
    """
    # Factor sensitivity of the base NII part (using numerical gradients)
    grad_yA = numerical_grad_y(cm, X_ref, tau_A)
    grad_yL = numerical_grad_y(cm, X_ref, tau_L)

    Hx_X = dt * (A_notional * grad_yA - L_notional * grad_yL)  # (3,)

    # Hedge = pay-fixed swaplet. Mean hedge sensitivity at the reference
    # is ~0 (struck at the par rate); the hedge's VALUE is variance reduction. The NII
    # factor loading l(h)=dt*(A*grad_yA-(L-h)*grad_yL) has risk l(h)'Sigma l(h)
    # minimized at h=H_VM. Locally risk ~ const + dt^2 (grad_yL'Sigma grad_yL)(h-H_VM)^2.
    # We penalize (h-H_VM)^2 -- the hedge state is shifted to (h - H_VM) in
    # run_lq_hedge_on_path / the rollout LQ branch.
    yL_ref = afns_yield_single_from_cm(cm, X_ref, tau_L)
    d_h = (yL_ref - K_STRIKE) * dt   # ~0 when struck at the par reference rate

    H_x = np.zeros(4)
    H_x[:3] = Hx_X
    H_x[3] = d_h

    q_var = alpha_nii * (dt**2) * float(grad_yL @ Sigma_hat @ grad_yL)

    # Quadratic cost (state = [X - X_ref ; h - H_VM])
    Q_s = alpha_nii * np.outer(H_x, H_x)
    Q_s[3, 3] = Q_s[3, 3] + q_var
    if lambda_h > 0:
        Q_s[3, 3] = Q_s[3, 3] + lambda_h

    R = np.array([[lambda_u]], dtype=float)

    return H_x, Q_s, R


def build_AB_from_Phi(Phi):
    A = np.zeros((4, 4))
    A[:3, :3] = Phi
    A[3, 3] = 1.0
    B = np.zeros((4, 1))
    B[3, 0] = 1.0
    return A, B


def solve_discrete_riccati(A, B, Q, R, max_iter=20000, tol=1e-12):
    """
    Iterative solution to discrete algebraic Riccati equation (DARE).
    Returns P, K where u_t = -K x_t.
    """
    A = np.asarray(A, float)
    B = np.asarray(B, float)
    Q = np.asarray(Q, float)
    R = np.asarray(R, float)

    P = Q.copy()
    for _ in range(max_iter):
        S = R + B.T @ P @ B
        K = np.linalg.solve(S, B.T @ P @ A)
        P_next = Q + A.T @ P @ A - A.T @ P @ B @ K
        if np.max(np.abs(P_next - P)) < tol:
            P = P_next
            break
        P = P_next

    S = R + B.T @ P @ B
    K = np.linalg.solve(S, B.T @ P @ A)
    return P, K


# ============================================================
# 5) Run LQ hedging on an exogenous factor path
# ============================================================

def run_lq_hedge_on_path(cm, X_path, Phi, H_x, Q_s, R,
                         A_notional, L_notional, tau_A, tau_L, dt):
    """
    Given an exogenous factor path X_path and VAR(1) Phi (for control design),
    compute LQ hedge policy and resulting hedge path h_t and NII.
    """
    A, B = build_AB_from_Phi(Phi)
    P, K = solve_discrete_riccati(A, B, Q_s, R)

    # simulate hedge
    T = X_path.shape[0]
    h = np.zeros(T)
    u = np.zeros(T-1)

    # Run the regulator on the CENTERED factor state X - X_ref, so the
    # LQ policy stabilizes NII around its mean (X_ref) instead of driving it to zero.
    for t in range(T-1):
        x_t = np.array([X_path[t,0] - X_ref[0],
                        X_path[t,1] - X_ref[1],
                        X_path[t,2] - X_ref[2],
                        h[t] - H_VM])   # Target the variance-min hedge ratio
        u_t = -float((K @ x_t).item())
        u[t] = u_t
        h[t+1] = h[t] + u_t

    NII_unhedged = compute_unhedged_nii_path(cm, X_path, A_notional, L_notional, tau_A, tau_L, dt)
    NII_hedged = compute_hedged_nii_path_FRA(cm, X_path, h, A_notional, L_notional, tau_A, tau_L, dt)

    return {
        "K": K,
        "h": h,
        "u": u,
        "NII_unhedged": NII_unhedged,
        "NII_hedged": NII_hedged,
    }


# ============================================================
# 6) Stress path builder (shock at t=0, then propagate with Phi, no noise)
# ============================================================

def make_stress_path_from_var1(X0, c, Phi, T, shock_vec=None):
    """
    Deterministic stressed path:
      X_0 = X0 + shock_vec
      X_{t+1} = c + Phi X_t
    """
    X = np.zeros((T, 3))
    if shock_vec is None:
        shock_vec = np.zeros(3)
    X[0] = X0 + shock_vec
    for t in range(T-1):
        X[t+1] = c + Phi @ X[t]
    return X


# ============================================================
# 7) MAIN RUN: baseline + stress
# ============================================================

# --- Fit factor dynamics from smoothed factors ---
c_hat, Phi_hat, Sigma_hat = fit_var1(X_smooth)
X_ref = X_smooth.mean(axis=0)

# The objective penalizes (NII - NII_TARGET)^2 (variance).
NII_TARGET = (A_notional * afns_yield_single_from_cm(cm, X_ref, tau_A) * dt
              - L_notional * afns_yield_single_from_cm(cm, X_ref, tau_L) * dt)
print("NII_TARGET (expected NII at X_ref):", NII_TARGET)

# Hedge = pay-fixed swaplet on the CURRENT reference rate: h*(y_t(tau_L) - K_STRIKE)*dt.
# Struck at the par reference rate K_STRIKE so the mean hedge carry is ~0.
# H_VM is the closed-form variance-minimizing hedge ratio (regression beta in the
# Sigma metric): the NII factor loading is l(h)=dt*(A*gA - (L-h)*gL), minimizing
# l(h)'Sigma l(h) over h gives h = L - A*(gL'Sigma gA)/(gL'Sigma gL).
_gA = numerical_grad_y(cm, X_ref, tau_A)
_gL = numerical_grad_y(cm, X_ref, tau_L)
K_STRIKE = afns_yield_single_from_cm(cm, X_ref, tau_L)
H_VM = L_notional - A_notional * float((_gL @ Sigma_hat @ _gA) / (_gL @ Sigma_hat @ _gL))
print("Variance-min hedge ratio H_VM:", H_VM, "| strike K_STRIKE:", K_STRIKE)

print("Fitted VAR(1) Phi:\n", Phi_hat)
print("Reference X_ref (mean of smoothed):", X_ref)

# --- Build LQ cost from NII sensitivities ---
H_x, Q_s, R = build_Hx_QR_from_nii(
    cm, X_ref,
    A_notional, L_notional, tau_A, tau_L, dt,
    alpha_nii, lambda_h, lambda_u
)

print("H_x =", H_x)
print("Q_s max abs =", np.max(np.abs(Q_s)))
print("R =", R)

# --- Baseline: use the observed smoothed factor path as exogenous ---
res_base = run_lq_hedge_on_path(
    cm, X_smooth, Phi_hat, H_x, Q_s, R,
    A_notional, L_notional, tau_A, tau_L, dt
)

print("LQ gain K =", res_base["K"])
print("max |h| baseline:", np.max(np.abs(res_base["h"])))
print("std NII unhedged baseline:", np.std(res_base["NII_unhedged"]))
print("std NII hedged baseline:", np.std(res_base["NII_hedged"]))

plt.figure(figsize=(9,4))
plt.plot(res_base["NII_unhedged"], label="Unhedged (baseline)")
plt.plot(res_base["NII_hedged"], label="LQ hedged (baseline)")
plt.title("Baseline path (smoothed factors): Unhedged vs LQ-hedged NII")
plt.xlabel("t (months)")
plt.ylabel("NII")
plt.legend()
plt.tight_layout()
plt.show()

plt.figure(figsize=(9,3))
plt.plot(res_base["h"], label="Hedge notional h_t (baseline)")
plt.title("LQ hedge position (baseline)")
plt.xlabel("t (months)")
plt.ylabel("h_t")
plt.legend()
plt.tight_layout()
plt.show()


# --- Stress: +200bp parallel shock implemented as a level-factor bump at t=0 ---
# You can also shock slope/curvature; start with level.
shock_vec = np.array([shock, 0.0, 0.0])

X_stress = make_stress_path_from_var1(
    X0=X_ref, c=c_hat, Phi=Phi_hat, T=X_smooth.shape[0], shock_vec=shock_vec
)

res_stress = run_lq_hedge_on_path(
    cm, X_stress, Phi_hat, H_x, Q_s, R,
    A_notional, L_notional, tau_A, tau_L, dt
)

print("\n--- STRESS RESULTS (+200bp level shock at t=0, deterministic VAR1 propagation) ---")
print("LQ gain K =", res_stress["K"])
print("max |h| stress:", np.max(np.abs(res_stress["h"])))
print("std NII unhedged stress:", np.std(res_stress["NII_unhedged"]))
print("std NII hedged stress:", np.std(res_stress["NII_hedged"]))

plt.figure(figsize=(9,4))
plt.plot(res_stress["NII_unhedged"], label="Unhedged (stress)")
plt.plot(res_stress["NII_hedged"], label="LQ hedged (stress)")
plt.title("Stress (+200bp level shock at t=0): Unhedged vs LQ-hedged NII")
plt.xlabel("t (months)")
plt.ylabel("NII")
plt.legend()
plt.tight_layout()
plt.show()

plt.figure(figsize=(9,3))
plt.plot(res_stress["h"], label="Hedge notional h_t (stress)")
plt.title("LQ hedge position under stress")
plt.xlabel("t (months)")
plt.ylabel("h_t")
plt.legend()
plt.tight_layout()
plt.show()

NII_TARGET (expected NII at X_ref): -0.08085383350892353
Variance-min hedge ratio H_VM: 41.09175222568043 | strike K_STRIKE: 0.030161579251736695
Fitted VAR(1) Phi:
 [[ 0.9938804   0.00462503 -0.00218918]
 [-0.01640262  0.99214422  0.01098002]
 [ 0.01286072 -0.0127359   0.95747698]]
Reference X_ref (mean of smoothed): [ 0.06421823 -0.03616098 -0.02123435]
H_x = [-3.33333333 -3.89124429  0.24925829  0.        ]
Q_s max abs = 15.14178216149656
R = [[0.001]]
LQ gain K = [[0.         0.         0.         0.03172053]]
max |h| baseline: 41.091715407056284
std NII unhedged baseline: 0.0903518153735064
std NII hedged baseline: 0.03459803902804702

--- STRESS RESULTS (+200bp level shock at t=0, deterministic VAR1 propagation) ---
LQ gain K = [[0.         0.         0.         0.03172053]]
max |h| stress: 41.091715407056284
std NII unhedged stress: 0.02002043748709172
std NII hedged stress: 0.011302768582036403

def summarize_paths(N0, N1, h, u):
    """
    Produce standard metrics for hedging quality and cost.
    N0: unhedged NII (T-1,)
    N1: hedged NII (T-1,)
    h: hedge notional (T,)
    u: hedge trades (T-1,)
    """
    out = {}
    out["std_unhedged"] = float(np.std(N0))
    out["std_hedged"] = float(np.std(N1))
    out["std_reduction_%"] = float(100.0 * (1.0 - out["std_hedged"] / out["std_unhedged"])) if out["std_unhedged"] > 0 else np.nan

    out["min_unhedged"] = float(np.min(N0))
    out["min_hedged"] = float(np.min(N1))
    out["p05_unhedged"] = float(np.quantile(N0, 0.05))
    out["p05_hedged"] = float(np.quantile(N1, 0.05))

    out["mean_unhedged"] = float(np.mean(N0))
    out["mean_hedged"] = float(np.mean(N1))

    # “Cost” proxies
    out["mean_abs_h"] = float(np.mean(np.abs(h)))
    out["max_abs_h"] = float(np.max(np.abs(h)))
    out["mean_abs_u"] = float(np.mean(np.abs(u)))
    out["max_abs_u"] = float(np.max(np.abs(u)))

    return out


def sweep_lq_hyperparams(cm, X_path, Phi_hat, X_ref,
                         A_notional, L_notional, tau_A, tau_L, dt,
                         alpha_nii,
                         lambda_u_grid,
                         lambda_h_grid):
    """
    Sweeps (lambda_u, lambda_h) and returns a DataFrame of metrics.

    Requires build_Hx_QR_from_nii(...) and run_lq_hedge_on_path(...).
    """
    rows = []

    for lambda_u in lambda_u_grid:
        for lambda_h in lambda_h_grid:

            H_x, Q_s, R = build_Hx_QR_from_nii(
                cm, X_ref,
                A_notional, L_notional, tau_A, tau_L, dt,
                alpha_nii=alpha_nii,
                lambda_h=lambda_h,
                lambda_u=lambda_u
            )

            res = run_lq_hedge_on_path(
                cm, X_path, Phi_hat, H_x, Q_s, R,
                A_notional, L_notional, tau_A, tau_L, dt
            )

            metrics = summarize_paths(
                res["NII_unhedged"], res["NII_hedged"], res["h"], res["u"]
            )

            row = {
                "lambda_u": float(lambda_u),
                "lambda_h": float(lambda_h),
                "Hx_hedge_sensitivity": float(H_x[3]),
                "K_max_abs": float(np.max(np.abs(res["K"]))),
                **metrics
            }
            rows.append(row)

    df = pd.DataFrame(rows)

    # useful derived columns
    df["turnover_proxy"] = df["mean_abs_u"]      # rename, but keep explicit too
    df["inventory_proxy"] = df["mean_abs_h"]

    # Sort by best hedging first (std reduction), then lower turnover
    df = df.sort_values(["std_reduction_%", "turnover_proxy"], ascending=[False, True]).reset_index(drop=True)

    return df

lambda_u_grid = [1e-9, 1e-6, 1e-3]
lambda_h_grid = [1e-9, 1e-6, 1e-3]

df_sweep = sweep_lq_hyperparams(
    cm=cm,
    X_path=X_smooth,          # or X_stress
    Phi_hat=Phi_hat,
    X_ref=X_ref,
    A_notional=A_NOTIONAL,
    L_notional=L_NOTIONAL,
    tau_A=3.0,
    tau_L=1.0,
    dt=1/12,
    alpha_nii=1.0,
    lambda_u_grid=lambda_u_grid,
    lambda_h_grid=lambda_h_grid
)

df_sweep

def plot_lq_tradeoff(df, title="LQ hyperparameter sweep: NII risk vs turnover"):
    """
    Scatter plot:
      x = turnover (mean_abs_u)
      y = hedged NII volatility (std_hedged)
    """
    x = df["mean_abs_u"].values
    y = df["std_hedged"].values

    plt.figure(figsize=(8,5))
    plt.scatter(x, y)

    # annotate points with (lambda_u, lambda_h)
    for _, row in df.iterrows():
        plt.annotate(
            f"u={row['lambda_u']:.0e}, h={row['lambda_h']:.0e}",
            (row["mean_abs_u"], row["std_hedged"]),
            fontsize=8,
            xytext=(4, 4),
            textcoords="offset points"
        )

    plt.xlabel("Turnover proxy: mean(|u_t|)")
    plt.ylabel("Risk proxy: std(NII_hedged)")
    plt.title(title)
    plt.tight_layout()
    plt.show()

plot_lq_tradeoff(df_sweep)

def summarize_paths(N0, N1, h, u):
    out = {}
    out["std_unhedged"] = float(np.std(N0))
    out["std_hedged"] = float(np.std(N1))
    out["std_reduction_%"] = float(100.0 * (1.0 - out["std_hedged"]/out["std_unhedged"])) if out["std_unhedged"] > 0 else np.nan

    out["p05_unhedged"] = float(np.quantile(N0, 0.05))
    out["p05_hedged"] = float(np.quantile(N1, 0.05))
    out["min_unhedged"] = float(np.min(N0))
    out["min_hedged"] = float(np.min(N1))

    out["mean_unhedged"] = float(np.mean(N0))
    out["mean_hedged"] = float(np.mean(N1))

    out["mean_abs_h"] = float(np.mean(np.abs(h)))
    out["max_abs_h"]  = float(np.max(np.abs(h)))
    out["mean_abs_u"] = float(np.mean(np.abs(u)))
    out["max_abs_u"]  = float(np.max(np.abs(u)))
    return out


def plot_lq_tradeoff(df, title="LQ sweep: NII risk vs turnover"):
    """
    Scatter: x=turnover (mean|u|), y=risk (std hedged).
    Annotate with lambdas.
    """
    x = df["mean_abs_u"].values
    y = df["std_hedged"].values

    plt.figure(figsize=(8,5))
    plt.scatter(x, y)
    for _, row in df.iterrows():
        plt.annotate(
            f"u={row['lambda_u']:.0e}, h={row['lambda_h']:.0e}",
            (row["mean_abs_u"], row["std_hedged"]),
            fontsize=8, xytext=(4,4), textcoords="offset points"
        )
    plt.xlabel("Turnover proxy: mean(|u_t|)")
    plt.ylabel("Risk proxy: std(NII_hedged)")
    plt.title(title)
    plt.tight_layout()
    plt.show()

# Monthly setup
dt = 1.0/12.0
tau_A = 3.0
tau_L = 1.0
A_notional = A_NOTIONAL
L_notional = L_NOTIONAL

# Fit VAR(1) to smoothed factors
c_hat, Phi_hat, Sigma_hat = fit_var1(X_smooth)
X_ref = X_smooth.mean(axis=0)

NII_TARGET = (A_notional * afns_yield_single_from_cm(cm, X_ref, tau_A) * dt
              - L_notional * afns_yield_single_from_cm(cm, X_ref, tau_L) * dt)

_gA = numerical_grad_y(cm, X_ref, tau_A)
_gL = numerical_grad_y(cm, X_ref, tau_L)
K_STRIKE = afns_yield_single_from_cm(cm, X_ref, tau_L)
H_VM = L_notional - A_notional * float((_gL @ Sigma_hat @ _gA) / (_gL @ Sigma_hat @ _gL))

print("Phi_hat:\n", Phi_hat)
print("X_ref:", X_ref)

# Stress scenario builders (deterministic paths for clean comparison)
def scenario_parallel_up(T, shock_bps=200):
    shock = shock_bps/10000.0
    shock_vec = np.array([shock, 0.0, 0.0])  # Level shock
    return make_stress_path_from_var1(X0=X_ref, c=c_hat, Phi=Phi_hat, T=T, shock_vec=shock_vec)

def scenario_bear_steepener(T, level_bps=200, slope_bps=100):
    # Simple stylized: +level and -slope (more steepness in NS sign convention)
    level = level_bps/10000.0
    slope = slope_bps/10000.0
    shock_vec = np.array([level, -slope, 0.0])
    return make_stress_path_from_var1(X0=X_ref, c=c_hat, Phi=Phi_hat, T=T, shock_vec=shock_vec)

def scenario_high_vol(T, vol_scale=3.0, seed=123):
    # Stochastic path: amplify innovations
    rng = np.random.default_rng(seed)
    X = np.zeros((T,3))
    X[0] = X_ref.copy()
    for t in range(T-1):
        eps = rng.multivariate_normal(np.zeros(3), vol_scale * Sigma_hat)
        X[t+1] = c_hat + Phi_hat @ X[t] + eps
    return X

T = X_smooth.shape[0]
X_parallel = scenario_parallel_up(T, shock_bps=200)
X_steepen  = scenario_bear_steepener(T, level_bps=200, slope_bps=100)
X_highvol  = scenario_high_vol(T, vol_scale=3.0, seed=123)

Phi_hat:
 [[ 0.9938804   0.00462503 -0.00218918]
 [-0.01640262  0.99214422  0.01098002]
 [ 0.01286072 -0.0127359   0.95747698]]
X_ref: [ 0.06421823 -0.03616098 -0.02123435]

alpha_nii = 1.0

lambda_u_grid = [1e-4, 1e-3, 1e-2]
lambda_h_grid = [1e-7, 1e-6, 1e-5]

rows = []
for lambda_u in lambda_u_grid:
    for lambda_h in lambda_h_grid:

        H_x, Q_s, R = build_Hx_QR_from_nii(
            cm, X_ref,
            A_notional, L_notional, tau_A, tau_L, dt,
            alpha_nii=alpha_nii,
            lambda_h=lambda_h,
            lambda_u=lambda_u
        )

        res = run_lq_hedge_on_path(
            cm, X_smooth, Phi_hat, H_x, Q_s, R,
            A_notional, L_notional, tau_A, tau_L, dt
        )

        metrics = summarize_paths(res["NII_unhedged"], res["NII_hedged"], res["h"], res["u"])
        rows.append({
            "lambda_u": float(lambda_u),
            "lambda_h": float(lambda_h),
            "Hx_hedge_sensitivity": float(H_x[3]),
            "K_max_abs": float(np.max(np.abs(res["K"]))),
            **metrics
        })

df_lq = pd.DataFrame(rows)

# NOTE: Hx_hedge_sensitivity (= H_x[3]) is ~0 by design --
# the swaplet is struck at par, so the hedge adds value through variance reduction
# (the q_var term on Q_s[3,3]), not through a nonzero mean loading.

# sort: prefer bigger std reduction and lower turnover
df_lq = df_lq.sort_values(["std_reduction_%", "mean_abs_u"], ascending=[False, True]).reset_index(drop=True)

df_lq

plot_lq_tradeoff(df_lq, title="Baseline path: LQ sweep (risk vs turnover)")

benchmark = df_lq.iloc[3].to_dict()
benchmark

{'lambda_u': 0.0001,
 'lambda_h': 1e-07,
 'Hx_hedge_sensitivity': 0.0,
 'K_max_abs': 0.036614131501450695,
 'std_unhedged': 0.0903518153735064,
 'std_hedged': 0.03419364385071961,
 'std_reduction_%': 62.15500074971807,
 'p05_unhedged': -0.2149733204098512,
 'p05_hedged': -0.13128531071927427,
 'min_unhedged': -0.2813655407987071,
 'min_hedged': -0.26736302199945294,
 'mean_unhedged': -0.0807641374102874,
 'mean_hedged': -0.08742536725652207,
 'mean_abs_h': 38.49985375524823,
 'max_abs_h': 41.09174810022436,
 'mean_abs_u': 0.09511978726903773,
 'max_abs_u': 1.5045388196160925}

def eval_policy_on_path(X_path, lambda_u, lambda_h, label):
    H_x, Q_s, R = build_Hx_QR_from_nii(
        cm, X_ref,
        A_notional, L_notional, tau_A, tau_L, dt,
        alpha_nii=alpha_nii,
        lambda_h=lambda_h,
        lambda_u=lambda_u
    )

    res = run_lq_hedge_on_path(
        cm, X_path, Phi_hat, H_x, Q_s, R,
        A_notional, L_notional, tau_A, tau_L, dt
    )

    metrics = summarize_paths(res["NII_unhedged"], res["NII_hedged"], res["h"], res["u"])
    return {
        "scenario": label,
        "lambda_u": float(lambda_u),
        "lambda_h": float(lambda_h),
        **metrics
    }, res


# Choose lambdas (use preferred benchmark)
lambda_u_star = float(benchmark["lambda_u"])
lambda_h_star = float(benchmark["lambda_h"])

rows = []
res_store = {}

for label, X_path in [
    ("baseline (smoothed)", X_smooth),
    ("stress: +200bp parallel", X_parallel),
    ("stress: bear steepener", X_steepen),
    ("stress: high vol x3", X_highvol),
]:
    row, res = eval_policy_on_path(X_path, lambda_u_star, lambda_h_star, label)
    rows.append(row)
    res_store[label] = res

df_stress = pd.DataFrame(rows)
df_stress

def plot_paths(res, title):
    plt.figure(figsize=(9,4))
    plt.plot(res["NII_unhedged"], label="Unhedged")
    plt.plot(res["NII_hedged"], label="LQ hedged")
    plt.title(title)
    plt.xlabel("t (months)")
    plt.ylabel("NII")
    plt.legend()
    plt.tight_layout()
    plt.show()

    plt.figure(figsize=(9,3))
    plt.plot(res["h"], label="h_t")
    plt.title(title + " — hedge notional")
    plt.xlabel("t (months)")
    plt.ylabel("h_t")
    plt.legend()
    plt.tight_layout()
    plt.show()

plot_paths(res_store["stress: +200bp parallel"], "LQ benchmark under +200bp parallel shock")
plot_paths(res_store["stress: bear steepener"], "LQ benchmark under bear steepener")
plot_paths(res_store["stress: high vol x3"], "LQ benchmark under high-vol regime")

import gymnasium as gym
from gymnasium import spaces

class IrrbbNiiHedgeEnv(gym.Env):
    """
    RL environment for IRRBB-style NII hedging using a FRA hedge (L2 / quadratic costs).

    State : [L, S, C, h] (factors centered around X_ref if center_state=True)
    Action: u_raw in [-1, 1]; the applied hedge change is  Delta h = u_raw * u_scale.
            (Issue 2 fix: identical action -> Delta h mapping in env and eval rollout.)

    Dynamics:
      X_{t+1} = c + Phi X_t + eps_t,  eps_t ~ N(0, Sigma)
      h_{t+1} = h_t + Delta h

    Reward (to maximize) -- Issue 3 fix: penalize NII *variance* around its target,
    not its level:
      r_t = -[ (NII_{t+1} - nii_target)^2 + lambda_h h_t^2 + lambda_u (Delta h)^2 ]
    """
    metadata = {"render_modes": []}

    def __init__(self,
                 cm,
                 c_hat, Phi_hat, Sigma_hat, X_ref,
                 A_notional=A_NOTIONAL, L_notional=L_NOTIONAL,
                 tau_A=3.0, tau_L=1.0, dt=1/12,
                 lambda_u=1e-3, lambda_h=1e-6,
                 u_scale=U_SCALE,
                 nii_target=None,
                 episode_len=240,
                 center_state=True,
                 seed=123):
        super().__init__()

        self.cm = cm
        self.c = np.asarray(c_hat, float)
        self.Phi = np.asarray(Phi_hat, float)
        self.Sigma = np.asarray(Sigma_hat, float)
        self.X_ref = np.asarray(X_ref, float)

        self.A_notional = float(A_notional)
        self.L_notional = float(L_notional)
        self.tau_A = float(tau_A)
        self.tau_L = float(tau_L)
        self.dt = float(dt)

        self.lambda_u = float(lambda_u)
        self.lambda_h = float(lambda_h)
        self.u_scale = float(u_scale)
        self.nii_target = float(NII_TARGET if nii_target is None else nii_target)

        self.episode_len = int(episode_len)
        self.center_state = bool(center_state)
        self.rng = np.random.default_rng(seed)

        # Observation: 4D continuous
        obs_high = np.full((4,), np.inf, dtype=np.float32)
        self.observation_space = spaces.Box(low=-obs_high, high=obs_high, dtype=np.float32)

        # Action is the *raw* control in [-1, 1]; scaled to Delta h by u_scale.
        self.action_space = spaces.Box(low=-1.0, high=1.0, shape=(1,), dtype=np.float32)

        self.t = 0
        self.X = None
        self.h = None

    # --- AFNS helpers ---
    def _y(self, X, tau):
        taus = np.array([tau], dtype=float)
        return float(afns_yields_from_factors(X, taus, self.cm.lam, self.cm.sig1, self.cm.sig2, self.cm.sig3)[0])

    def _fwd_cc(self, X, tau1, tau2):
        y1 = self._y(X, tau1)
        y2 = self._y(X, tau2)
        return (tau2 * y2 - tau1 * y1) / (tau2 - tau1)

    def _nii_one_step(self, X_t, X_next, h_t):
        yA = self._y(X_t, self.tau_A)
        yL = self._y(X_t, self.tau_L)
        # Issue 10: pay-fixed swaplet on the current reference rate y_t(tau_L)
        return (self.A_notional * yA * self.dt
                - self.L_notional * yL * self.dt
                + h_t * (yL - K_STRIKE) * self.dt)

    def _obs(self):
        x = np.array([self.X[0], self.X[1], self.X[2], self.h], dtype=np.float32)
        if self.center_state:
            x[:3] = x[:3] - self.X_ref.astype(np.float32)
        return x

    def reset(self, *, seed=None, options=None):
        if seed is not None:
            self.rng = np.random.default_rng(seed)
        self.t = 0
        X0 = self.X_ref + self.rng.multivariate_normal(np.zeros(3), 0.1 * self.Sigma)
        self.X = np.asarray(X0, float)
        self.h = 0.0
        return self._obs(), {}

    def step(self, action):
        u_raw = float(np.clip(np.asarray(action, dtype=float)[0], -1.0, 1.0))
        u = u_raw * self.u_scale          # Delta h (Issue 2: same mapping as eval rollout)

        # Next factors
        eps = self.rng.multivariate_normal(np.zeros(3), self.Sigma)
        X_next = self.c + self.Phi @ self.X + eps

        # NII uses h_t (pre-trade)
        nii = self._nii_one_step(self.X, X_next, self.h)

        # Update hedge
        h_next = self.h + u

        # Issue 3: variance objective around nii_target (not NII^2)
        reward = -((nii - self.nii_target)**2
                   + self.lambda_h * (self.h**2)
                   + self.lambda_u * (u**2))

        self.X = X_next
        self.h = h_next
        self.t += 1

        terminated = False
        truncated = (self.t >= self.episode_len)
        info = {"nii": nii, "h": self.h, "u": u}
        return self._obs(), float(reward), terminated, truncated, info

from stable_baselines3 import SAC
from stable_baselines3.common.env_util import make_vec_env

np.random.seed(123)

# Fit VAR(1) once
c_hat, Phi_hat, Sigma_hat = fit_var1(X_smooth)
X_ref = X_smooth.mean(axis=0)

# Use the same lambdas previously selected for LQ benchmark
lambda_u_rl = float(lambda_u_star)
lambda_h_rl = float(lambda_h_star)

def make_env():
    return IrrbbNiiHedgeEnv(
        cm=cm,
        c_hat=c_hat, Phi_hat=Phi_hat, Sigma_hat=Sigma_hat, X_ref=X_ref,
        A_notional=100.0, L_notional=100.0,
        tau_A=3.0, tau_L=1.0, dt=1/12,
        lambda_u=lambda_u_rl, lambda_h=lambda_h_rl,
        u_scale=U_SCALE,
        episode_len=240,          # 20 years monthly
        center_state=True,
        seed=123
    )

vec_env = make_vec_env(make_env, n_envs=8)  # parallel rollouts

model = SAC(
    "MlpPolicy",
    vec_env,
    verbose=0,
    learning_rate=3e-4,
    batch_size=256,
    buffer_size=200_000,
    train_freq=1,
    gradient_steps=1,
    gamma=0.99,
    tau=0.005,
)

model.learn(total_timesteps=300_000)

# Save if you want
model.save("sac_irrbb_nii_hedge")

def nii_path_from_X_and_h(cm, X_path, h_path, A_notional=A_NOTIONAL, L_notional=L_NOTIONAL,
                         tau_A=3.0, tau_L=1.0, dt=1/12):
    T = X_path.shape[0]
    NII = np.zeros(T-1)

    for t in range(T-1):
        X_t, X_next = X_path[t], X_path[t+1]

        yA = float(afns_yields_from_factors(X_t, np.array([tau_A]), cm.lam, cm.sig1, cm.sig2, cm.sig3)[0])
        yL = float(afns_yields_from_factors(X_t, np.array([tau_L]), cm.lam, cm.sig1, cm.sig2, cm.sig3)[0])

        # Issue 10: pay-fixed swaplet on the current reference rate y_t(tau_L).
        NII[t] = (A_notional*yA*dt - L_notional*yL*dt + h_path[t]*(yL - K_STRIKE)*dt)

    return NII

def rollout_policy_on_exogenous_X(cm, X_path, policy,
                                  K_lq=None, rl_model=None,
                                  X_ref=None, center_state=True,
                                  u_scale=U_SCALE,
                                  A_notional=A_NOTIONAL, L_notional=L_NOTIONAL,
                                  tau_A=3.0, tau_L=1.0, dt=1/12):
    """
    Evaluate a policy on an exogenous factor path.
    policy: "unhedged" | "lq" | "rl"

    Issue 2 fix: the action -> Delta h mapping matches the training env exactly.
      - RL: the agent outputs u_raw in [-1, 1]; applied Delta h = u_raw * u_scale.
      - LQ: u = -K x_tilde with x_tilde = [X - X_ref, h] (centered, Issue 3),
            Delta h clipped to +-u_scale.
    Returns dict with h, u, NII.
    """
    T = X_path.shape[0]
    h = np.zeros(T)
    u = np.zeros(T-1)
    xr = np.zeros(3) if X_ref is None else np.asarray(X_ref, float)

    for t in range(T-1):
        if policy == "unhedged":
            u_t = 0.0
        elif policy == "lq":
            if K_lq is None:
                raise ValueError("Need K_lq for LQ.")
            # Issue 3: centered state so LQ targets the mean curve, not zero.
            x_t = np.array([X_path[t,0] - xr[0],
                            X_path[t,1] - xr[1],
                            X_path[t,2] - xr[2],
                            h[t] - H_VM], dtype=float)   # Issue 10: target H_VM
            u_t = -float((K_lq @ x_t).item())
            u_t = float(np.clip(u_t, -u_scale, u_scale))
        elif policy == "rl":
            if rl_model is None:
                raise ValueError("Need rl_model for RL.")
            obs = np.array([X_path[t,0], X_path[t,1], X_path[t,2], h[t]], dtype=np.float32)
            if center_state and (X_ref is not None):
                obs[:3] -= xr.astype(np.float32)
            act, _ = rl_model.predict(obs, deterministic=True)
            u_raw = float(np.clip(float(np.asarray(act).ravel()[0]), -1.0, 1.0))
            u_t = u_raw * u_scale          # Issue 2: apply u_scale exactly like the env
        else:
            raise ValueError("Unknown policy.")

        u[t] = u_t
        h[t+1] = h[t] + u_t

    NII = nii_path_from_X_and_h(cm, X_path, h, A_notional, L_notional, tau_A, tau_L, dt)
    return {"h": h, "u": u, "NII": NII}

def summarize(N):
    return {
        "std": float(np.std(N)),
        "p05": float(np.quantile(N, 0.05)),
        "min": float(np.min(N)),
        "mean": float(np.mean(N)),
    }

def compare_on_scenario(name, X_path, K_lq, rl_model, u_scale=U_SCALE):
    res0 = rollout_policy_on_exogenous_X(cm, X_path, "unhedged", u_scale=u_scale, X_ref=X_ref)
    resL = rollout_policy_on_exogenous_X(cm, X_path, "lq", K_lq=K_lq, u_scale=u_scale, X_ref=X_ref)
    resR = rollout_policy_on_exogenous_X(cm, X_path, "rl", rl_model=rl_model, u_scale=u_scale, X_ref=X_ref)

    row = {
        "scenario": name,
        "unhedged_std": summarize(res0["NII"])["std"],
        "lq_std": summarize(resL["NII"])["std"],
        "rl_std": summarize(resR["NII"])["std"],
        "unhedged_p05": summarize(res0["NII"])["p05"],
        "lq_p05": summarize(resL["NII"])["p05"],
        "rl_p05": summarize(resR["NII"])["p05"],
        "lq_turnover": float(np.mean(np.abs(resL["u"]))),
        "rl_turnover": float(np.mean(np.abs(resR["u"]))),
        "lq_inv": float(np.mean(np.abs(resL["h"]))),
        "rl_inv": float(np.mean(np.abs(resR["h"]))),
    }

    plt.figure(figsize=(9,4))
    plt.plot(res0["NII"], label="Unhedged")
    plt.plot(resL["NII"], label="LQ")
    plt.plot(resR["NII"], label="RL (SAC)")
    plt.title(name)
    plt.xlabel("t (months)")
    plt.ylabel("NII")
    plt.legend()
    plt.tight_layout()
    plt.show()
    return row

# Build K_lq once using chosen lambdas and the same build_Hx_QR_from_nii + Riccati pipeline
H_x, Q_s, R = build_Hx_QR_from_nii(
    cm, X_ref, A_notional, L_notional, tau_A, tau_L, dt,
    alpha_nii=1.0, lambda_h=lambda_h_star, lambda_u=lambda_u_star
)
A_lq, B_lq = build_AB_from_Phi(Phi_hat)
P_lq, K_lq = solve_discrete_riccati(A_lq, B_lq, Q_s, R)

rows = []
rows.append(compare_on_scenario("Stress: +200bp parallel", X_parallel, K_lq, model))
rows.append(compare_on_scenario("Stress: bear steepener", X_steepen, K_lq, model))
rows.append(compare_on_scenario("Stress: high vol x3", X_highvol, K_lq, model))

import pandas as pd
df_compare = pd.DataFrame(rows)
df_compare

compare_on_scenario("Stress: high vol x3", X_highvol, K_lq, model)

{'scenario': 'Stress: high vol x3',
 'unhedged_std': 0.05929270644490275,
 'lq_std': 0.03276737389012783,
 'rl_std': 0.05327312544774718,
 'unhedged_p05': -0.1889067930597452,
 'lq_p05': -0.1696746955554819,
 'rl_p05': -0.22120991439574125,
 'lq_turnover': 0.09511978726903773,
 'rl_turnover': 1.8364524951687566,
 'lq_inv': 38.49985375524823,
 'rl_inv': 13.353505131829419}

class IrrbbNiiHedgeEnvL1(gym.Env):
    """
    IRRBB NII hedging environment with L1 transaction costs.

    State: [L, S, C, h]  (AFNS factors + hedge inventory)
    Action: u in [-1,1], scaled to Δh via u_scale

    Reward:
        r_t = - [ NII_{t+1}^2 + lambda_h * h_t^2 + kappa_u * |u_t| ]

    This breaks LQ assumptions and favors sparse trading.
    """

    metadata = {"render_modes": []}

    def __init__(
        self,
        cm,
        c_hat,
        Phi_hat,
        Sigma_hat,
        X_ref,
        A_notional=A_NOTIONAL,
        L_notional=L_NOTIONAL,
        tau_A=3.0,
        tau_L=1.0,
        dt=1 / 12,
        lambda_h=1e-6,
        kappa_u=3e-4,
        nii_target=None,
        u_scale=U_SCALE,
        no_trade_eps=0.0,
        episode_len=240,
        center_state=True,
        seed=123,
    ):
        super().__init__()

        self.cm = cm
        self.c = np.asarray(c_hat, float)
        self.Phi = np.asarray(Phi_hat, float)
        self.Sigma = np.asarray(Sigma_hat, float)
        self.X_ref = np.asarray(X_ref, float)

        self.A_notional = A_notional
        self.L_notional = L_notional
        self.tau_A = tau_A
        self.tau_L = tau_L
        self.dt = dt

        self.lambda_h = lambda_h
        self.kappa_u = kappa_u
        self.nii_target = float(NII_TARGET if nii_target is None else nii_target)
        self.u_scale = u_scale
        self.no_trade_eps = no_trade_eps

        self.episode_len = episode_len
        self.center_state = center_state
        self.rng = np.random.default_rng(seed)

        # State: (L, S, C, h)
        self.observation_space = spaces.Box(
            low=-np.inf, high=np.inf, shape=(4,), dtype=np.float32
        )

        # Action: scalar in [-1, 1]
        self.action_space = spaces.Box(
            low=-1.0, high=1.0, shape=(1,), dtype=np.float32
        )

        self.t = 0
        self.X = None
        self.h = None

    # ---------- AFNS helpers ----------

    def _y(self, X, tau):
        return float(
            afns_yields_from_factors(
                X,
                np.array([tau]),
                self.cm.lam,
                self.cm.sig1,
                self.cm.sig2,
                self.cm.sig3,
            )[0]
        )

    def _fwd_cc(self, X, tau1, tau2):
        y1 = self._y(X, tau1)
        y2 = self._y(X, tau2)
        return (tau2 * y2 - tau1 * y1) / (tau2 - tau1)

    def _nii_step(self, X_t, X_next, h_t):
        yA = self._y(X_t, self.tau_A)
        yL = self._y(X_t, self.tau_L)

        # Issue 10: pay-fixed swaplet on the current reference rate y_t(tau_L)
        return (
            self.A_notional * yA * self.dt
            - self.L_notional * yL * self.dt
            + h_t * (yL - K_STRIKE) * self.dt
        )

    # ---------- Gym API ----------

    def _obs(self):
        obs = np.array([self.X[0], self.X[1], self.X[2], self.h], dtype=np.float32)
        if self.center_state:
            obs[:3] -= self.X_ref.astype(np.float32)
        return obs

    def reset(self, *, seed=None, options=None):
        if seed is not None:
            self.rng = np.random.default_rng(seed)

        self.t = 0
        self.X = self.X_ref + self.rng.multivariate_normal(
            np.zeros(3), 0.1 * self.Sigma
        )
        self.h = 0.0
        return self._obs(), {}

    def step(self, action):
        u_raw = float(np.clip(action[0], -1.0, 1.0))
        u = u_raw * self.u_scale

        if abs(u) < self.no_trade_eps:
            u = 0.0

        # Next factors
        eps = self.rng.multivariate_normal(np.zeros(3), self.Sigma)
        X_next = self.c + self.Phi @ self.X + eps

        nii = self._nii_step(self.X, X_next, self.h)

        # Update hedge
        h_next = self.h + u

        # L1 reward
        # Issue 3: variance objective around nii_target (not NII^2)
        reward = -(
            (nii - self.nii_target)**2 + self.lambda_h * (self.h**2) + self.kappa_u * abs(u)
        )

        self.X = X_next
        self.h = h_next
        self.t += 1

        terminated = False
        truncated = self.t >= self.episode_len

        info = {"nii": nii, "h": self.h, "u": u}

        return self._obs(), float(reward), terminated, truncated, info

def train_sac_l1(
    cm,
    X_smooth,
    c_hat,
    Phi_hat,
    Sigma_hat,
    X_ref,
    lambda_h,
    kappa_u,
    total_timesteps=800_000,
    n_envs=8,
    u_scale=25.0,
    no_trade_eps=0.0,
    seed=123,
):
    def make_env():
        return IrrbbNiiHedgeEnvL1(
            cm=cm,
            c_hat=c_hat,
            Phi_hat=Phi_hat,
            Sigma_hat=Sigma_hat,
            X_ref=X_ref,
            lambda_h=lambda_h,
            kappa_u=kappa_u,
            u_scale=u_scale,
            no_trade_eps=no_trade_eps,
            seed=seed,
        )

    vec_env = make_vec_env(make_env, n_envs=n_envs)

    model = SAC(
        "MlpPolicy",
        vec_env,
        learning_rate=3e-4,
        batch_size=256,
        buffer_size=300_000,
        gamma=0.99,
        tau=0.005,
        train_freq=1,
        gradient_steps=1,
        verbose=0,
    )

    model.learn(total_timesteps=total_timesteps)
    return model

def eval_L1_cost(NII, h, u, lambda_h, kappa_u):
    """
    Average L1-based economic cost.
    """
    return float(
        np.mean(NII**2 + lambda_h * (h[:-1] ** 2) + kappa_u * np.abs(u))
    )

lambda_h_rl = lambda_h_star   # keep inventory discipline
kappa_u_rl  = 3e-4            # L1 trading cost
u_scale     = 25.0            # hedge impact scale

model_l1 = train_sac_l1(
    cm=cm,
    X_smooth=X_smooth,
    c_hat=c_hat,
    Phi_hat=Phi_hat,
    Sigma_hat=Sigma_hat,
    X_ref=X_ref,
    lambda_h=lambda_h_rl,
    kappa_u=kappa_u_rl,
    total_timesteps=800_000,   # increase to 1–2M if needed
    n_envs=8,
    u_scale=u_scale,
    no_trade_eps=0.0           # later try 0.5 or 1.0
)

res_unhedged = rollout_policy_on_exogenous_X(
    cm, X_parallel, policy="unhedged", u_scale=u_scale, X_ref=X_ref
)

res_lq = rollout_policy_on_exogenous_X(
    cm, X_parallel, policy="lq",
    K_lq=K_lq, u_scale=u_scale, X_ref=X_ref
)

res_rl = rollout_policy_on_exogenous_X(
    cm, X_parallel, policy="rl",
    rl_model=model_l1, u_scale=u_scale, X_ref=X_ref
)

eval_unhedged = eval_L1_cost(
    res_unhedged["NII"], res_unhedged["h"], res_unhedged["u"],
    lambda_h=lambda_h_star, kappa_u=kappa_u_rl
)

eval_lq = eval_L1_cost(
    res_lq["NII"], res_lq["h"], res_lq["u"],
    lambda_h=lambda_h_star, kappa_u=kappa_u_rl
)

eval_rl = eval_L1_cost(
    res_rl["NII"], res_rl["h"], res_rl["u"],
    lambda_h=lambda_h_star, kappa_u=kappa_u_rl
)

print("L1 cost (lower is better):")
print("Unhedged:", eval_unhedged)
print("LQ:", eval_lq)
print("RL:", eval_rl)

L1 cost (lower is better):
Unhedged: 0.004343321730425727
LQ: 0.006914232403692727
RL: 0.0062331473370535705

plt.figure(figsize=(9,4))
plt.plot(res_unhedged["NII"], label="Unhedged")
plt.plot(res_lq["NII"], label="LQ (L2)")
plt.plot(res_rl["NII"], label="RL (L1)")
plt.title("Stress: +200bp parallel")
plt.xlabel("t (months)")
plt.ylabel("NII")
plt.legend()
plt.tight_layout()
plt.show()

plt.figure(figsize=(9,3))
plt.plot(res_lq["h"], label="LQ hedge")
plt.plot(res_rl["h"], label="RL hedge")
plt.title("Hedge inventory")
plt.legend()
plt.tight_layout()
plt.show()

# ============================================================
# 0) YOU PROVIDE THESE
# ============================================================
# You need a function that returns a rollout dict:
# {"NII": (T-1,), "h": (T,), "u": (T-1,)}
#

def get_rollout(cm, X_path, policy_name,
                K_lq=None, rl_model=None,
                X_ref=None, u_scale=U_SCALE):
    """
    Adapter around existing rollout_policy_on_exogenous_X.
    policy_name in {"unhedged","lq","rl"}.
    Returns dict with keys: NII, h, u.
    """
    if policy_name == "unhedged":
        res = rollout_policy_on_exogenous_X(cm, X_path, "unhedged",
                                            u_scale=u_scale, X_ref=X_ref)
    elif policy_name == "lq":
        res = rollout_policy_on_exogenous_X(cm, X_path, "lq",
                                            K_lq=K_lq, u_scale=u_scale, X_ref=X_ref)
    elif policy_name == "rl":
        res = rollout_policy_on_exogenous_X(cm, X_path, "rl",
                                            rl_model=rl_model, u_scale=u_scale, X_ref=X_ref)
    else:
        raise ValueError("policy_name must be 'unhedged', 'lq', or 'rl'")
    return {"NII": np.asarray(res["NII"]), "h": np.asarray(res["h"]), "u": np.asarray(res["u"])}


# ============================================================
# 1) METRICS + TABLES
# ============================================================

def l1_economic_cost(nii, h, u, lambda_h, kappa_u):
    """
    Evaluation functional for the L1 setting:
      mean( NII^2 + lambda_h*h^2 + kappa_u*|u| )
    (h aligned to t where NII,u exist => h[:-1])
    """
    nii = np.asarray(nii)
    h = np.asarray(h)
    u = np.asarray(u)
    return float(np.mean(nii**2 + lambda_h*(h[:-1]**2) + kappa_u*np.abs(u)))

def summarize_rollout(roll, lambda_h=None, kappa_u=None):
    nii, h, u = roll["NII"], roll["h"], roll["u"]
    out = {
        "std_NII": float(np.std(nii)),
        "p05_NII": float(np.quantile(nii, 0.05)),
        "min_NII": float(np.min(nii)),
        "mean_NII": float(np.mean(nii)),
        "mean_abs_u": float(np.mean(np.abs(u))),
        "max_abs_u": float(np.max(np.abs(u))),
        "mean_abs_h": float(np.mean(np.abs(h))),
        "max_abs_h": float(np.max(np.abs(h))),
    }
    if (lambda_h is not None) and (kappa_u is not None):
        out["L1_cost"] = l1_economic_cost(nii, h, u, lambda_h, kappa_u)
    return out

def build_metrics_tables(cm, scenarios, K_lq, rl_model,
                         X_ref, u_scale,
                         lambda_h_eval, kappa_u_eval):
    """
    scenarios: dict name -> X_path
    Returns:
      df_metrics: multiindex (scenario, policy)
      df_norm: normalized to unhedged per scenario
    """
    rows = []
    for scen_name, X_path in scenarios.items():
        for pol, label in [("unhedged","Unhedged"), ("lq","LQ (L2)"), ("rl","RL (L1)")]:
            roll = get_rollout(cm, X_path, pol, K_lq=K_lq, rl_model=rl_model,
                               X_ref=X_ref, u_scale=u_scale)
            met = summarize_rollout(roll, lambda_h=lambda_h_eval, kappa_u=kappa_u_eval)
            rows.append({"scenario": scen_name, "policy": label, **met})

    df = pd.DataFrame(rows).set_index(["scenario","policy"]).sort_index()

    # Normalized table: divide each metric by unhedged metric (per scenario)
    metrics_cols = [c for c in df.columns if c not in []]
    df_norm = df.copy()
    for scen in df.index.get_level_values(0).unique():
        base = df.loc[(scen, "Unhedged"), metrics_cols]
        df_norm.loc[scen, metrics_cols] = (df.loc[scen, metrics_cols].values / base.values) * 100.0

    return df, df_norm


# ============================================================
# 2) PLOTS
# ============================================================

def plot_nii_overlay(rolls, title):
    plt.figure(figsize=(9,4))
    for label, roll in rolls.items():
        plt.plot(roll["NII"], label=label)
    plt.title(title)
    plt.xlabel("t (months)")
    plt.ylabel("NII")
    plt.legend()
    plt.tight_layout()
    plt.show()

def plot_h_overlay(rolls, title):
    plt.figure(figsize=(9,3))
    for label, roll in rolls.items():
        plt.plot(roll["h"], label=label)
    plt.title(title)
    plt.xlabel("t (months)")
    plt.ylabel("Hedge inventory h_t")
    plt.legend()
    plt.tight_layout()
    plt.show()

def plot_absu_histogram(roll_lq, roll_rl, title):
    plt.figure(figsize=(8,4))
    plt.hist(np.abs(roll_lq["u"]), bins=40, alpha=0.6, label="LQ (L2)")
    plt.hist(np.abs(roll_rl["u"]), bins=40, alpha=0.6, label="RL (L1)")
    plt.title(title)
    plt.xlabel("|u_t| (absolute hedge change)")
    plt.ylabel("Frequency")
    plt.legend()
    plt.tight_layout()
    plt.show()

def plot_risk_turnover_scatter(df_metrics, title="Risk–turnover scatter"):
    """
    Scatter per scenario/policy:
      x = mean_abs_u, y = std_NII
    """
    plt.figure(figsize=(8,5))
    for (scen, pol), row in df_metrics.iterrows():
        x = row["mean_abs_u"]
        y = row["std_NII"]
        plt.scatter(x, y)
        plt.annotate(f"{scen}\n{pol}", (x, y), fontsize=8, xytext=(4,4), textcoords="offset points")

    plt.xlabel("Turnover proxy: mean(|u_t|)")
    plt.ylabel("Risk proxy: std(NII)")
    plt.title(title)
    plt.tight_layout()
    plt.show()


# ============================================================
# 3) RUN IT: define scenario dictionary + produce outputs
# ============================================================

# --- Provide scenario paths ---
scenarios = {
    "Baseline (smoothed)": X_smooth,
    "Stress: +200bp parallel": X_parallel,
    "Stress: high vol x3": X_highvol,
}

# --- Evaluation settings ---
ACTION_MAX = U_SCALE
LAMBDA_H_EVAL = lambda_h_star      # chosen hedge-inventory penalty
KAPPA_U_EVAL  = 3e-4               # same kappa_u from L1 RL training/eval

# --- Models from before ---
# K_lq: from Riccati solution for LQ benchmark
# model_l1: trained SAC model for RL (L1)

df_metrics, df_norm = build_metrics_tables(
    cm=cm,
    scenarios=scenarios,
    K_lq=K_lq,
    rl_model=model_l1,
    X_ref=X_ref,
    u_scale=ACTION_MAX,
    lambda_h_eval=LAMBDA_H_EVAL,
    kappa_u_eval=KAPPA_U_EVAL
)

print("=== TABLE 1: Core metrics ===")
display(df_metrics)

print("=== TABLE 2: Normalized to Unhedged (Unhedged=100) ===")
display(df_norm)

# ============================================================
# 4) SELECTED PLOTS (6 total)
# ============================================================

# Helper to fetch rolls for a scenario
def get_rolls_for_scenario(X_path, scen_label):
    roll_u = get_rollout(cm, X_path, "unhedged", X_ref=X_ref, u_scale=ACTION_MAX)
    roll_l = get_rollout(cm, X_path, "lq", K_lq=K_lq, X_ref=X_ref, u_scale=ACTION_MAX)
    roll_r = get_rollout(cm, X_path, "rl", rl_model=model_l1, X_ref=X_ref, u_scale=ACTION_MAX)
    return {
        "Unhedged": roll_u,
        "LQ (L2)": roll_l,
        "RL (L1)": roll_r
    }

# (1) Baseline NII overlay
rolls_base = get_rolls_for_scenario(X_smooth, "Baseline (smoothed)")
plot_nii_overlay(rolls_base, "Baseline: NII paths (Unhedged vs LQ vs RL-L1)")

# (2) Stress +200bp NII overlay
rolls_par = get_rolls_for_scenario(X_parallel, "Stress: +200bp parallel")
plot_nii_overlay(rolls_par, "Stress (+200bp parallel): NII paths")

# (3) Stress high-vol NII overlay
rolls_hv = get_rolls_for_scenario(X_highvol, "Stress: high vol x3")
plot_nii_overlay(rolls_hv, "Stress (high vol x3): NII paths")

# (4) Baseline hedge inventory overlay (LQ vs RL)
plot_h_overlay({"LQ (L2)": rolls_base["LQ (L2)"], "RL (L1)": rolls_base["RL (L1)"]},
               "Baseline: Hedge inventory (LQ vs RL-L1)")

# (5) Histogram of |u| (baseline) to show sparse trading
plot_absu_histogram(rolls_base["LQ (L2)"], rolls_base["RL (L1)"],
                    "Baseline: Trading sparsity (|u_t| histogram)")

# (6) Risk–turnover scatter using df_metrics (all scenarios/policies)
plot_risk_turnover_scatter(df_metrics, title="Risk–turnover map across scenarios (policy dots)")

=== TABLE 1: Core metrics ===

C:\Users\thoma\AppData\Local\Temp\ipykernel_24368\1485188669.py:85: RuntimeWarning: divide by zero encountered in divide
  df_norm.loc[scen, metrics_cols] = (df.loc[scen, metrics_cols].values / base.values) * 100.0
C:\Users\thoma\AppData\Local\Temp\ipykernel_24368\1485188669.py:85: RuntimeWarning: invalid value encountered in divide
  df_norm.loc[scen, metrics_cols] = (df.loc[scen, metrics_cols].values / base.values) * 100.0
C:\Users\thoma\AppData\Local\Temp\ipykernel_24368\1485188669.py:85: RuntimeWarning: divide by zero encountered in divide
  df_norm.loc[scen, metrics_cols] = (df.loc[scen, metrics_cols].values / base.values) * 100.0
C:\Users\thoma\AppData\Local\Temp\ipykernel_24368\1485188669.py:85: RuntimeWarning: invalid value encountered in divide
  df_norm.loc[scen, metrics_cols] = (df.loc[scen, metrics_cols].values / base.values) * 100.0
C:\Users\thoma\AppData\Local\Temp\ipykernel_24368\1485188669.py:85: RuntimeWarning: divide by zero encountered in divide
  df_norm.loc[scen, metrics_cols] = (df.loc[scen, metrics_cols].values / base.values) * 100.0
C:\Users\thoma\AppData\Local\Temp\ipykernel_24368\1485188669.py:85: RuntimeWarning: invalid value encountered in divide
  df_norm.loc[scen, metrics_cols] = (df.loc[scen, metrics_cols].values / base.values) * 100.0

=== TABLE 2: Normalized to Unhedged (Unhedged=100) ===

# ============================================================
# Monte-Carlo stress testing
# ============================================================
# Deterministic single-path "std" is not a risk measure. Here we draw many noisy
# factor paths per scenario (a persistent initial shock + ongoing N(0, vol_scale*Sigma)
# innovations) and report the DISTRIBUTION of NII risk across paths, with bootstrap
# confidence intervals and pooled tail metrics (5% quantile, CVaR-5%).

def simulate_scenario_paths(n_paths, T, shock_vec=None, vol_scale=1.0, seed=0):
    """N noisy factor paths: X0 = X_ref + shock_vec; X_{t+1} = c + Phi X_t + N(0, vol_scale*Sigma)."""
    rng = np.random.default_rng(seed)
    shock_vec = np.zeros(3) if shock_vec is None else np.asarray(shock_vec, float)
    paths = np.zeros((n_paths, T, 3))
    for i in range(n_paths):
        X = np.zeros((T, 3))
        X[0] = X_ref + shock_vec
        for t in range(T - 1):
            eps = rng.multivariate_normal(np.zeros(3), vol_scale * Sigma_hat)
            X[t + 1] = c_hat + Phi_hat @ X[t] + eps
        paths[i] = X
    return paths


def _boot_ci(vals, n_boot=2000, alpha=0.05, seed=0):
    rng = np.random.default_rng(seed)
    vals = np.asarray(vals, float)
    means = np.array([rng.choice(vals, size=vals.size, replace=True).mean()
                      for _ in range(n_boot)])
    return float(np.quantile(means, alpha / 2)), float(np.quantile(means, 1 - alpha / 2))


def eval_policies_mc(paths, K_lq, rl_model, u_scale=U_SCALE, cvar_q=0.05):
    """
    For each path, roll out Unhedged / LQ / RL and record per-path std(NII).
    Aggregate: mean +- bootstrap CI of per-path std, plus pooled tail metrics
    (5% quantile and CVaR-5%) across all paths.
    """
    policies = {"Unhedged": dict(policy="unhedged"),
                "LQ (L2)":  dict(policy="lq", K_lq=K_lq),
                "RL (L1)":  dict(policy="rl", rl_model=rl_model)}
    per_path_std = {k: [] for k in policies}
    per_path_turn = {k: [] for k in policies}
    pooled_nii = {k: [] for k in policies}

    for i in range(paths.shape[0]):
        Xp = paths[i]
        for name, kw in policies.items():
            res = rollout_policy_on_exogenous_X(cm, Xp, X_ref=X_ref, u_scale=u_scale, **kw)
            per_path_std[name].append(float(np.std(res["NII"])))
            per_path_turn[name].append(float(np.mean(np.abs(res["u"]))))
            pooled_nii[name].append(res["NII"])

    rows = []
    for name in policies:
        stds = np.asarray(per_path_std[name])
        pooled = np.concatenate(pooled_nii[name])
        lo, hi = _boot_ci(stds)
        q = float(np.quantile(pooled, cvar_q))
        cvar = float(pooled[pooled <= q].mean()) if np.any(pooled <= q) else q
        rows.append({
            "policy": name,
            "mean_std_NII": float(stds.mean()),
            "std_CI_low": lo,
            "std_CI_high": hi,
            "p05_NII": q,
            "CVaR05_NII": cvar,
            "mean_turnover": float(np.mean(per_path_turn[name])),
        })
    return pd.DataFrame(rows).set_index("policy")

# Run the Monte-Carlo stress comparison. Uses K_lq (LQ benchmark) and model_l1 (RL under L1 costs).
N_PATHS = 200
T_MC = X_smooth.shape[0]

mc_scenarios = {
    "Baseline":                 dict(shock_vec=None,                          vol_scale=1.0),
    "Stress: +200bp parallel":  dict(shock_vec=np.array([0.02, 0.0, 0.0]),    vol_scale=1.0),
    "Stress: bear steepener":   dict(shock_vec=np.array([0.02, -0.01, 0.0]),  vol_scale=1.0),
    "Stress: high vol x3":      dict(shock_vec=None,                          vol_scale=3.0),
}

mc_tables = {}
for name, cfg in mc_scenarios.items():
    paths = simulate_scenario_paths(N_PATHS, T_MC, seed=12345, **cfg)
    tbl = eval_policies_mc(paths, K_lq, model_l1, u_scale=U_SCALE)
    mc_tables[name] = tbl
    print("=" * 72)
    print(f"{name}  (N={N_PATHS} paths)")
    display(tbl)

# Bar chart: mean per-path std(NII) with bootstrap 95% CI, by scenario and policy.
pols = ["Unhedged", "LQ (L2)", "RL (L1)"]
scen_names = list(mc_tables.keys())
x = np.arange(len(scen_names))
width = 0.25

fig, ax = plt.subplots(figsize=(11, 5))
for j, pol in enumerate(pols):
    means = np.array([mc_tables[s].loc[pol, "mean_std_NII"] for s in scen_names])
    lo = means - np.array([mc_tables[s].loc[pol, "std_CI_low"] for s in scen_names])
    hi = np.array([mc_tables[s].loc[pol, "std_CI_high"] for s in scen_names]) - means
    ax.bar(x + (j - 1) * width, means, width, yerr=[lo, hi], capsize=3, label=pol)

ax.set_xticks(x)
ax.set_xticklabels(scen_names, rotation=15, ha="right")
ax.set_ylabel("Mean per-path std(NII)  (bootstrap 95% CI)")
ax.set_title("Monte-Carlo NII risk by scenario and policy (Issue 4)")
ax.legend()
plt.tight_layout()
plt.show()

========================================================================
Baseline  (N=200 paths)

========================================================================
Stress: +200bp parallel  (N=200 paths)

========================================================================
Stress: bear steepener  (N=200 paths)

========================================================================
Stress: high vol x3  (N=200 paths)

import json, numpy as np, pandas as pd
import matplotlib.pyplot as plt

# ---- single-seed kappa_u sweep ------------------------------------------------
sweep = json.load(open("_sweep/sweep_results.json"))
rows = []
for r in sweep:
    k = r["kappa_u"]
    for pol in ["Unhedged", "LQ", "RL"]:
        d = r[pol]
        rows.append(dict(kappa_u=k, policy=pol, cost=d["cost"], var=d["var"],
                         inv=d.get("inv", np.nan), trade=d["trade"],
                         std=d["std"], turnover=d["turnover"]))
df_sweep = pd.DataFrame(rows)

ratio = (df_sweep.pivot(index="kappa_u", columns="policy", values="cost")
         .assign(**{"RL/LQ": lambda x: x["RL"] / x["LQ"]}))
print("Single-seed kappa_u sweep — total L1 economic cost (300 MC paths)")
display(ratio[["Unhedged", "LQ", "RL", "RL/LQ"]].style.format("{:.3e}".format,
        subset=["Unhedged", "LQ", "RL"]).format("{:.3f}", subset=["RL/LQ"]))

# ---- 3-seed robustness --------------------------------------------------------
seeds = json.load(open("_sweep/seed_results.json"))
srows = []
for kstr, blk in seeds.items():
    k = float(kstr)
    lq = blk["lq_cost"]
    ratios = [blk["seeds"][s]["ratio"] for s in sorted(blk["seeds"])]
    srows.append(dict(kappa_u=k, lq_cost=lq,
                      **{f"seed{i}": ratios[i] for i in range(len(ratios))},
                      mean=np.mean(ratios), std=np.std(ratios)))
df_seed = pd.DataFrame(srows).set_index("kappa_u")
print("\n3-seed robustness — RL/LQ cost ratio (both env & SAC seeded, 800k each)")
display(df_seed.style.format("{:.3f}"))

# ---- figure: ratio vs kappa_u with seed spread --------------------------------
fig, ax = plt.subplots(figsize=(8, 4.5))
ks = sorted(df_sweep["kappa_u"].unique())
ax.plot(ks, [ratio.loc[k, "RL/LQ"] for k in ks], "o-", color="C0",
        label="single-seed sweep (RL/LQ)")
for k in df_seed.index:
    seed_cols = [c for c in df_seed.columns if c.startswith("seed")]
    rs = df_seed.loc[k, seed_cols].values.astype(float)
    ax.scatter([k] * len(rs), rs, color="C3", zorder=5,
               label="3-seed retrains" if k == df_seed.index[0] else None)
    ax.scatter([k], [df_seed.loc[k, "mean"]], marker="_", s=400, color="C3",
               label="3-seed mean" if k == df_seed.index[0] else None)
ax.axhline(1.0, ls="--", color="k", lw=1, label="RL = LQ")
ax.set_xscale("log"); ax.set_xlabel(r"$\kappa_u$ (L1 trading cost)")
ax.set_ylabel("RL / LQ economic-cost ratio  (<1 = RL wins)")
ax.set_title("RL vs LQ: apparent single-seed wins vanish under multi-seed evaluation")
ax.legend(fontsize=8); plt.tight_layout(); plt.show()

Single-seed kappa_u sweep — total L1 economic cost (300 MC paths)

3-seed robustness — RL/LQ cost ratio (both env & SAC seeded, 800k each)

import json, numpy as np, pandas as pd
import matplotlib.pyplot as plt

imp = json.load(open("_sweep/impulse_results.json"))   # premise C: RL vs LQ under fixed costs
cap = json.load(open("_sweep/capstone.json"))          # figure artifacts
kf, cont, band = imp["kappa_fixed"], imp["continuous_LQ"], imp["band_LQ"]
seeds = imp["rl_seeds"]
rl_tot = np.array([seeds[s]["total"] for s in sorted(seeds)])

rows = [("LQ continuous (band 0)", cont["total"], cont["trade_freq"], cont["std"]),
        (f"LQ best fixed band ({imp['best_band']})", band["total"], band["trade_freq"], band["std"])]
for s in sorted(seeds):
    d = seeds[s]; rows.append((f"RL seed {s}", d["total"], d["trade_freq"], d["std"]))
df = pd.DataFrame(rows, columns=["policy", "total cost", "trade freq", "NII std"]).set_index("policy")
df["vs cont %"] = 100 * (1 - df["total cost"] / cont["total"])
df["vs band %"] = 100 * (1 - df["total cost"] / band["total"])
display(df.style.format({"total cost": "{:.3e}", "trade freq": "{:.2f}", "NII std": "{:.4f}",
                         "vs cont %": "{:+.0f}", "vs band %": "{:+.0f}"}))

n = len(rl_tot)
print(f"Fixed cost kappa={kf:.0e}.  RL beats continuous LQ in {int((rl_tot<cont['total']).sum())}/{n} seeds, "
      f"beats the best fixed band in {int((rl_tot<band['total']).sum())}/{n}.  "
      f"Median RL total cost = {np.median(rl_tot):.2e} "
      f"({100*(1-np.median(rl_tot)/cont['total']):.0f}% below continuous LQ).")

fig, ax = plt.subplots(1, 3, figsize=(15, 4.2))
# (A) gate: no-trade band vs total cost
gb, gt = cap["gate_bands"], cap["gate_totals"]
ax[0].plot(gb, gt, "o-", color="C0")
ax[0].axhline(gt[0], ls="--", color="grey", lw=1, label="continuous LQ (band 0)")
ax[0].scatter([gb[int(np.argmin(gt))]], [min(gt)], color="C3", zorder=5, s=60, label="best band")
ax[0].set_xlabel("no-trade band half-width"); ax[0].set_ylabel("total economic cost")
ax[0].set_title(f"(s,S) band beats continuous rebalancing\n(fixed cost $\\kappa$={kf:.0e})"); ax[0].legend(fontsize=8)
# (B) representative trajectory
tj = cap["traj"]; W = slice(0, 180); t = np.arange(len(tj["h_star"]))[W]
ax[1].plot(t, np.array(tj["h_star"])[W], color="grey", lw=1.0, label="variance-min target $h^*$")
ax[1].plot(t, np.array(tj["h_cont"])[W], color="C0", lw=0.8, alpha=.7, label=f"LQ continuous (trades {tj['trades_cont']:.0%})")
ax[1].plot(t, np.array(tj["h_rl"])[W], color="C3", lw=1.6, label=f"RL (s,S) (trades {tj['trades_rl']:.0%})")
ax[1].set_xlabel("month"); ax[1].set_ylabel("hedge notional $h_t$")
ax[1].set_title("RL learns to trade rarely, not every month"); ax[1].legend(fontsize=8)
# (C) total cost bars
labels = ["LQ\ncontinuous", f"LQ band\n({imp['best_band']})", "RL\n(median)"]
vals = [cont["total"], band["total"], float(np.median(rl_tot))]
ax[2].bar(labels, vals, color=["C0", "C1", "C3"])
ax[2].set_ylabel("total economic cost")
ax[2].set_title("Total cost = NII variance + trading frictions")
plt.tight_layout(); plt.show()

Fixed cost kappa=4e-03.  RL beats continuous LQ in 5/6 seeds, beats the best fixed band in 5/6.  Median RL total cost = 2.36e-03 (66% below continuous LQ).

	1.0	2.0	3.0	4.0	5.0	6.0	7.0	8.0	9.0	10.0	...	21.0	22.0	23.0	24.0	25.0	26.0	27.0	28.0	29.0	30.0
date
1990-01-31	0.080998	0.081567	0.082178	0.082620	0.082924	0.083137	0.083292	0.083409	0.083500	0.083573	...	0.083916	0.083930	0.083943	0.083955	0.083966	0.083976	0.083985	0.083994	0.084002	0.084010
1990-02-28	0.080925	0.082517	0.083358	0.083810	0.084082	0.084263	0.084392	0.084489	0.084564	0.084624	...	0.084907	0.084918	0.084929	0.084939	0.084948	0.084956	0.084964	0.084971	0.084977	0.084984
1990-03-31	0.083192	0.084778	0.085420	0.085639	0.085698	0.085701	0.085687	0.085670	0.085654	0.085640	...	0.085570	0.085567	0.085565	0.085562	0.085560	0.085558	0.085556	0.085554	0.085553	0.085551
1990-04-30	0.085684	0.087799	0.088629	0.088956	0.089093	0.089156	0.089190	0.089210	0.089225	0.089235	...	0.089284	0.089286	0.089288	0.089290	0.089291	0.089293	0.089294	0.089295	0.089296	0.089297
1990-05-31	0.081363	0.083009	0.084012	0.084568	0.084881	0.085058	0.085157	0.085208	0.085232	0.085239	...	0.085137	0.085130	0.085123	0.085117	0.085111	0.085106	0.085101	0.085096	0.085092	0.085088

	level	slope	curvature
1990-01-31	0.078016	0.003215	0.014766
1990-02-28	0.075872	0.006037	0.020238
1990-03-31	0.079267	0.005030	0.012670
1990-04-30	0.077276	0.009562	0.023612
1990-05-31	0.072786	0.009663	0.024977

	1.0	2.0	3.0	5.0	7.0	10.0	20.0	30.0
date
1990-01-31	0.027290	0.020870	0.023581	0.040637	0.072917	0.168779	1.818128	-430.320652
1990-02-28	0.026877	0.020070	0.022486	0.038854	0.070082	0.162952	1.734290	-431.862708
1990-03-31	0.027016	0.019273	0.020994	0.035967	0.065246	0.152740	1.585390	-434.602759
1990-04-30	0.026891	0.018409	0.019664	0.033667	0.061521	0.145011	1.473642	-436.658533
1990-05-31	0.025187	0.017341	0.018848	0.033008	0.060822	0.143961	1.461309	-436.883605

	lambda_u	lambda_h	K_max_abs	std_unhedged	std_hedged	std_reduction_%	min_unhedged	min_hedged	p05_unhedged	p05_hedged	mean_unhedged	mean_hedged	mean_abs_h	max_abs_h	mean_abs_u	max_abs_u	turnover_proxy	inventory_proxy
0	1.000000e-09	1.000000e-03	0.999999	0.090352	0.027062	70.048187	-0.281366	-0.267363	-0.214973	-0.119902	-0.080764	-0.081215	40.996852	41.091752	0.095120	41.091711	0.095120	40.996852
1	1.000000e-09	1.000000e-06	0.999040	0.090352	0.027062	70.048029	-0.281366	-0.267363	-0.214973	-0.119902	-0.080764	-0.081215	40.996761	41.091752	0.095120	41.052285	0.095120	40.996761
2	1.000000e-06	1.000000e-03	0.999002	0.090352	0.027062	70.048023	-0.281366	-0.267363	-0.214973	-0.119902	-0.080764	-0.081215	40.996757	41.091752	0.095120	41.050744	0.095120	40.996757
3	1.000000e-09	1.000000e-09	0.976264	0.090352	0.027066	70.043392	-0.281366	-0.267363	-0.214973	-0.119902	-0.080764	-0.081225	40.994545	41.091752	0.095120	40.116413	0.095120	40.994545
4	1.000000e-06	1.000000e-06	0.624588	0.090352	0.027368	69.709015	-0.281366	-0.267363	-0.214973	-0.120923	-0.080764	-0.081465	40.939812	41.091752	0.095120	25.665431	0.095120	40.939812
5	1.000000e-03	1.000000e-03	0.618041	0.090352	0.027380	69.696416	-0.281366	-0.267363	-0.214973	-0.120923	-0.080764	-0.081472	40.938202	41.091752	0.095120	25.396374	0.095120	40.938202
6	1.000000e-06	1.000000e-09	0.181312	0.090352	0.029812	67.004217	-0.281366	-0.267363	-0.214973	-0.125203	-0.080764	-0.082946	40.568344	41.091752	0.095120	7.450423	0.095120	40.568344
7	1.000000e-03	1.000000e-06	0.031721	0.090352	0.034598	61.707422	-0.281366	-0.267363	-0.214973	-0.134125	-0.080764	-0.088107	38.099998	41.091715	0.095120	1.303452	0.095120	38.099998
8	1.000000e-03	1.000000e-09	0.006317	0.090352	0.048593	46.218237	-0.281366	-0.277733	-0.214973	-0.159986	-0.080764	-0.094081	27.034485	38.431733	0.088962	0.259565	0.088962	27.034485

	lambda_u	lambda_h	K_max_abs	std_unhedged	std_hedged	std_reduction_%	p05_unhedged	p05_hedged	min_unhedged	min_hedged	mean_unhedged	mean_hedged	mean_abs_h	max_abs_h	mean_abs_u	max_abs_u
0	0.0001	1.000000e-05	0.270602	0.090352	0.028806	68.118500	-0.214973	-0.124468	-0.281366	-0.267363	-0.080764	-0.082304	40.741052	41.091752	0.095120	11.119507
1	0.0001	1.000000e-06	0.096875	0.090352	0.031575	65.053047	-0.214973	-0.127158	-0.281366	-0.267363	-0.080764	-0.084280	40.112143	41.091752	0.095120	3.980783
2	0.0010	1.000000e-05	0.095302	0.090352	0.031622	65.001393	-0.214973	-0.127158	-0.281366	-0.267363	-0.080764	-0.084320	40.095966	41.091752	0.095120	3.916112
3	0.0001	1.000000e-07	0.036614	0.090352	0.034194	62.155001	-0.214973	-0.131285	-0.281366	-0.267363	-0.080764	-0.087425	38.499854	41.091748	0.095120	1.504539
4	0.0010	1.000000e-06	0.031721	0.090352	0.034598	61.707422	-0.214973	-0.134125	-0.281366	-0.267363	-0.080764	-0.088107	38.099998	41.091715	0.095120	1.303452
5	0.0100	1.000000e-05	0.031187	0.090352	0.034649	61.651508	-0.214973	-0.134447	-0.281366	-0.267363	-0.080764	-0.088192	38.048782	41.091706	0.095120	1.281514
6	0.0010	1.000000e-07	0.011727	0.090352	0.040334	55.359414	-0.214973	-0.143385	-0.281366	-0.274659	-0.080764	-0.093469	33.048262	40.840223	0.094538	0.481883
7	0.0100	1.000000e-06	0.010142	0.090352	0.041919	53.605021	-0.214973	-0.146817	-0.281366	-0.275556	-0.080764	-0.093931	31.847964	40.589139	0.093956	0.416755
8	0.0100	1.000000e-07	0.003723	0.090352	0.057830	35.994237	-0.214973	-0.178365	-0.281366	-0.279219	-0.080764	-0.092117	20.672189	32.890339	0.076135	0.153000

	Date	BETA0	BETA1	BETA2	SVEN1F01	SVEN1F04	SVEN1F09	SVENF01	SVENF02	...	SVENY23	SVENY24	SVENY25	SVENY26	SVENY27	SVENY28	SVENY29	SVENY30	TAU1	TAU2
0	1961-06-14	3.917606	-1.277955	-1.949397	3.8067	3.9562	NaN	3.5492	3.8825	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	0.339218	-999.99
1	1961-06-15	3.978498	-1.257404	-2.247617	3.8694	4.0183	NaN	3.5997	3.9460	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	0.325775	-999.99
2	1961-06-16	3.984350	-1.429538	-1.885024	3.8634	4.0242	NaN	3.5957	3.9448	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	0.348817	-999.99
3	1961-06-19	4.004379	-0.723311	-3.310743	3.9196	4.0447	NaN	3.6447	3.9842	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	0.282087	-999.99
4	1961-06-20	3.985789	-0.900432	-2.844809	3.8732	4.0257	NaN	3.5845	3.9552	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	0.310316	-999.99

	SVENY01	SVENY02	SVENY03	SVENY04	SVENY05	SVENY06	SVENY07	SVENY08	SVENY09	SVENY10	...	SVENY21	SVENY22	SVENY23	SVENY24	SVENY25	SVENY26	SVENY27	SVENY28	SVENY29	SVENY30
date
1961-06-14	2.9825	3.3771	3.5530	3.6439	3.6987	3.7351	3.7612	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-06-15	2.9941	3.4137	3.5981	3.6930	3.7501	3.7882	3.8154	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-06-16	3.0012	3.4142	3.5994	3.6953	3.7531	3.7917	3.8192	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-06-19	2.9949	3.4386	3.6252	3.7199	3.7768	3.8147	3.8418	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-06-20	2.9833	3.4101	3.5986	3.6952	3.7533	3.7921	3.8198	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN

	1.0	2.0	3.0	4.0	5.0	6.0	7.0	8.0	9.0	10.0	...	21.0	22.0	23.0	24.0	25.0	26.0	27.0	28.0	29.0	30.0
date
1961-06-14	0.029825	0.033771	0.035530	0.036439	0.036987	0.037351	0.037612	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-06-15	0.029941	0.034137	0.035981	0.036930	0.037501	0.037882	0.038154	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-06-16	0.030012	0.034142	0.035994	0.036953	0.037531	0.037917	0.038192	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-06-19	0.029949	0.034386	0.036252	0.037199	0.037768	0.038147	0.038418	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-06-20	0.029833	0.034101	0.035986	0.036952	0.037533	0.037921	0.038198	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN

	1.0	2.0	3.0	4.0	5.0	6.0	7.0	8.0	9.0	10.0	...	21.0	22.0	23.0	24.0	25.0	26.0	27.0	28.0	29.0	30.0
date
1961-06-30	0.029011	0.032795	0.035036	0.036316	0.037109	0.037640	0.038020	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-07-31	0.027780	0.032304	0.035068	0.036787	0.037907	0.038678	0.039234	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-08-31	0.029863	0.033990	0.036481	0.037919	0.038812	0.039412	0.039841	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-09-30	0.029358	0.033250	0.035412	0.036661	0.037442	0.037968	0.038345	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN
1961-10-31	0.028936	0.032396	0.034616	0.036087	0.037096	0.037813	0.038339	NaN	NaN	NaN	...	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN	NaN

	scenario	lambda_u	lambda_h	std_unhedged	std_hedged	std_reduction_%	p05_unhedged	p05_hedged	min_unhedged	min_hedged	mean_unhedged	mean_hedged	mean_abs_h	max_abs_h	mean_abs_u	max_abs_u
0	baseline (smoothed)	0.0001	1.000000e-07	0.090352	0.034194	62.155001	-0.214973	-0.131285	-0.281366	-0.267363	-0.080764	-0.087425	38.499854	41.091748	0.09512	1.504539
1	stress: +200bp parallel	0.0001	1.000000e-07	0.020020	0.011062	44.748189	-0.108548	-0.088227	-0.147521	-0.147521	-0.062789	-0.081304	38.499854	41.091748	0.09512	1.504539
2	stress: bear steepener	0.0001	1.000000e-07	0.016309	0.010078	38.209566	-0.073737	-0.087777	-0.108608	-0.108608	-0.055667	-0.078735	38.499854	41.091748	0.09512	1.504539
3	stress: high vol x3	0.0001	1.000000e-07	0.059293	0.032767	44.736249	-0.188907	-0.169675	-0.230237	-0.197051	-0.095596	-0.101848	38.499854	41.091748	0.09512	1.504539

		std_NII	p05_NII	min_NII	mean_NII	mean_abs_u	max_abs_u	mean_abs_h	max_abs_h	L1_cost
scenario	policy
Baseline (smoothed)	LQ (L2)	0.034194	-0.131285	-0.267363	-0.087425	0.095120	1.504539	38.499854	41.091748	0.008994
	RL (L1)	0.024854	-0.113054	-0.267363	-0.073811	0.604044	22.924119	42.376044	53.488478	0.006431
	Unhedged	0.090352	-0.214973	-0.281366	-0.080764	0.000000	0.000000	0.000000	0.000000	0.014686
Stress: +200bp parallel	LQ (L2)	0.011062	-0.088227	-0.147521	-0.081304	0.095120	1.504539	38.499854	41.091748	0.006914
	RL (L1)	0.008841	-0.086189	-0.147521	-0.077409	0.108291	11.489066	35.984236	38.595587	0.006233
	Unhedged	0.020020	-0.108548	-0.147521	-0.062789	0.000000	0.000000	0.000000	0.000000	0.004343
Stress: high vol x3	LQ (L2)	0.032767	-0.169675	-0.197051	-0.101848	0.095120	1.504539	38.499854	41.091748	0.011628
	RL (L1)	0.035016	-0.177900	-0.201576	-0.098451	0.608999	9.825009	42.257005	54.273592	0.011284
	Unhedged	0.059293	-0.188907	-0.230237	-0.095596	0.000000	0.000000	0.000000	0.000000	0.012654

	mean_std_NII	std_CI_low	std_CI_high	p05_NII	CVaR05_NII	mean_turnover
policy
Unhedged	0.059981	0.057534	0.062452	-0.187804	-0.222097	0.000000
LQ (L2)	0.020693	0.019954	0.021468	-0.124007	-0.134397	0.095120
RL (L1)	0.018174	0.017604	0.018747	-0.114930	-0.123653	0.432978

	mean_std_NII	std_CI_low	std_CI_high	p05_NII	CVaR05_NII	mean_turnover
policy
Unhedged	0.062761	0.060299	0.065374	-0.181598	-0.214378	0.00000
LQ (L2)	0.023181	0.022381	0.023953	-0.125773	-0.138188	0.09512
RL (L1)	0.019133	0.018554	0.019717	-0.113787	-0.123807	0.43086

	mean_std_NII	std_CI_low	std_CI_high	p05_NII	CVaR05_NII	mean_turnover
policy
Unhedged	0.061748	0.059125	0.064538	-0.177427	-0.211559	0.000000
LQ (L2)	0.022804	0.021927	0.023644	-0.121076	-0.131687	0.095120
RL (L1)	0.019001	0.018367	0.019640	-0.112508	-0.121649	0.435617

	mean_std_NII	std_CI_low	std_CI_high	p05_NII	CVaR05_NII	mean_turnover
policy
Unhedged	0.103814	0.099658	0.108052	-0.275759	-0.335187	0.000000
LQ (L2)	0.035528	0.034248	0.036857	-0.152461	-0.170352	0.095120
RL (L1)	0.029953	0.028955	0.030928	-0.131485	-0.147547	0.656499

		std_NII	p05_NII	min_NII	mean_NII	mean_abs_u	max_abs_u	mean_abs_h	max_abs_h	L1_cost
scenario	policy
Baseline (smoothed)	LQ (L2)	37.844999	61.070514	95.023371	108.247757	inf	inf	inf	inf	61.239919
	RL (L1)	27.508079	52.589671	95.023371	91.390947	inf	inf	inf	inf	43.788940
	Unhedged	100.000000	100.000000	100.000000	100.000000	NaN	NaN	NaN	NaN	100.000000
Stress: +200bp parallel	LQ (L2)	55.251811	81.279512	100.000000	129.487453	inf	inf	inf	inf	159.192269
	RL (L1)	44.157799	79.402264	100.000000	123.283909	inf	inf	inf	inf	143.511067
	Unhedged	100.000000	100.000000	100.000000	100.000000	NaN	NaN	NaN	NaN	100.000000
Stress: high vol x3	LQ (L2)	55.263751	89.819266	85.586050	106.539481	inf	inf	inf	inf	91.891075
	RL (L1)	59.056446	94.173220	87.551192	102.986308	inf	inf	inf	inf	89.173138
	Unhedged	100.000000	100.000000	100.000000	100.000000	NaN	NaN	NaN	NaN	100.000000

policy	Unhedged	LQ	RL	RL/LQ
kappa_u
0.000300	5.214e-03	2.129e-03	3.729e-03	1.752
0.001000	5.214e-03	2.195e-03	1.949e-03	0.888
0.003000	5.214e-03	2.384e-03	2.330e-03	0.977
0.010000	5.214e-03	3.047e-03	3.331e-03	1.093

	lq_cost	seed0	seed1	seed2	mean	std
kappa_u
0.001000	0.002	2.077	0.893	0.948	1.306	0.546
0.003000	0.002	1.049	0.948	3.198	1.732	1.038

Deformation	Optimal hedge	Winner	Why
Base one-period NII hedge	constant ratio	tie	near-static. Riccati is optimal
A: rate-dependent balance sheet	linear in state	LQ	hedging is covariance → linear; risk drivers are observable in the state
B: downside / CVaR objective	linear in state	LQ	only shifts the linear hedge ratio, doesn't make it nonlinear
+ Instrument: add an interest-rate floor	linear in state	LQ	the swap, sized on state, already spans the (observable) risk
C: fixed-cost / impulse frictions	nonlinear (s,S) band	RL	a smooth quadratic control cost cannot encode a no-trade region

	total cost	trade freq	NII std	vs cont %	vs band %
policy
LQ continuous (band 0)	6.915e-03	1.00	0.0456	+0	-111
LQ best fixed band (13)	3.276e-03	0.04	0.0473	+53	+0
RL seed 0	8.120e-03	0.52	0.0636	-17	-148
RL seed 1	2.334e-03	0.10	0.0389	+66	+29
RL seed 2	2.442e-03	0.10	0.0386	+65	+25
RL seed 3	2.386e-03	0.10	0.0369	+66	+27
RL seed 4	2.310e-03	0.12	0.0368	+67	+29
RL seed 5	2.272e-03	0.08	0.0381	+67	+31

Reinforcement Learning and Optimal Control for IRRBB Hedging Under Uncertainty¶

1) Yield curve dataset (GSW) and preprocessing¶

2) DNS (Diebold–Li) factor model: cross-sectional OLS¶

3) Time-series dynamics for the factors¶

State (Transition) Equation¶

Measurement (Observation) Equation¶

4) Measurement equation and residual covariance¶

5) Kalman filtering and smoothing (DNS)¶

6) Applying the Kalman filter/smoother¶

7) OLS vs Kalman factors (what should we expect?)¶

8) Fully latent state-space estimation via MLE¶

9) Initialization choices for MLE¶

10) Maximum likelihood estimation¶

11) Estimated parameters and interpretation¶

12) Smoothed factors under MLE¶

Part II: AFNS (Arbitrage-Free Nelson–Siegel)¶

From Diebold–Li to AFNS: No-Arbitrage Term Structure Modeling¶

The Diebold–Li (Dynamic Nelson–Siegel) Model¶

$ y_t(\tau)¶

L_t + S_t \frac{1 - e^{-\lambda \tau}}{\lambda \tau} + C_t \left( \frac{1 - e^{-\lambda \tau}}{\lambda \tau}¶

$ \mathbf{X}_{t+1}¶

The Key Limitation: Lack of No-Arbitrage¶

Risk-Neutral Pricing and No-Arbitrage¶

$ P_t(\tau)¶

$ y_t(\tau)¶

The AFNS Model: Making Nelson–Siegel Arbitrage-Free¶

Short Rate Specification¶

Risk-Neutral Dynamics¶

$ d\mathbf{X}_t¶

The Yield Adjustment Term¶

Observed yields satisfy: $ y_t(\tau)¶

Relationship Between DNS and AFNS¶

Why AFNS is important in this project¶

Conceptual Summary¶

AFNS approach and implementation used here¶

AFNS approach used here¶

Utility functions: DNS yields and parameter vector¶

CurveModelAFNS: reusable curve model object¶

AFNS estimation step¶

Maturities and yield matrix used for AFNS¶

Sanity-checking the fitted AFNS model¶

State space: filtering and smoothing for AFNS¶

Predicted vs filtered vs smoothed states¶

AFNS filtering/smoothing output¶

Part III: Banking book and NII hedging experiment design¶

From State-Space Modeling to Optimal Control: Kalman Filtering and LQ Control¶

State-Space Structure as the Unifying Framework¶

State (transition) equation¶

$ \mathbf{X}_{t+1}¶

Measurement equation¶

$ \mathbf{y}_t¶

Kalman Filtering: Inference on Latent Yield Factors¶

Prediction¶

Filtering¶

After observing yields $\mathbf{y}_t$, the prediction is updated: $ \hat{\mathbf{X}}_{t|t}¶

Smoothing¶

From Estimation to Decision-Making: Augmenting the State¶

Optimal Control Theory in This Context¶

Linear–Quadratic (LQ) Control with L2 Costs¶

Quadratic Objective¶

Optimal Policy¶

Why LQ Control Is a Natural Benchmark¶

Role in This Project¶

Conceptual Summary¶

LQ control benchmark (quadratic costs)¶

Choosing penalties: risk–turnover trade-off curve¶

Simplified NII model and hedge instrument¶

Stress testing protocol¶

Policy evaluation on identical scenarios¶

Part IV: Reinforcement Learning¶

Beyond Quadratic Costs: L1 Transaction Costs and the Transition to Reinforcement Learning¶

Economic Motivation for L1 Transaction Costs¶

The Control Problem with L1 Costs¶

What Breaks in Classical Optimal Control¶

State-Space Dynamics Remain Valid¶

Why Reinforcement Learning Is Appropriate¶

Relationship Between LQ Control and RL¶

Role in This Project¶

Conceptual Summary¶

L1 transaction costs (bid–ask / fees): a realistic non-quadratic objective¶

`CurveModelAFNS`: reusable curve model object¶