repricing.py

"""Cross-model repricing and RMSE comparison

Compares market implied volatilities against:
1. Flat Black-Scholes (using ATM vol)
2. Local Volatility (Dupire via Crank-Nicolson PDE)
3. Stochastic Volatility (Heston via Lewis CF)
"""

import numpy as np
import pandas as pd
from scipy.interpolate import RegularGridInterpolator

from black_scholes import bs_price, implied_vol
from heston import HestonParameters, HestonPricerCF
from market_data import MarketSnapshot


def calculate_rmse(iv_market: np.ndarray, iv_model: np.ndarray) -> float:
    """Calculates the Root Mean Square Error, ignoring NaNs"""
    mask = ~np.isnan(iv_market) & ~np.isnan(iv_model)
    if not np.any(mask):
        return np.nan
    return np.sqrt(np.mean((iv_market[mask] - iv_model[mask]) ** 2))


def reprice_pde_to_iv(
    snapshot: MarketSnapshot, k_grid: np.ndarray, tau_grid: np.ndarray, V_surface: np.ndarray
) -> np.ndarray:
    """Interpolates PDE call prices and converts back to implied volatility"""
    # V_surface from DupireForwardPricer is shape (nk, nt)
    interpolator = RegularGridInterpolator(
        (k_grid, tau_grid), V_surface, bounds_error=False, fill_value=np.nan
    )

    pts = np.column_stack((snapshot.log_moneyness, snapshot.tau))
    call_prices_model = interpolator(pts)

    # Put-Call parity for puts
    forward = snapshot.spot * np.exp(snapshot.r * snapshot.tau)
    strikes = forward * np.exp(snapshot.log_moneyness)

    model_prices = np.where(
        snapshot.option_type == "call",
        call_prices_model,
        call_prices_model - snapshot.spot + (strikes * np.exp(-snapshot.r * snapshot.tau)),
    )

    model_iv = np.array(
        [
            implied_vol(price, K, tau, k, otype, snapshot.spot, snapshot.r)
            for price, K, tau, k, otype in zip(
                model_prices,
                strikes,
                snapshot.tau,
                snapshot.log_moneyness,
                snapshot.option_type,
                strict=True,
            )
        ]
    )
    return model_iv


def reprice_heston_to_iv(snapshot: MarketSnapshot, params: HestonParameters) -> np.ndarray:
    """Prices the whole chain under Heston and converts back to IV"""
    pricer = HestonPricerCF(params)
    model_iv = np.zeros(len(snapshot.tau))

    # Forward calculation
    forward = snapshot.spot * np.exp(snapshot.r * snapshot.tau)
    strikes = forward * np.exp(snapshot.log_moneyness)

    for i in range(len(snapshot.tau)):
        tau, k, otype = snapshot.tau[i], snapshot.log_moneyness[i], snapshot.option_type[i]

        price = pricer.call_price(tau, k) if otype == "call" else pricer.put_price(tau, k)

        model_iv[i] = implied_vol(price, strikes[i], tau, k, otype, snapshot.spot, snapshot.r)

    return model_iv


def run_model_comparison(
    snapshot: MarketSnapshot,
    iv_market: np.ndarray,
    atm_vol: float,
    k_grid: np.ndarray,
    tau_grid: np.ndarray,
    pde_prices: np.ndarray,
    heston_params: HestonParameters,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Reprice the chain under all three models

    Returns the global RMSE table and the bucketed RMSE table (by maturity and
    moneyness). Repricing is done in implied-vol space against the market IV
    """
    iv_bs = np.full_like(iv_market, atm_vol)
    iv_dupire = reprice_pde_to_iv(snapshot, k_grid, tau_grid, pde_prices)
    iv_heston = reprice_heston_to_iv(snapshot, heston_params)

    # global RMSE
    global_rmse = pd.DataFrame(
        {
            "Model": ["Black-Scholes (ATM vol)", "Dupire Local Vol (PDE)", "Heston Stoch Vol (CF)"],
            "Global RMSE": [
                calculate_rmse(iv_market, iv_bs),
                calculate_rmse(iv_market, iv_dupire),
                calculate_rmse(iv_market, iv_heston),
            ],
        }
    )
    global_rmse["Global RMSE"] = global_rmse["Global RMSE"].apply(lambda x: f"{x:.2%}")

    # bucketed RMSE by maturity and moneyness
    errors = pd.DataFrame(
        {
            "tau": snapshot.tau,
            "k": snapshot.log_moneyness,
            "err_bs": (iv_market - iv_bs) ** 2,
            "err_dupire": (iv_market - iv_dupire) ** 2,
            "err_heston": (iv_market - iv_heston) ** 2,
        }
    )
    errors["Maturity"] = pd.cut(
        errors["tau"],
        bins=[0, 0.3, 1.0, 5.0],
        labels=["Short (<3M)", "Medium (3M-1Y)", "Long (>1Y)"],
    )
    errors["Moneyness"] = pd.cut(
        errors["k"],
        bins=[-np.inf, -0.05, 0.05, np.inf],
        labels=["ITM Put / OTM Call", "ATM", "OTM Put / ITM Call"],
    )

    bucketed = (
        errors.groupby(["Maturity", "Moneyness"], observed=True)[
            ["err_bs", "err_dupire", "err_heston"]
        ]
        .mean()
        .apply(np.sqrt)
    )
    bucketed = bucketed.map(lambda x: f"{x:.2%}" if pd.notnull(x) else "NaN")
    bucketed.columns = ["BS RMSE", "Dupire RMSE", "Heston RMSE"]

    return global_rmse, bucketed


def benchmark_price_table(
    snapshot: MarketSnapshot,
    k_grid: np.ndarray,
    tau_grid: np.ndarray,
    pde_prices: np.ndarray,
    heston_params: HestonParameters,
    atm_vol: float,
) -> pd.DataFrame:
    """Build a price comparison table on a few benchmark quotes

    Picks the market quote closest to each (target maturity, target moneyness)
    reference point and reports market mid vs each model price, with the
    market-minus-model difference for every model
    """
    targets = [
        ("3M ATM", 0.25, 0.0),
        ("1Y ATM", 1.0, 0.0),
        ("3M 10% OTM put", 0.25, -0.10),
        ("1Y 10% OTM put", 1.0, -0.10),
        ("3M 10% OTM call", 0.25, 0.10),
    ]

    interp = RegularGridInterpolator(
        (k_grid, tau_grid), pde_prices, bounds_error=False, fill_value=np.nan
    )
    heston = HestonPricerCF(heston_params)
    spot, r = snapshot.spot, snapshot.r

    rows = []
    for label, tau_t, k_t in targets:
        # restrict to the OTM side matching the target (puts for k<0, calls for k>0)
        want_put = k_t < 0
        side = snapshot.option_type == ("put" if want_put else "call")
        dist = np.where(
            side,
            (snapshot.tau - tau_t) ** 2 + (snapshot.log_moneyness - k_t) ** 2,
            np.inf,
        )
        i = int(np.argmin(dist))
        tau, k, otype = snapshot.tau[i], snapshot.log_moneyness[i], snapshot.option_type[i]
        forward = spot * np.exp(r * tau)
        strike = forward * np.exp(k)

        price_market = snapshot.mid[i]

        # Black-Scholes at the flat ATM vol
        price_bs = float(bs_price(strike, tau, k, otype, spot, r, atm_vol))

        # Dupire: interpolate the call surface, put-call parity for puts
        call_pde = float(interp([[k, tau]])[0])
        price_dupire = call_pde if otype == "call" else call_pde - spot + strike * np.exp(-r * tau)

        # Heston
        price_heston = heston.call_price(tau, k) if otype == "call" else heston.put_price(tau, k)

        rows.append(
            {
                "Benchmark": label,
                "Type": otype,
                "tau": round(float(tau), 2),
                "Market": round(price_market, 2),
                "BS": round(price_bs, 2),
                "Dupire": round(price_dupire, 2),
                "Heston": round(price_heston, 2),
                "Mkt-BS": round(price_market - price_bs, 2),
                "Mkt-Dupire": round(price_market - price_dupire, 2),
                "Mkt-Heston": round(price_market - price_heston, 2),
            }
        )

    return pd.DataFrame(rows)