better plots

Author: lsbardel

.vscode/settings.json

@@ -0,0 +1,3 @@
+{
+    "python.defaultInterpreterPath": ".venv/bin/python"
+}

notebooks/applications/volatility_surface.md

@@ -6,14 +6,14 @@ jupytext:
     format_version: 0.13
     jupytext_version: 1.16.6
 kernelspec:
-  display_name: Python 3 (ipykernel)
+  display_name: .venv
   language: python
   name: python3
 ---
 
 # Volatility Surface
 
-In this notebook we illustrate the use of the Volatility Surface tool in the library. We use [deribit](https://docs.deribit.com/) options on BTCUSD as example.
+In this notebook we illustrate the use of the Volatility Surface tool in the library. We use [deribit](https://docs.deribit.com/) options on ETHUSD as example.
 
 First thing, fetch the data
 
@@ -28,7 +28,7 @@ Once we have loaded the data, we create the surface and display the term-structu
 
 ```{code-cell} ipython3
 vs = loader.surface()
-vs.maturities = vs.maturities[1:]
+vs.maturities = vs.maturities
 vs.term_structure()
 ```
 
@@ -58,7 +58,22 @@ df
 The plot function is enabled only if [plotly](https://plotly.com/python/) is installed
 
 ```{code-cell} ipython3
-vs.plot().update_layout(height=500, title="BTC Volatility Surface")
+from plotly.subplots import make_subplots
+
+# consider 6 expiries
+vs6 = vs.trim(6)
+
+titles = []
+for row in range(2):
+    for col in range(3):
+        index = row * 3 + col
+        titles.append(f"Expiry {vs6.maturities[index].maturity}")
+fig = make_subplots(rows=2, cols=3, subplot_titles=titles).update_layout(height=600, title="ETH Volatility Surface")
+for row in range(2):
+    for col in range(3):
+        index = row * 3 + col
+        vs6.plot(index=index, fig=fig, showlegend=False, fig_params=dict(row=row+1, col=col+1))
+fig
 ```
 
 The `moneyness_ttm` is defined as
@@ -70,24 +85,26 @@ The `moneyness_ttm` is defined as
 where $T$ is the time-to-maturity.
 
 ```{code-cell} ipython3
-vs.plot3d().update_layout(height=800, title="BTC Volatility Surface", scene_camera=dict(eye=dict(x=1, y=-2, z=1)))
+vs6.plot3d().update_layout(height=800, title="ETH Volatility Surface", scene_camera=dict(eye=dict(x=1, y=-2, z=1)))
 ```
 
 ## Model Calibration
 
 We can now use the Vol Surface to calibrate the Heston stochastic volatility model.
 
 ```{code-cell} ipython3
-from quantflow.options.calibration import HestonCalibration, OptionPricer
-from quantflow.sp.heston import Heston
+from quantflow.options.calibration import HestonJCalibration, OptionPricer
+from quantflow.utils.distributions import DoubleExponential
+from quantflow.sp.heston import HestonJ
 
-pricer = OptionPricer(Heston.create(vol=0.5))
-cal = HestonCalibration(pricer=pricer, vol_surface=vs, moneyness_weight=-0)
+model = HestonJ.create(DoubleExponential, vol=0.8, sigma=1.5, kappa=0.5, rho=0.1, jump_intensity=50, jump_fraction=0.3)
+pricer = OptionPricer(model=model)
+cal = HestonJCalibration(pricer=pricer, vol_surface=vs6, moneyness_weight=-0)
 len(cal.options)
 ```
 
 ```{code-cell} ipython3
-cal.model
+cal.model.model_dump()
 ```
 
 ```{code-cell} ipython3
@@ -99,7 +116,7 @@ pricer.model
 ```
 
 ```{code-cell} ipython3
-cal.plot(index=6, max_moneyness_ttm=1)
+cal.plot(index=5, max_moneyness_ttm=1)
 ```
 
 ## 

quantflow/options/surface.py

@@ -2,10 +2,10 @@
 
 import enum
 import warnings
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from datetime import datetime, timedelta
 from decimal import Decimal
-from typing import Any, Generic, Iterator, NamedTuple, Protocol, TypeVar
+from typing import Any, Generic, Iterator, NamedTuple, Protocol, Self, TypeVar
 
 import numpy as np
 import pandas as pd
@@ -430,6 +430,10 @@ def term_structure(self, frequency: float = 0) -> pd.DataFrame:
             cross.info_dict(self.ref_date, self.spot) for cross in self.maturities
         )
 
+    def trim(self, num_maturities: int) -> Self:
+        """Create a new volatility surface with the last `num_maturities` maturities"""
+        return replace(self, maturities=self.maturities[-num_maturities:])
+
     def option_prices(
         self,
         *,

quantflow/utils/plot.py

@@ -95,34 +95,48 @@ def plot_characteristic(m: Marginal1D, n: int | None = None, **kwargs: Any) -> A
 
 
 def plot_vol_surface(
-    data: Any,
+    data: pd.DataFrame,
     *,
     model: pd.DataFrame | None = None,
-    template: str = PLOTLY_THEME,
     marker_size: int = 10,
+    x_series: str = "moneyness_ttm",
     series: str = "implied_vol",
+    color_series: str = "side",
+    fig: Any | None = None,
+    fig_params: dict | None = None,
     **kwargs: Any
 ) -> Any:
     check_plotly()
+    # Define a color map for the categorical values
+    color_map = {"bid": "blue", "ask": "red"}
+    colors = data[color_series].map(color_map)
+    fig_params = fig_params or {}
+    fig_: go.Figure = fig or go.Figure()
     params = dict(
-        x="moneyness_ttm",
-        y=series,
-        color="side",
-        template=template,
+        mode="markers",
+        marker=dict(color=colors),
+        **kwargs,
+    )
+    fig_.add_trace(
+        go.Scatter(
+            x=data[x_series],
+            y=data[series],
+            **params,
+        ),
+        **fig_params,
     )
-    params.update(kwargs)
-    fig = px.scatter(data, **params)
     if model is not None:
-        fig.add_trace(
+        fig_.add_trace(
             go.Scatter(
                 x=model["moneyness_ttm"],
                 y=model[series],
                 name="model",
                 mode="lines",
-            )
+            ),
+            **fig_params,
         )
-    fig.update_traces(marker_size=marker_size)
-    return fig
+    fig_.update_traces(marker_size=marker_size)
+    return fig_
 
 
 def plot_vol_surface_3d(