plot_smile.py

"""Plot the implied volatility smile under Heston for several correlations.

Run with:

    python plot_smile.py
"""

from __future__ import annotations

import matplotlib.pyplot as plt
import numpy as np

from black_scholes import implied_vol
from model import HestonMonteCarlo, HestonParameters

CORRELATION_SCENARIOS: dict[str, float] = {
    "rho = -0.90 (strong)": -0.90,
    "rho = -0.70 (standard)": -0.70,
    "rho = -0.50 (weak)": -0.50,
    "rho =  0.00 (none)": 0.00,
}


def smile(
    simulator: HestonMonteCarlo,
    strikes: np.ndarray,
    maturity: float,
    rho: float,
) -> np.ndarray:
    """Return the implied volatility curve for one correlation scenario"""
    prices = simulator.call_prices(strikes, maturity, rho)
    return np.array(
        [
            implied_vol(simulator.params.spot, strike, maturity, simulator.params.rate, price)
            for strike, price in zip(strikes, prices, strict=True)
        ]
    )


def main() -> None:
    params = HestonParameters(
        spot=100.0,
        rate=0.03,
        kappa=2.0,
        theta=0.04,
        sigma_v=0.30,
        v0=0.04,
    )
    maturity = 1.0
    strikes = np.linspace(80.0, 120.0, 40)

    simulator = HestonMonteCarlo(params, n_paths=300_000, n_steps=100, seed=42)

    print(f"Feller condition satisfied: {params.satisfies_feller}")

    plt.figure(figsize=(12, 8))
    for label, rho in CORRELATION_SCENARIOS.items():
        iv = smile(simulator, strikes, maturity, rho)
        plt.plot(strikes, iv, label=label, linewidth=2)

    plt.title("Heston implied volatility smile as a function of rho", fontsize=14)
    plt.xlabel("Strike", fontsize=12)
    plt.ylabel("Implied volatility", fontsize=12)
    plt.legend(loc="best", fontsize=10)
    plt.grid(visible=True, linestyle="--", alpha=0.7)
    plt.tight_layout()
    plt.show()


if __name__ == "__main__":
    main()