visualize.py

import numpy as np
import matplotlib.pyplot as plt
from pricing import black_scholes
from greeks import compute_greeks

def plot_option_payoff(S, K, premium, option_type="call"):
    spot_range = np.linspace(K*0.5, K*1.5, 200)
    if option_type == "call":
        payoff = np.maximum(spot_range - K, 0) - premium
    else:
        payoff = np.maximum(K - spot_range, 0) - premium
    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(spot_range, payoff, color="#1D9E75", linewidth=2.5, label="P&L at expiry")
    ax.axhline(0, color="black", linewidth=0.8, linestyle="--")
    ax.axvline(K, color="#E24B4A", linewidth=1.5, linestyle=":", label=f"Strike = {K}")
    ax.fill_between(spot_range, payoff, 0, where=payoff>0, alpha=0.2, color="#1D9E75")
    ax.fill_between(spot_range, payoff, 0, where=payoff<0, alpha=0.2, color="#E24B4A")
    ax.set_title(f"{option_type.capitalize()} Option Payoff Diagram")
    ax.set_xlabel("Stock Price at Expiry")
    ax.set_ylabel("Profit / Loss")
    ax.legend()
    ax.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig("../results/payoff_diagram.png", dpi=150, bbox_inches="tight")
    plt.show()

def plot_greeks_surface(S, K, r, sigma, option_type="call"):
    T_range = np.linspace(0.05, 1.0, 50)
    S_range = np.linspace(S*0.7, S*1.3, 50)
    TT, SS = np.meshgrid(T_range, S_range)
    delta_surface = np.zeros_like(TT)
    for i in range(TT.shape[0]):
        for j in range(TT.shape[1]):
            g = compute_greeks(SS[i,j], K, TT[i,j], r, sigma, option_type)
            delta_surface[i,j] = g["Delta"]
    fig = plt.figure(figsize=(12, 7))
    ax = fig.add_subplot(111, projection="3d")
    ax.plot_surface(TT, SS, delta_surface, cmap="viridis", alpha=0.85)
    ax.set_xlabel("Time to Expiry (years)")
    ax.set_ylabel("Stock Price")
    ax.set_zlabel("Delta")
    ax.set_title(f"Delta Surface — {option_type.capitalize()} Option")
    plt.tight_layout()
    plt.savefig("../results/delta_surface.png", dpi=150, bbox_inches="tight")
    plt.show()
    print("Greeks surface saved.")