option_pnl_plot.py

from pydantic import BaseModel
from enum import Enum
from datetime import date, datetime, timedelta
import matplotlib.pyplot as plt
from matplotlib.ticker import FormatStrFormatter


class OptionType(str, Enum):
    Call = 'Call'
    Put = 'Put'


class TradeDirection(str, Enum):
    Buy = 'Buy'
    Sell = 'Sell'


class Option(BaseModel):
    strike: int
    spot: float
    premium: float
    type: OptionType
    direction: TradeDirection

    @property
    def pnl(self) -> float:
        if self.direction == TradeDirection.Buy:
            if self.type == OptionType.Call:
                return max(self.spot - self.strike, 0) - self.premium
            if self.type == OptionType.Put:
                return max(self.strike - self.spot, 0) - self.premium
        
        if self.direction == TradeDirection.Sell:
            if self.type == OptionType.Call:
                return min(self.strike - self.spot, 0) + self.premium
            if self.type == OptionType.Put:
                return min(self.spot - self.strike, 0) + self.premium

    @property
    def summary(self):
        return f'{self.direction} {self.strike:.0f} {self.type} @{self.premium:.1f}'


class Portfolio(BaseModel):
    positions: list[Option]

    @property
    def pnl(self) -> float:
        return sum([trade.pnl for trade in self.positions])


def option_pnl_plot(
    pnl_curves: dict[str, dict[float, float]], 
    strategy_name: str, 
    fill:bool = False, 
    grid:bool = False, 
    show_all:bool = False
):

    GAIN_FILLED_COLOR = "#15bf2c"
    LOSS_FILLED_COLOR = "#e62222"

    underlying_price = pnl_curves.get('total').keys()

    fig = plt.figure(figsize=(6, 4))
    ax = fig.add_subplot()
    fig.subplots_adjust(top=0.9)

    profit = pnl_curves.get('total').values()
    if fill:
        gain = [x > 0 for x in profit]
        loss = [x <= 0 for x in profit]
        ax.fill_between(
                underlying_price, 0, profit, where=loss, color=LOSS_FILLED_COLOR, alpha=0.5
            ),
        ax.fill_between(
            underlying_price, 0, profit, where=gain, color=GAIN_FILLED_COLOR, alpha=0.5
        )

    ls, color, line_width = '-', 'black', 1
    ax.plot(underlying_price, profit, label='total', ls=ls, linewidth=line_width, color=color)


    if show_all:
        for summary, curve in pnl_curves.items():
            if summary != 'total':
                profit = curve.values()
                ls, color, line_width = '--', 'grey', 0.8
                ax.plot(underlying_price, profit, label=summary, ls=ls, linewidth=line_width, color=color)

    # Add labels and title
    plt.xlabel('Spot Price')
    plt.ylabel('Profit / Loss')
    plt.title(f'{strategy_name}', fontweight='bold', fontsize='16')
    plt.axhline(0, color='black', linestyle='--', linewidth=0.8, label='Zero Profit Line')

    if grid:
        plt.grid(linestyle='-.')

    ax.xaxis.set_major_formatter(FormatStrFormatter("%.0f"))
    ax.spines["right"].set_visible(False)
    ax.spines["top"].set_visible(False)

    return fig




current_spot = 17383.99
option = Option(
    strike=17450,
    spot=current_spot,
    premium=70,
    type='Call',
    direction='Buy'
)

positions = [option]
portfolio = Portfolio(positions=positions)


pnl_curves = {'total':{}}
for i in range(-500, 500, 1):
    new_price = current_spot + i

    for trade in portfolio.positions:
        trade.spot = new_price
        pnl_curves.setdefault(trade.summary, {})
        pnl_curves[trade.summary][new_price] = trade.pnl
    
    pnl_curves['total'][new_price] = portfolio.pnl


plot = option_pnl_plot(pnl_curves, 'Buy Call', fill=True, grid=True)
plot.savefig('buy_call.png', dpi=300)