src/option_implied_lab/density.py

from __future__ import annotations

from datetime import datetime

import numpy as np
import pandas as pd

from .types import DensityEstimate


def select_density_slice(
    options: pd.DataFrame,
    spot: float,
    expiry: datetime | None = None,
    moneyness_band: float = 0.2,
    min_open_interest: int = 1,
    min_volume: int = 0,
) -> tuple[pd.Series, pd.Series, datetime]:
    calls = options[options["option_type"] == "call"].copy()
    if calls.empty:
        raise ValueError("No call options available")

    if expiry is None:
        counts = calls.groupby("expiry").size().sort_values(ascending=False)
        expiry = pd.to_datetime(counts.index[0]).to_pydatetime()

    c = calls[calls["expiry"] == pd.to_datetime(expiry)].copy()
    c = c[
        (c["strike"] >= (1.0 - moneyness_band) * spot)
        & (c["strike"] <= (1.0 + moneyness_band) * spot)
    ]

    if "openInterest" in c.columns:
        c = c[c["openInterest"].fillna(0) >= min_open_interest]
    if "volume" in c.columns:
        c = c[c["volume"].fillna(0) >= min_volume]

    c = c.sort_values("strike")
    if len(c) < 4:
        raise ValueError("Not enough filtered strikes for density estimation")
    return (
        c["strike"].astype(float),
        c["mid"].astype(float),
        pd.to_datetime(expiry).to_pydatetime(),
    )


def estimate_rn_density(
    strikes: pd.Series,
    call_prices: pd.Series,
    expiry: datetime,
    smooth_window: int = 3,
) -> DensityEstimate:
    k = np.asarray(strikes, dtype=float)
    c = np.asarray(call_prices, dtype=float)
    if len(k) < 4:
        raise ValueError("Need at least 4 strikes to estimate density")
    if np.any(np.diff(k) <= 0):
        raise ValueError("Strikes must be strictly increasing")

    if smooth_window > 1:
        c = (
            pd.Series(c)
            .rolling(window=smooth_window, min_periods=1, center=True)
            .mean()
            .to_numpy()
        )

    d2 = np.gradient(np.gradient(c, k), k)
    rho = np.clip(d2, 0.0, None)
    mass = np.trapezoid(rho, k)
    if mass <= 0:
        raise ValueError("Estimated density has non-positive mass")
    rho = rho / mass

    return DensityEstimate(
        strikes=pd.Series(k),
        density=pd.Series(rho),
        expiry=expiry,
    )