stochastic-rs-quant 2.5.3

//! # Implied Volatility Surface
//!
//! Constructs an implied volatility surface from market option prices.
//!
//! Given a grid of call (or put) prices $C(K_i, T_j)$ over strikes $K$ and
//! maturities $T$, inverts the Black-Scholes formula to obtain
//! $\sigma_{\mathrm{imp}}(K_i, T_j)$.
//!
//! Uses the [`implied_vol`] crate (Jäckel's rational approximation) for
//! robust, high-precision inversion.
//!
//! Reference: Jäckel (2017), "Let's Be Rational"

use implied_vol::DefaultSpecialFn;
use implied_vol::ImpliedBlackVolatility;
use ndarray::Array2;

/// Error returned by the falliable [`ImpliedVolSurface`] constructors.
#[derive(Debug, Clone)]
pub enum ImpliedSurfaceError {
  /// A quote's maturity has no matching forward in the supplied
  /// `forwards` slice. Add a `(tau, forward)` entry to `forwards` for
  /// every distinct quote maturity.
  MissingForward { tau: f64 },
  /// `flat_ivs.len()` did not equal `N_T * N_K`.
  FlatLengthMismatch {
    got: usize,
    nt: usize,
    nk: usize,
    expected: usize,
  },
}

impl std::fmt::Display for ImpliedSurfaceError {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    match self {
      Self::MissingForward { tau } => write!(
        f,
        "missing forward for tau={tau}; supply (tau, forward) for every distinct quote maturity"
      ),
      Self::FlatLengthMismatch {
        got,
        nt,
        nk,
        expected,
      } => write!(
        f,
        "flat_ivs length {got} must equal N_T * N_K = {nt} * {nk} = {expected}"
      ),
    }
  }
}

impl std::error::Error for ImpliedSurfaceError {}

/// Market data for a single option quote.
#[derive(Clone, Debug)]
pub struct OptionQuote {
  /// Strike price
  pub strike: f64,
  /// Time to expiry in years
  pub tau: f64,
  /// Market price of the option
  pub price: f64,
  /// `true` for call, `false` for put
  pub is_call: bool,
}

/// Implied volatility surface built from market data.
#[derive(Clone, Debug)]
pub struct ImpliedVolSurface {
  /// Strikes (ascending), length N_K
  pub strikes: Vec<f64>,
  /// Maturities in years (ascending), length N_T
  pub maturities: Vec<f64>,
  /// Forward prices for each maturity, length N_T
  pub forwards: Vec<f64>,
  /// Implied volatility grid (N_T, N_K); NaN where inversion failed
  pub ivs: Array2<f64>,
  /// Total implied variance grid: $w(k, T) = \sigma^2 T$, shape (N_T, N_K)
  pub total_variance: Array2<f64>,
  /// Log-forward moneyness grid: $k = \ln(K / F)$, shape (N_T, N_K)
  pub log_moneyness: Array2<f64>,
}

impl ImpliedVolSurface {
  /// Build an implied volatility surface from a grid of option prices.
  ///
  /// # Arguments
  /// * `strikes` - Strike prices (ascending)
  /// * `maturities` - Maturities in years (ascending)
  /// * `forwards` - Forward prices for each maturity
  /// * `prices` - **Undiscounted** option price grid (N_T, N_K)
  /// * `is_call` - Whether prices are call (`true`) or put (`false`)
  #[must_use]
  pub fn from_prices(
    strikes: Vec<f64>,
    maturities: Vec<f64>,
    forwards: Vec<f64>,
    prices: &Array2<f64>,
    is_call: bool,
  ) -> Self {
    let nt = maturities.len();
    let nk = strikes.len();
    assert_eq!(prices.dim(), (nt, nk), "prices shape must be (N_T, N_K)");
    assert_eq!(forwards.len(), nt, "forwards length must match maturities");

    let mut ivs = Array2::<f64>::from_elem((nt, nk), f64::NAN);
    let mut total_variance = Array2::<f64>::from_elem((nt, nk), f64::NAN);
    let mut log_moneyness = Array2::<f64>::zeros((nt, nk));

    for j in 0..nt {
      let f = forwards[j];
      let t = maturities[j];
      for i in 0..nk {
        let k = strikes[i];
        log_moneyness[[j, i]] = (k / f).ln();

        let iv = ImpliedBlackVolatility::builder()
          .option_price(prices[[j, i]])
          .forward(f)
          .strike(k)
          .expiry(t)
          .is_call(is_call)
          .build()
          .and_then(|v| v.calculate::<DefaultSpecialFn>())
          .unwrap_or(f64::NAN);

        if iv.is_finite() && iv > 0.0 {
          ivs[[j, i]] = iv;
          total_variance[[j, i]] = iv * iv * t;
        }
      }
    }

    Self {
      strikes,
      maturities,
      forwards,
      ivs,
      total_variance,
      log_moneyness,
    }
  }

  /// Build a surface directly from a pre-computed implied-vol grid.
  ///
  /// Useful when the IVs come from an upstream source that already inverted
  /// (or never needed to invert) Black-Scholes — for example AI surrogates
  /// such as `stochastic_rs_ai::volatility::HestonNn::predict_surface`,
  /// which output IVs directly in the standard `(N_T, N_K)` layout.
  ///
  /// # Arguments
  /// * `strikes` — strike prices in ascending order, length `N_K`
  /// * `maturities` — expiries in years, length `N_T`
  /// * `forwards` — forward prices for each maturity, length `N_T`
  /// * `ivs` — implied volatility grid of shape `(N_T, N_K)`
  #[must_use]
  pub fn from_iv_grid(
    strikes: Vec<f64>,
    maturities: Vec<f64>,
    forwards: Vec<f64>,
    ivs: Array2<f64>,
  ) -> Self {
    let nt = maturities.len();
    let nk = strikes.len();
    assert_eq!(ivs.dim(), (nt, nk), "ivs shape must be (N_T, N_K)");
    assert_eq!(forwards.len(), nt, "forwards length must match maturities");

    let mut total_variance = Array2::<f64>::from_elem((nt, nk), f64::NAN);
    let mut log_moneyness = Array2::<f64>::zeros((nt, nk));

    for j in 0..nt {
      let f = forwards[j];
      let t = maturities[j];
      for i in 0..nk {
        log_moneyness[[j, i]] = (strikes[i] / f).ln();
        let iv = ivs[[j, i]];
        if iv.is_finite() && iv > 0.0 {
          total_variance[[j, i]] = iv * iv * t;
        }
      }
    }

    Self {
      strikes,
      maturities,
      forwards,
      ivs,
      total_variance,
      log_moneyness,
    }
  }

  /// Build a surface from an AI surrogate's flat output vector.
  ///
  /// Bridges [`stochastic_rs_ai::volatility::StochVolNn::predict_surface`]
  /// (and the specialized `HestonNn` / `RBergomiNn` / `OneFactorNn` wrappers)
  /// to the vol-surface pipeline. The neural network returns a flat
  /// `Vec<f32>` of length `N_T * N_K` in row-major `(maturity, strike)`
  /// order; this constructor reshapes and lifts to `f64`.
  ///
  /// # Arguments
  /// * `strikes` — strike prices in ascending order, length `N_K`
  /// * `maturities` — expiries in years, length `N_T`
  /// * `forwards` — forward prices for each maturity, length `N_T`
  /// * `flat_ivs` — flat row-major IV grid of length `N_T * N_K`
  ///
  /// # Example
  ///
  /// ```ignore
  /// let flat = nn.predict_surface(&params)?;          // Vec<f32>, len N_T*N_K
  /// let surf = ImpliedVolSurface::from_flat_iv_grid(
  ///     strikes, maturities, forwards, &flat,
  /// );
  /// ```
  #[must_use]
  pub fn from_flat_iv_grid(
    strikes: Vec<f64>,
    maturities: Vec<f64>,
    forwards: Vec<f64>,
    flat_ivs: &[f32],
  ) -> Self {
    let nt = maturities.len();
    let nk = strikes.len();
    assert_eq!(
      flat_ivs.len(),
      nt * nk,
      "flat_ivs length must equal N_T * N_K = {} * {} = {}",
      nt,
      nk,
      nt * nk,
    );
    let ivs = Array2::from_shape_vec((nt, nk), flat_ivs.iter().map(|&v| v as f64).collect())
      .expect("shape (N_T, N_K) is consistent with flat_ivs length");
    Self::from_iv_grid(strikes, maturities, forwards, ivs)
  }

  /// Build from a flat list of [`OptionQuote`]s.
  ///
  /// Quotes are sorted and grouped by maturity, then by strike.
  /// Forward prices are required for each unique maturity.
  ///
  /// Panics on missing forwards. Use [`Self::try_from_quotes`] for a
  /// falliable variant returning [`ImpliedSurfaceError`].
  #[must_use]
  pub fn from_quotes(quotes: &[OptionQuote], forwards: &[(f64, f64)]) -> Self {
    Self::try_from_quotes(quotes, forwards)
      .expect("from_quotes: forwards missing — use try_from_quotes for the Result variant")
  }

  /// Falliable variant of [`Self::from_quotes`] returning
  /// [`ImpliedSurfaceError::MissingForward`] when a quote's maturity
  /// has no matching forward in `forwards`.
  pub fn try_from_quotes(
    quotes: &[OptionQuote],
    forwards: &[(f64, f64)],
  ) -> Result<Self, ImpliedSurfaceError> {
    let mut tau_set: Vec<f64> = quotes.iter().map(|q| q.tau).collect();
    tau_set.sort_by(|a, b| a.partial_cmp(b).unwrap());
    tau_set.dedup_by(|a, b| (*a - *b).abs() < 1e-12);

    let mut strike_set: Vec<f64> = quotes.iter().map(|q| q.strike).collect();
    strike_set.sort_by(|a, b| a.partial_cmp(b).unwrap());
    strike_set.dedup_by(|a, b| (*a - *b).abs() < 1e-12);

    let fwd_map: std::collections::HashMap<u64, f64> =
      forwards.iter().map(|&(t, f)| (t.to_bits(), f)).collect();

    let nt = tau_set.len();
    let nk = strike_set.len();
    let mut prices = Array2::<f64>::from_elem((nt, nk), f64::NAN);
    let mut is_call_grid = vec![true; nt * nk];
    let fwd_vec: Vec<f64> = tau_set
      .iter()
      .map(|t| {
        fwd_map
          .get(&t.to_bits())
          .copied()
          .ok_or(ImpliedSurfaceError::MissingForward { tau: *t })
      })
      .collect::<Result<_, _>>()?;

    for q in quotes {
      let j = tau_set
        .iter()
        .position(|t| (t - q.tau).abs() < 1e-12)
        .unwrap();
      let i = strike_set
        .iter()
        .position(|k| (k - q.strike).abs() < 1e-12)
        .unwrap();
      prices[[j, i]] = q.price;
      is_call_grid[j * nk + i] = q.is_call;
    }

    let mut ivs = Array2::<f64>::from_elem((nt, nk), f64::NAN);
    let mut total_variance = Array2::<f64>::from_elem((nt, nk), f64::NAN);
    let mut log_moneyness = Array2::<f64>::zeros((nt, nk));

    for j in 0..nt {
      let f = fwd_vec[j];
      let t = tau_set[j];
      for i in 0..nk {
        let p = prices[[j, i]];
        if p.is_nan() {
          continue;
        }
        let k = strike_set[i];
        log_moneyness[[j, i]] = (k / f).ln();

        let iv = ImpliedBlackVolatility::builder()
          .option_price(p)
          .forward(f)
          .strike(k)
          .expiry(t)
          .is_call(is_call_grid[j * nk + i])
          .build()
          .and_then(|v| v.calculate::<DefaultSpecialFn>())
          .unwrap_or(f64::NAN);

        if iv.is_finite() && iv > 0.0 {
          ivs[[j, i]] = iv;
          total_variance[[j, i]] = iv * iv * t;
        }
      }
    }

    Ok(Self {
      strikes: strike_set,
      maturities: tau_set,
      forwards: fwd_vec,
      ivs,
      total_variance,
      log_moneyness,
    })
  }

  /// Build an implied-vol surface directly from a market-data provider's
  /// option chain. Fetches the chain for `symbol`, converts the quotes to
  /// the `(OptionQuote, forwards)` inputs via
  /// [`OptionChain::to_surface_inputs`](crate::market::provider::OptionChain::to_surface_inputs)
  /// with carry $(r, q)$, then runs [`Self::try_from_quotes`].
  ///
  /// Works against any [`MarketDataProvider`](crate::market::provider::MarketDataProvider) —
  /// the in-memory `MockProvider` for offline tests / examples, or the live
  /// Yahoo connector behind the `yahoo` feature.
  pub fn from_provider<P: crate::market::provider::MarketDataProvider>(
    provider: &P,
    symbol: &str,
    r: f64,
    q: f64,
  ) -> anyhow::Result<Self> {
    let chain = provider.option_chain(symbol)?;
    let (quotes, forwards) = chain.to_surface_inputs(r, q);
    if quotes.is_empty() {
      anyhow::bail!("from_provider: option chain for '{symbol}' yielded no usable quotes");
    }
    Self::try_from_quotes(&quotes, &forwards)
      .map_err(|e| anyhow::anyhow!("from_provider: surface build failed: {e}"))
  }

  /// Extract a single smile slice (implied vols for one maturity).
  #[must_use]
  pub fn smile_slice(&self, maturity_idx: usize) -> SmileSlice {
    assert!(maturity_idx < self.maturities.len());
    let nk = self.strikes.len();
    let mut ks = Vec::with_capacity(nk);
    let mut vols = Vec::with_capacity(nk);
    let mut ws = Vec::with_capacity(nk);
    let f = self.forwards[maturity_idx];
    let t = self.maturities[maturity_idx];

    for i in 0..nk {
      let iv = self.ivs[[maturity_idx, i]];
      if iv.is_finite() {
        ks.push(self.log_moneyness[[maturity_idx, i]]);
        vols.push(iv);
        ws.push(iv * iv * t);
      }
    }

    SmileSlice {
      log_moneyness: ks,
      implied_vols: vols,
      total_variance: ws,
      forward: f,
      tau: t,
    }
  }
}

/// A single maturity smile slice with market-observed data.
#[derive(Clone, Debug)]
pub struct SmileSlice {
  /// Log-forward moneyness $k = \ln(K/F)$
  pub log_moneyness: Vec<f64>,
  /// Implied volatilities
  pub implied_vols: Vec<f64>,
  /// Total implied variance $w = \sigma^2 T$
  pub total_variance: Vec<f64>,
  /// Forward price
  pub forward: f64,
  /// Time to expiry in years
  pub tau: f64,
}

impl SmileSlice {
  /// Fit SVI raw parameters to this smile slice.
  pub fn fit_svi(
    &self,
    initial: Option<super::svi::SviRawParams<f64>>,
  ) -> super::svi::SviRawParams<f64> {
    super::svi::calibrate_svi(&self.log_moneyness, &self.total_variance, initial)
  }

  /// Convert to an SSVI slice using the ATM total variance from the data.
  ///
  /// ATM total variance $\theta$ is interpolated at $k = 0$.
  pub fn to_ssvi_slice(&self) -> super::ssvi::SsviSlice<f64> {
    let theta = self.atm_total_variance();
    super::ssvi::SsviSlice {
      log_moneyness: self.log_moneyness.clone(),
      total_variance: self.total_variance.clone(),
      theta,
    }
  }

  /// Interpolate ATM total variance ($k = 0$) from the data.
  ///
  /// When the smile straddles $k = 0$, returns a linear interpolation between
  /// the two adjacent grid points.
  ///
  /// When the smile is **one-sided** (all strikes ITM or all strikes OTM),
  /// extrapolates linearly from the two innermost data points (closest to zero)
  /// rather than returning a non-ATM endpoint as ATM. The previous behaviour
  /// (returning `total_variance[0]` or `total_variance[n-1]`) silently biased
  /// downstream global fits (e.g. SSVI's $\theta_t$ estimate) — see audit
  /// §1.2.8.
  fn atm_total_variance(&self) -> f64 {
    let n = self.log_moneyness.len();
    if n == 0 {
      return 0.0;
    }
    if n == 1 {
      return self.total_variance[0];
    }

    let idx = self.log_moneyness.partition_point(|&k| k < 0.0);

    if idx == 0 {
      // All strikes are k > 0 (one-sided OTM-call / ITM-put).
      // Extrapolate linearly to k=0 using the two innermost (smallest k) points.
      let k0 = self.log_moneyness[0];
      let k1 = self.log_moneyness[1];
      let w0 = self.total_variance[0];
      let w1 = self.total_variance[1];
      if (k1 - k0).abs() < 1e-14 {
        return w0;
      }
      let slope = (w1 - w0) / (k1 - k0);
      return w0 + slope * (0.0 - k0);
    }
    if idx >= n {
      // All strikes are k < 0 (one-sided ITM-call / OTM-put).
      // Extrapolate linearly to k=0 using the two innermost (largest k) points.
      let k0 = self.log_moneyness[n - 2];
      let k1 = self.log_moneyness[n - 1];
      let w0 = self.total_variance[n - 2];
      let w1 = self.total_variance[n - 1];
      if (k1 - k0).abs() < 1e-14 {
        return w1;
      }
      let slope = (w1 - w0) / (k1 - k0);
      return w1 + slope * (0.0 - k1);
    }

    let k0 = self.log_moneyness[idx - 1];
    let k1 = self.log_moneyness[idx];
    let w0 = self.total_variance[idx - 1];
    let w1 = self.total_variance[idx];

    if (k1 - k0).abs() < 1e-14 {
      return w0;
    }

    let alpha = (0.0 - k0) / (k1 - k0);
    w0 * (1.0 - alpha) + w1 * alpha
  }
}

impl ImpliedVolSurface {
  /// Fit SVI parameters to each maturity slice independently.
  pub fn fit_svi_slices(&self) -> Vec<super::svi::SviRawParams<f64>> {
    let nt = self.maturities.len();
    (0..nt).map(|j| self.smile_slice(j).fit_svi(None)).collect()
  }

  /// Fit SSVI surface to all maturity slices simultaneously.
  ///
  /// First extracts ATM total variance $\theta_t$ per slice, then
  /// calibrates global SSVI parameters $(\rho, \eta, \gamma)$.
  pub fn fit_ssvi(
    &self,
    initial: Option<super::ssvi::SsviParams<f64>>,
  ) -> super::ssvi::SsviSurface<f64> {
    let nt = self.maturities.len();
    let slices: Vec<super::ssvi::SsviSlice<f64>> = (0..nt)
      .map(|j| self.smile_slice(j).to_ssvi_slice())
      .collect();

    let params = super::ssvi::calibrate_ssvi(&slices, initial);

    let thetas: Vec<f64> = slices.iter().map(|s| s.theta).collect();

    super::ssvi::SsviSurface::new(params, thetas, self.maturities.clone())
  }
}

#[cfg(test)]
mod tests {
  use ndarray::array;

  use super::*;

  /// `from_provider` against a `MockProvider` whose option chain carries
  /// undiscounted-Black call prices: the recovered implied vols must match
  /// the σ those prices were generated from. Exercises the full Tier 3
  /// provider → surface wiring offline.
  #[test]
  fn from_provider_recovers_surface_from_mock_chain() {
    use stochastic_rs_distributions::special::norm_cdf;

    use crate::market::provider::ChainQuote;
    use crate::market::provider::MockProvider;
    use crate::market::provider::OptionChain;

    let s = 100.0_f64;
    let r = 0.04_f64;
    let sigma = 0.22_f64;
    let strikes = [90.0_f64, 100.0, 110.0];
    let taus = [0.5_f64, 1.0];

    let mut quotes = Vec::new();
    for &t in &taus {
      let f = s * (r * t).exp();
      for &k in &strikes {
        let d1 = ((f / k).ln() + 0.5 * sigma * sigma * t) / (sigma * t.sqrt());
        let d2 = d1 - sigma * t.sqrt();
        let price = f * norm_cdf(d1) - k * norm_cdf(d2); // undiscounted Black
        quotes.push(ChainQuote {
          strike: k,
          tau: t,
          last: price,
          bid: price,
          ask: price,
          implied_vol: sigma,
          is_call: true,
        });
      }
    }
    let mut mp = MockProvider::new();
    mp.insert_option_chain(OptionChain {
      symbol: "ACME".to_string(),
      spot: s,
      quotes,
    });

    let surface = ImpliedVolSurface::from_provider(&mp, "ACME", r, 0.0).unwrap();
    for j in 0..taus.len() {
      for i in 0..strikes.len() {
        let iv = surface.ivs[[j, i]];
        assert!(
          (iv - sigma).abs() < 1e-4,
          "recovered iv={iv} vs σ={sigma} at T={}, K={}",
          taus[j],
          strikes[i]
        );
      }
    }
  }

  /// `from_provider` propagates a provider error for a missing symbol.
  #[test]
  fn from_provider_errors_on_missing_symbol() {
    use crate::market::provider::MockProvider;

    let mp = MockProvider::new();
    let res = ImpliedVolSurface::from_provider(&mp, "NOPE", 0.05, 0.0);
    assert!(res.is_err());
  }

  #[test]
  fn from_prices_round_trip() {
    use stochastic_rs_distributions::special::norm_cdf;

    let s = 100.0;
    let r = 0.05;
    let sigma = 0.20;

    let strikes = vec![90.0, 95.0, 100.0, 105.0, 110.0];
    let maturities = vec![0.25, 0.50, 1.0];
    let forwards: Vec<f64> = maturities.iter().map(|&t| s * f64::exp(r * t)).collect();

    let mut prices = Array2::<f64>::zeros((maturities.len(), strikes.len()));
    for (j, &t) in maturities.iter().enumerate() {
      let f = forwards[j];
      for (i, &k) in strikes.iter().enumerate() {
        let d1 = ((f / k).ln() + 0.5 * sigma * sigma * t) / (sigma * t.sqrt());
        let d2 = d1 - sigma * t.sqrt();
        // Undiscounted Black price
        prices[[j, i]] = f * norm_cdf(d1) - k * norm_cdf(d2);
      }
    }

    let surface = ImpliedVolSurface::from_prices(
      strikes.clone(),
      maturities.clone(),
      forwards.clone(),
      &prices,
      true,
    );

    for j in 0..maturities.len() {
      for i in 0..strikes.len() {
        let iv = surface.ivs[[j, i]];
        assert!(
          (iv - sigma).abs() < 1e-6,
          "iv={iv} vs sigma={sigma} at T={}, K={}",
          maturities[j],
          strikes[i]
        );
      }
    }
  }

  #[test]
  fn smile_slice_filters_nans() {
    let surface = ImpliedVolSurface {
      strikes: vec![90.0, 100.0, 110.0],
      maturities: vec![0.5],
      forwards: vec![100.0],
      ivs: array![[0.22, f64::NAN, 0.20]],
      total_variance: array![[0.0242, f64::NAN, 0.02]],
      log_moneyness: array![[(90.0_f64 / 100.0).ln(), 0.0, (110.0_f64 / 100.0).ln()]],
    };

    let slice = surface.smile_slice(0);
    assert_eq!(slice.implied_vols.len(), 2);
    assert!((slice.implied_vols[0] - 0.22).abs() < 1e-12);
    assert!((slice.implied_vols[1] - 0.20).abs() < 1e-12);
  }
}