xsabr.py

#!/usr/bin/env python3
"""
Fit SABR smile per expiration from European option prices in a CSV file.

Expected columns (minimum):
  expiration, strike, option_type, bid, ask
Optional: lastPrice

Usage:
  python xsabr.py options.csv [--expiry YYYY-MM-DD|YYYYMMDD:YYYYMMDD] [--as-of YYYY-MM-DD]
      [--fwd-range L:U] [--both-sides] [--beta B] [--fit-beta]
      [--outfile out.csv] [--plot [file.png]] [--plot-density [file.png]]
      [--plot-density-log-s]
"""

from __future__ import annotations

import sys
from datetime import date, datetime, timedelta

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

from black_scholes import implied_vol
from option_utils import infer_forward_df, normalize_expiry, parse_date, price_from_row
from sabr import fit_sabr, sabr_iv_hagan


def main(argv: list[str]) -> None:
    """CLI entrypoint."""
    if len(argv) < 2:
        print("usage: python xsabr.py options.csv [--expiry YYYY-MM-DD|YYYYMMDD:YYYYMMDD] [--as-of YYYY-MM-DD]")
        print("       [--fwd-range L:U] [--both-sides] [--beta B] [--fit-beta] [--outfile out.csv]")
        print("       [--plot [file.png]] [--plot-density [file.png]] [--plot-density-log-s]")
        sys.exit(1)

    infile = argv[1]
    as_of = None
    outfile = None
    plot = "--plot" in argv
    plot_file = None
    plot_density = "--plot-density" in argv
    plot_density_log_s = "--plot-density-log-s" in argv
    if plot_density_log_s:
        plot_density = True
    plot_density_file = None
    otm_only = "--both-sides" not in argv
    fwd_range = None
    expiry = None
    expiry_list = None
    expiry_range = None
    beta = 1.0
    fit_beta = "--fit-beta" in argv
    print_per_expiry = False

    if "--as-of" in argv:
        as_of = normalize_expiry(argv[argv.index("--as-of") + 1])
    if "--outfile" in argv:
        outfile = argv[argv.index("--outfile") + 1]
    if "--plot" in argv:
        pidx = argv.index("--plot")
        if pidx + 1 < len(argv) and not argv[pidx + 1].startswith("--"):
            plot_file = argv[pidx + 1]
    if "--plot-density" in argv:
        pidx = argv.index("--plot-density")
        if pidx + 1 < len(argv) and not argv[pidx + 1].startswith("--"):
            plot_density_file = argv[pidx + 1]
    if "--plot-density-log-s" in argv:
        pidx = argv.index("--plot-density-log-s")
        if pidx + 1 < len(argv) and not argv[pidx + 1].startswith("--"):
            plot_density_file = argv[pidx + 1]
    if "--fwd-range" in argv:
        spec = argv[argv.index("--fwd-range") + 1]
        if ":" not in spec:
            raise ValueError("fwd-range must be L:U")
        left, right = spec.split(":", 1)
        fwd_range = (float(left), float(right))
    if "--expiry" in argv:
        idx = argv.index("--expiry") + 1
        vals = []
        while idx < len(argv) and not argv[idx].startswith("--"):
            vals.append(argv[idx])
            idx += 1
        if len(vals) == 1:
            token = vals[0]
            if ":" in token:
                expiry_range = token
            else:
                expiry = normalize_expiry(token)
        elif len(vals) > 1:
            expiry_list = [normalize_expiry(v) for v in vals]
    if "--beta" in argv:
        beta = float(argv[argv.index("--beta") + 1])

    df = pd.read_csv(infile)
    lower_cols = {c.lower(): c for c in df.columns}
    has_bid = "bid" in lower_cols
    has_ask = "ask" in lower_cols
    if has_bid and has_ask:
        if lower_cols["bid"] != "bid" or lower_cols["ask"] != "ask":
            df = df.rename(columns={lower_cols["bid"]: "bid", lower_cols["ask"]: "ask"})
    else:
        price_cols = [c for c in df.columns if "price" in c.lower()]
        if not price_cols:
            missing = []
            if not has_bid:
                missing.append("bid")
            if not has_ask:
                missing.append("ask")
            print(f"error: missing columns {', '.join(missing)} and no price column found", file=sys.stderr)
            sys.exit(1)
        price_col = None
        for c in price_cols:
            if c.lower() in ("lastprice", "price"):
                price_col = c
                break
        if price_col is None:
            price_col = price_cols[0]
        if price_col != "lastPrice":
            df = df.rename(columns={price_col: "lastPrice"})
    symbol = None
    if "contractSymbol" in df.columns:
        first = df["contractSymbol"].dropna().astype(str).head(1)
        if not first.empty and len(first.iloc[0]) >= 3:
            symbol = first.iloc[0][:3]
    df["expiration"] = df["expiration"].astype(str).apply(normalize_expiry)
    df["option_type"] = df["option_type"].astype(str).str.lower()
    df["price"] = df.apply(price_from_row, axis=1, args=(True,))
    df = df.dropna(subset=["price", "strike", "option_type"])

    t0 = parse_date(as_of) if as_of else date.today()
    results = []
    fit_rows = []
    plot_data = []
    density_data = []

    expiries = sorted(df["expiration"].unique())
    if expiry_range is not None:
        start_s, end_s = expiry_range.split(":", 1)
        start_d = parse_date(start_s)
        end_d = parse_date(end_s)
        if end_d < start_d:
            start_d, end_d = end_d, start_d
        expiries = [e for e in expiries if parse_date(e) >= start_d and parse_date(e) <= end_d]
    if expiry is not None:
        expiries = [expiry] if expiry in expiries else []
    if expiry_list is not None:
        expiries = [e for e in expiries if e in set(expiry_list)]

    for exp in expiries:
        work = df[df["expiration"] == exp].copy()
        if work.empty:
            continue
        try:
            forward, dfactor = infer_forward_df(work)
        except Exception as e:
            print(f"{exp}: parity inference failed ({e})", file=sys.stderr)
            continue

        t1 = parse_date(exp)
        if t1 == t0:
            t0_eff = t0 - timedelta(days=1)
        else:
            t0_eff = t0
        t = (t1 - t0_eff).days / 365.25
        if t <= 0:
            print(f"{exp}: non-positive time to expiry", file=sys.stderr)
            continue
        rate = -np.log(dfactor) / t
        spot = dfactor * forward
        q = 0.0

        def _iv(row: pd.Series) -> float:
            opt = "c" if row["option_type"] == "call" else "p"
            return implied_vol(opt, row["price"], spot, row["strike"], rate, t, q, model="european")

        work["iv"] = work.apply(_iv, axis=1)
        work = work[np.isfinite(work["iv"]) & (work["iv"] > 0.01) & (work["iv"] < 5.0)]
        if work.empty:
            print(f"{exp}: no valid implied vols", file=sys.stderr)
            continue

        if otm_only:
            work = work[
                ((work["option_type"] == "call") & (work["strike"] >= forward)) |
                ((work["option_type"] == "put") & (work["strike"] <= forward))
            ]
        if work.empty:
            print(f"{exp}: no OTM implied vols", file=sys.stderr)
            continue
        if fwd_range is not None:
            low, high = fwd_range
            low_k = forward * low
            high_k = forward * high
            work = work[(work["strike"] >= low_k) & (work["strike"] <= high_k)]
            if work.empty:
                print(f"{exp}: no strikes in fwd-range", file=sys.stderr)
                continue

        strikes = work["strike"].values.astype(float)
        iv = work["iv"].values.astype(float)

        params = fit_sabr(forward, strikes, t, iv, beta=beta, fit_beta=fit_beta)
        results.append(
            {
                "expiry": exp,
                "forward": forward,
                "df": dfactor,
                "alpha": params["alpha"],
                "beta": params["beta"],
                "rho": params["rho"],
                "nu": params["nu"],
                "rmse": params["rmse"],
            }
        )
        if print_per_expiry:
            print(
                f"{exp}: alpha={params['alpha']:.6f} beta={params['beta']:.6f} "
                f"rho={params['rho']:.6f} nu={params['nu']:.6f} rmse={params['rmse']:.6f}"
            )

        strikes_sorted = np.array(sorted(work["strike"].unique()), dtype=float)
        iv_fit = np.array([sabr_iv_hagan(forward, k, t, params["alpha"], params["beta"],
                                         params["rho"], params["nu"]) for k in strikes_sorted])
        for strike, ivf in zip(strikes_sorted, iv_fit):
            fit_rows.append(
                {
                    "expiry": exp,
                    "strike": float(strike),
                    "iv_fit": float(ivf),
                    "iv_fit_pct": float(ivf * 100.0),
                }
            )
        plot_data.append((exp, strikes_sorted, iv_fit * 100.0))
        density_data.append((exp, strikes_sorted, forward, dfactor, t, params))

    if results:
        print("\nSABR summary:")
        print(pd.DataFrame(results).to_string(index=False))

    if outfile and fit_rows:
        out_df = pd.DataFrame(fit_rows)
        out_df.to_csv(outfile, index=False)
        print(f"\nWrote fitted vols to {outfile}")

    if plot and plot_data:
        fig, ax = plt.subplots(figsize=(10, 6))
        for exp, strikes, iv_pct in plot_data:
            ax.plot(strikes, iv_pct, label=exp)
        title_date = as_of if as_of else datetime.now().date().isoformat()
        title_sym = symbol if symbol is not None else "data"
        ax.set_title(f"SABR implied vol vs strike by expiry ({title_sym}, {title_date})")
        ax.set_xlabel("strike")
        ax.set_ylabel("implied vol (pct)")
        ax.legend(fontsize="small", ncol=2)
        ax.grid(True, alpha=0.3)
        plt.tight_layout()
        if plot_file:
            fig.savefig(plot_file, dpi=150)
            print(f"\nWrote plot to {plot_file}")
        else:
            plt.show()

    if plot_density and density_data:
        from black_scholes import bs_price
        fig, ax = plt.subplots(figsize=(10, 6))
        single_tenor = len(density_data) == 1
        ref_x = None
        ref_density = None
        for exp, strikes, forward, dfactor, t, params in density_data:
            strikes = np.asarray(strikes, dtype=float)
            if strikes.size < 5:
                continue
            if fwd_range is not None:
                k_min = forward * fwd_range[0]
                k_max = forward * fwd_range[1]
            else:
                k_min = forward * 0.8
                k_max = forward * 1.2
                k_min = max(k_min, float(np.nanmin(strikes)))
                k_max = min(k_max, float(np.nanmax(strikes)))
            if not np.isfinite(k_min) or not np.isfinite(k_max) or k_max <= k_min:
                continue
            grid = np.linspace(k_min, k_max, 200)
            iv_fit = np.array([sabr_iv_hagan(forward, k, t, params["alpha"], params["beta"],
                                             params["rho"], params["nu"]) for k in grid])
            rate = -np.log(dfactor) / t
            q = 0.0
            spot = forward * dfactor
            call_prices = np.array([bs_price("c", spot, kk, rate, t, v, q)
                                    for kk, v in zip(grid, iv_fit)])
            dK = grid[1] - grid[0]
            dC_dK = np.gradient(call_prices, dK)
            d2C_dK2 = np.gradient(dC_dK, dK)
            f_s = np.maximum(d2C_dK2 * np.exp(rate * t), 0.0)
            if plot_density_log_s:
                x = np.log(grid)
                density = f_s * grid
                ax.plot(x, density, label=exp)
                if single_tenor:
                    ref_x = x
                    ref_density = density
            else:
                ax.plot(grid, f_s, label=exp)
        if plot_density_log_s and single_tenor and ref_x is not None and ref_density is not None:
            integral = np.trapezoid(ref_density, ref_x)
            if integral > 0:
                dens = ref_density / integral
                mu = float(np.trapezoid(ref_x * dens, ref_x))
                var = float(np.trapezoid((ref_x - mu) ** 2 * dens, ref_x))
                if var > 0:
                    sd = np.sqrt(var)
                    normal = (1.0 / (sd * np.sqrt(2.0 * np.pi))) * np.exp(-0.5 * ((ref_x - mu) / sd) ** 2)
                    ax.plot(ref_x, normal, linestyle="--", color="black", label="normal ref")
        title_date = as_of if as_of else datetime.now().date().isoformat()
        title_sym = symbol if symbol is not None else "data"
        if plot_density_log_s:
            ax.set_title(f"Implied density vs log stock (SABR, {title_sym}, {title_date})")
            ax.set_xlabel("log(stock)")
        else:
            ax.set_title(f"Implied density vs stock (SABR, {title_sym}, {title_date})")
            ax.set_xlabel("stock")
        ax.set_ylabel("density (arb units)")
        ax.legend(fontsize="small", ncol=2)
        ax.grid(True, alpha=0.3)
        plt.tight_layout()
        if plot_density_file:
            fig.savefig(plot_density_file, dpi=150)
            print(f"\nWrote density plot to {plot_density_file}")
        else:
            plt.show()


if __name__ == "__main__":
    main(sys.argv)