solver/models/

heston.rs

use super::control::{ControlProblem, StateDerivatives};
use super::market_making::MarketMakingControl;
use super::normal_cdf;
use crate::numeric::finite_difference::discretization::{DimensionKind, Transport};
use crate::numeric::finite_difference::pde::PdeProblem;

// Re-export so existing `solver::models::heston::TerminalCondition` paths still compile.
pub use super::TerminalCondition;

/// Heston Stochastic Volatility Model for Market Making
///
/// Extends the Avellaneda-Stoikov model by adding stochastic volatility to the mid-price.
/// The value function depends on both inventory and current variance.
///
/// # Dynamics
/// * $dS_t = S_0 \sqrt{v_t} dW^S_t$
/// * $dv_t = \kappa_v (\theta - v_t) dt + \xi \sqrt{v_t} dW^v_t$
///
/// This reduced HJB model has state [q, v] only (no spot state). To align with
/// a multiplicative market process where price diffusion scales with S_t, we
/// freeze spot at a reference level `price_scale` (typically initial price).
#[derive(Clone, Copy, Debug)]
pub struct Heston {
    /// Risk aversion for inventory holding.
    pub gamma: f64,
    /// Order filling intensity decay ($\kappa$).
    pub kappa_m: f64, // Market kappa (liquidity)
    /// Base order arrival intensity ($A$).
    pub a: f64, // Order arrival parameter
    /// Mean reversion speed for variance ($\kappa_v$).
    pub v_kappa: f64, // Variance mean reversion speed
    /// Long-run mean variance ($\theta$).
    pub v_theta: f64, // Variance long-run mean
    /// Volatility of variance (vol-of-vol, $\xi$).
    pub v_xi: f64, // Variance vol of vol
    /// Correlation between price and variance Brownian motions ($\rho$).
    pub rho: f64,
    /// Reference spot level used to map relative variance into absolute price-risk scale.
    pub price_scale: f64,
    /// Grid step size for the variance dimension.
    pub dv: f64, // Grid step for variance
    /// Grid step for inventory dimension (defaults to 1.0 for integer grid).
    pub dq: f64,
    /// Minimum inventory (lower hard boundary). Defaults to -infinity (no constraint).
    pub q_min: f64,
    /// Maximum inventory (upper hard boundary). Defaults to +infinity (no constraint).
    pub q_max: f64,
    /// Terminal condition for the value function at T. Defaults to Zero.
    pub terminal_condition: TerminalCondition,
}

impl Heston {
    pub fn new(gamma: f64, kappa_m: f64, a: f64) -> Self {
        Self {
            gamma,
            kappa_m,
            a,
            v_kappa: 2.0,
            v_theta: 0.25,
            v_xi: 0.3,
            rho: 0.0,
            price_scale: 1.0,
            dv: 1.0,
            dq: 1.0,
            q_min: f64::NEG_INFINITY,
            q_max: f64::INFINITY,
            terminal_condition: TerminalCondition::Zero,
        }
    }

    pub fn with_terminal_condition(mut self, terminal_condition: TerminalCondition) -> Self {
        self.terminal_condition = terminal_condition;
        self
    }

    pub fn with_inventory_bounds(mut self, q_min: f64, q_max: f64) -> Self {
        self.q_min = q_min;
        self.q_max = q_max;
        self
    }

    pub fn with_variance_params(mut self, v_kappa: f64, v_theta: f64, v_xi: f64) -> Self {
        self.v_kappa = v_kappa;
        self.v_theta = v_theta;
        self.v_xi = v_xi;
        self
    }

    pub fn with_rho(mut self, rho: f64) -> Self {
        self.rho = rho;
        self
    }

    pub fn with_price_scale(mut self, price_scale: f64) -> Self {
        self.price_scale = price_scale.abs().max(1e-8);
        self
    }

    pub fn with_dv(mut self, dv: f64) -> Self {
        self.dv = dv.abs().max(1e-8);
        self
    }

    pub fn with_dq(mut self, dq: f64) -> Self {
        self.dq = dq.abs().max(1e-8);
        self
    }

    /// Sets grid step sizes from a Grid<2>.
    /// Dimension 0 = inventory (dq), Dimension 1 = variance (dv).
    pub fn with_grid_steps(mut self, dx: &[f64; 2]) -> Self {
        self.dq = dx[0].abs().max(1e-8);
        self.dv = dx[1].abs().max(1e-8);
        self
    }

    pub fn get_spreads(&self, derivs: &StateDerivatives<2>) -> (f64, f64) {
        // HJB-derived spread formula (Gueant-Lehalle-Tapia certainty-equivalent):
        //   d_bid* = (1/gamma)*ln(1+gamma/K) - h_fwd
        //   d_ask* = (1/gamma)*ln(1+gamma/K) + h_bwd
        // The base_spread = (1/gamma)*ln(1+gamma/K) is the optimal symmetric spread
        // from the GLT exponential-utility Hamiltonian. All variance and rho effects
        // enter through the value function gradients h_fwd, h_bwd.
        let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa_m).ln();

        let dv_dq_buy = derivs.fwd[0];
        let dv_dq_sell = derivs.bwd[0];

        let mut delta_bid = base_spread - dv_dq_buy;
        let mut delta_ask = base_spread + dv_dq_sell;

        let min_spread = -5.0;
        let max_spread = 10.0;
        delta_bid = delta_bid.max(min_spread).min(max_spread);
        delta_ask = delta_ask.max(min_spread).min(max_spread);

        (delta_bid, delta_ask)
    }

    /// The base arrival intensity used for intensity-to-spread conversion.
    pub fn fill_rate_base(&self, _state: &[f64; 2]) -> f64 {
        self.a
    }

    /// The fill-rate decay parameter used for intensity-to-spread conversion.
    pub fn fill_rate_decay(&self) -> f64 {
        self.kappa_m
    }
}

impl ControlProblem<2> for Heston {
    type Control = MarketMakingControl;

    fn optimize(&self, _t: f64, state: &[f64; 2], derivs: &StateDerivatives<2>) -> Self::Control {
        let q = state[0];
        let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa_m).ln();
        let delta_bid = (base_spread - derivs.fwd[0]).clamp(-5.0, 10.0);
        let delta_ask = (base_spread + derivs.bwd[0]).clamp(-5.0, 10.0);

        let lambda_bid = if q >= self.q_max {
            0.0
        } else {
            self.a * (-self.kappa_m * delta_bid).exp()
        };
        let lambda_ask = if q <= self.q_min {
            0.0
        } else {
            self.a * (-self.kappa_m * delta_ask).exp()
        };

        MarketMakingControl::new(lambda_bid, lambda_ask)
    }

    fn running_reward(&self, _t: f64, _state: &[f64; 2], control: &Self::Control) -> f64 {
        (control.bid_intensity + control.ask_intensity) / (self.gamma + self.kappa_m)
    }

    fn bsde_driver(
        &self,
        _t: f64,
        state: &[f64; 2],
        control: &Self::Control,
        _derivs: &StateDerivatives<2>,
        dt: f64,
    ) -> f64 {
        // Full-value problem: the variance drift/diffusion transport is already
        // simulated by the forward SDE, so the backward driver adds only the
        // running reward plus the local inventory-risk source. Including the
        // transport terms from `generator` would double-count them.
        //
        // The running reward is bounded at `1/dt` per fill side so it cannot
        // count fills the forward Euler step cannot represent: the forward
        // fill probability `lambda * dt` saturates at `1`.
        let q = state[0];
        let v = state[1];
        let local = -0.5 * self.gamma * self.price_scale.powi(2) * v * q.powi(2);
        let rate_cap = 1.0 / dt.max(1e-12);
        let reward = (control.bid_intensity.min(rate_cap) + control.ask_intensity.min(rate_cap))
            / (self.gamma + self.kappa_m);
        reward + local
    }

    fn generator(
        &self,
        _t: f64,
        state: &[f64; 2],
        _control: &Self::Control,
        derivs: &StateDerivatives<2>,
    ) -> f64 {
        // Risk penalty with stochastic volatility plus the continuous variance
        // generator. The variance drift includes the HJB cross-variation
        // correction -gamma * price_scale * rho * xi * v * q.
        let q = state[0];
        let v = state[1];
        let mu_v = self.v_kappa * (self.v_theta - v)
            - self.gamma * self.price_scale * self.rho * self.v_xi * v * q;
        let sigma2_v = self.v_xi.powi(2) * v;

        -0.5 * self.gamma * self.price_scale.powi(2) * v * q.powi(2)
            + mu_v * derivs.grad[1]
            + 0.5 * sigma2_v * derivs.hessian[1]
    }

    fn terminal(&self, state: &[f64; 2]) -> f64 {
        match self.terminal_condition {
            TerminalCondition::Zero => 0.0,
            TerminalCondition::LiquidationCost => {
                let q = state[0];
                let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa_m).ln();
                -q.abs() * base_spread
            }
        }
    }

    fn discount_rate(&self, _state: &[f64; 2]) -> f64 {
        0.0
    }

    fn constant_discount_rate(&self) -> Option<f64> {
        Some(0.0)
    }

    fn next_step(&self, _t: f64, state: &[f64; 2], dt: f64, noise: &[f64; 2]) -> [f64; 2] {
        let mut next = *state;

        let u = normal_cdf(noise[0]);
        let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa_m).ln();
        let lambda_bid = self.a * (-self.kappa_m * base_spread).exp();
        let lambda_ask = self.a * (-self.kappa_m * base_spread).exp();
        let p_bid = (lambda_bid * dt).clamp(0.0, 1.0);
        let p_ask = (lambda_ask * dt).clamp(0.0, 1.0);
        if u < p_bid {
            next[0] += 1.0;
        } else if u > 1.0 - p_ask {
            next[0] -= 1.0;
        }

        let v = state[1];
        let v_pos = v.max(1e-5);
        let mu_v = self.v_kappa * (self.v_theta - v_pos)
            - self.gamma * self.price_scale * self.rho * self.v_xi * v_pos * next[0];
        let drift = mu_v * dt;
        let diffusion = self.v_xi * v_pos.sqrt() * dt.sqrt() * noise[1];
        let mut next_v = v_pos + drift + diffusion;
        if next_v < 0.0 {
            next_v = -next_v;
        }
        next[1] = next_v;
        next
    }

    fn next_step_controlled(
        &self,
        _t: f64,
        state: &[f64; 2],
        control: &Self::Control,
        dt: f64,
        noise: &[f64; 2],
    ) -> [f64; 2] {
        let mut next = *state;

        // The optimal control already carries the fill intensities, so the
        // Bernoulli fill events use them instead of a frozen proxy quote.
        let u = normal_cdf(noise[0]);
        let p_bid = (control.bid_intensity * dt).clamp(0.0, 1.0);
        let p_ask = (control.ask_intensity * dt).clamp(0.0, 1.0);
        if u < p_bid {
            next[0] += 1.0;
        } else if u > 1.0 - p_ask {
            next[0] -= 1.0;
        }

        let v = state[1];
        let v_pos = v.max(1e-5);
        let mu_v = self.v_kappa * (self.v_theta - v_pos)
            - self.gamma * self.price_scale * self.rho * self.v_xi * v_pos * next[0];
        let drift = mu_v * dt;
        let diffusion = self.v_xi * v_pos.sqrt() * dt.sqrt() * noise[1];
        let mut next_v = v_pos + drift + diffusion;
        if next_v < 0.0 {
            next_v = -next_v;
        }
        next[1] = next_v;
        next
    }

    fn is_diffusion_dimension(&self, dim: usize) -> bool {
        dim == 1
    }

    fn gradient_step(&self, dim: usize) -> f64 {
        if dim == 1 {
            self.dv.abs().max(1e-8)
        } else {
            1.0
        }
    }
}

impl PdeProblem<2> for Heston {
    fn dimension_kind(&self, dim: usize) -> DimensionKind {
        match dim {
            0 => DimensionKind::DiscreteJump,
            _ => DimensionKind::Diffusion,
        }
    }

    fn transport(
        &self,
        _t: f64,
        state: &[f64; 2],
        control: &Self::Control,
        derivs: &StateDerivatives<2>,
    ) -> Transport<2> {
        let q = state[0];
        let v = state[1].max(1e-5);
        let mu_v = self.v_kappa * (self.v_theta - v)
            - self.gamma * self.price_scale * self.rho * self.v_xi * v * q;
        let sigma2_v = self.v_xi.powi(2) * v;
        let dv = self.dv.abs().max(1e-12);

        let diff_term = sigma2_v / (2.0 * dv * dv);
        let drift_abs = mu_v.abs() / dv;
        let (plus_v, minus_v) = if mu_v >= 0.0 {
            (diff_term + drift_abs, diff_term)
        } else {
            (diff_term, diff_term + drift_abs)
        };

        let hamiltonian =
            (control.bid_intensity + control.ask_intensity) / (self.gamma + self.kappa_m);
        let jump_transport = self.dq
            * (control.bid_intensity * derivs.fwd[0] - control.ask_intensity * derivs.bwd[0]);
        let local = -0.5 * self.gamma * self.price_scale.powi(2) * v * q.powi(2);

        Transport::new(
            [control.bid_intensity, plus_v],
            [control.ask_intensity, minus_v],
            hamiltonian + local - jump_transport,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn zero_derivs() -> StateDerivatives<2> {
        StateDerivatives::new([0.0; 2], [0.0; 2])
    }

    fn default_model() -> Heston {
        Heston {
            gamma: 0.5,
            kappa_m: 1.5,
            a: 10.0,
            v_kappa: 2.0,
            v_theta: 0.25,
            v_xi: 0.3,
            rho: 0.0,
            price_scale: 1.0,
            dv: 0.02,
            dq: 1.0,
            q_min: f64::NEG_INFINITY,
            q_max: f64::INFINITY,
            terminal_condition: TerminalCondition::Zero,
        }
    }

    #[test]
    fn zero_gradient_gives_base_spread() {
        let m = default_model();
        let derivs = zero_derivs();
        let (bid, ask) = m.get_spreads(&derivs);
        let base = (1.0 / m.gamma) * (1.0 + m.gamma / m.kappa_m).ln();
        assert!((bid - base).abs() < 1e-12);
        assert!((ask - base).abs() < 1e-12);
    }

    #[test]
    fn bsde_driver_excludes_variance_transport() {
        let m = default_model();
        let state = [2.0, 0.25];
        let control = ControlProblem::optimize(&m, 0.0, &state, &zero_derivs());

        // A non-zero variance gradient isolates the transport contribution.
        let mut derivs = zero_derivs();
        derivs.grad[1] = 1.0;
        derivs.hessian[1] = -2.0;

        let dt = 0.01;
        let q = state[0];
        let v = state[1];
        let local = -0.5 * m.gamma * m.price_scale.powi(2) * v * q.powi(2);
        let rate_cap = 1.0 / dt;
        let reward = (control.bid_intensity.min(rate_cap) + control.ask_intensity.min(rate_cap))
            / (m.gamma + m.kappa_m);
        let expected = reward + local;

        assert_eq!(
            ControlProblem::bsde_driver(&m, 0.0, &state, &control, &derivs, dt),
            expected
        );
        assert!(
            ControlProblem::bsde_driver(&m, 0.0, &state, &control, &derivs, dt)
                != ControlProblem::driver(&m, 0.0, &state, &control, &derivs),
            "full driver adds variance transport, bsde driver must not"
        );
    }

    #[test]
    fn next_step_controlled_uses_optimal_intensities() {
        let m = default_model();
        let state = [1.0, 0.25];

        let ctrl = MarketMakingControl::new(0.0, 0.0);
        let next = ControlProblem::next_step_controlled(&m, 0.0, &state, &ctrl, 0.01, &[0.0, 0.0]);
        assert_eq!(next[0], 1.0);

        let ctrl = MarketMakingControl::new(1000.0, 0.0);
        let next = ControlProblem::next_step_controlled(&m, 0.0, &state, &ctrl, 0.01, &[0.0, 0.0]);
        assert_eq!(next[0], 2.0);
    }

    #[test]
    fn base_spread_independent_of_variance_at_zero_gradient() {
        let m = default_model();
        let derivs = zero_derivs();
        let (bid_low, ask_low) = m.get_spreads(&derivs);
        let (bid_high, ask_high) = m.get_spreads(&derivs);
        assert!((bid_low - bid_high).abs() < 1e-12);
        assert!((ask_low - ask_high).abs() < 1e-12);
    }

    #[test]
    fn risk_penalty_scales_with_variance() {
        let m = default_model();
        let derivs = zero_derivs();
        let gen_low = ControlProblem::generator(
            &m,
            0.0,
            &[2.0, 0.05],
            &MarketMakingControl::new(0.0, 0.0),
            &derivs,
        );
        let gen_high = ControlProblem::generator(
            &m,
            0.0,
            &[2.0, 0.50],
            &MarketMakingControl::new(0.0, 0.0),
            &derivs,
        );
        assert!(
            gen_high < gen_low,
            "Higher variance should decrease the generator: low={}, high={}",
            gen_low,
            gen_high
        );
    }

    #[test]
    fn variance_transport_uses_upwinding() {
        let m = default_model();
        let mut derivs = zero_derivs();
        derivs.grad[1] = 1.0;
        let control = MarketMakingControl::new(0.0, 0.0);
        let g_low = ControlProblem::generator(&m, 0.0, &[0.0, 0.05], &control, &derivs);
        let g_high = ControlProblem::generator(&m, 0.0, &[0.0, 0.80], &control, &derivs);
        assert!(
            g_low > g_high,
            "Positive variance drift at low v should raise the generator more than at high v"
        );
    }

    #[test]
    fn rho_affects_variance_drift() {
        let mut derivs = zero_derivs();
        derivs.grad[1] = 1.0;
        let state = [2.0, 0.25];
        let control = MarketMakingControl::new(0.0, 0.0);

        let g0 = ControlProblem::generator(
            &Heston {
                rho: 0.0,
                ..default_model()
            },
            0.0,
            &state,
            &control,
            &derivs,
        );
        let g_pos = ControlProblem::generator(
            &Heston {
                rho: 0.5,
                ..default_model()
            },
            0.0,
            &state,
            &control,
            &derivs,
        );
        let g_neg = ControlProblem::generator(
            &Heston {
                rho: -0.5,
                ..default_model()
            },
            0.0,
            &state,
            &control,
            &derivs,
        );

        assert!(
            g_neg > g0 && g0 > g_pos,
            "Positive rho lowers the effective variance drift; negative rho raises it"
        );
    }

    #[test]
    fn rho_has_no_effect_at_zero_inventory() {
        let mut derivs = zero_derivs();
        derivs.grad[1] = 1.0;
        let state = [0.0, 0.25];
        let control = MarketMakingControl::new(0.0, 0.0);

        let g0 = ControlProblem::generator(
            &Heston {
                rho: 0.0,
                ..default_model()
            },
            0.0,
            &state,
            &control,
            &derivs,
        );
        let g_pos = ControlProblem::generator(
            &Heston {
                rho: 0.8,
                ..default_model()
            },
            0.0,
            &state,
            &control,
            &derivs,
        );

        assert!((g0 - g_pos).abs() < 1e-12);
    }

    #[test]
    fn cir_next_step_stays_positive() {
        let m = default_model();
        let state = [0.0, 0.001];
        let next = ControlProblem::next_step(&m, 0.0, &state, 0.01, &[0.0, -5.0]);
        assert!(
            next[1] > 0.0,
            "CIR reflection should keep variance positive: {}",
            next[1]
        );
    }

    #[test]
    fn next_step_mean_reverts_variance() {
        let m = default_model();
        let state = [0.0, 0.05];
        let mut sum = 0.0;
        for i in 0..1000 {
            let noise = [0.0, (i as f64 * 0.01).sin()];
            let next = ControlProblem::next_step(&m, 0.0, &state, 0.01, &noise);
            sum += next[1] - state[1];
        }
        assert!(
            sum / 1000.0 > 0.0,
            "Variance should drift upward when below theta"
        );
    }

    #[test]
    fn builder_methods_work() {
        let m = Heston::new(0.5, 1.5, 10.0)
            .with_variance_params(3.0, 0.1, 0.5)
            .with_rho(-0.3)
            .with_dv(0.05)
            .with_dq(2.0);
        assert_eq!(m.v_kappa, 3.0);
        assert_eq!(m.rho, -0.3);
        assert_eq!(m.dv, 0.05);
    }

    #[test]
    fn with_grid_steps_sets_both() {
        let m = Heston::new(0.5, 1.5, 10.0).with_grid_steps(&[2.0, 0.05]);
        assert_eq!(m.dq, 2.0);
        assert_eq!(m.dv, 0.05);
    }

    #[test]
    fn risk_penalty_scales_with_price_scale_squared() {
        let derivs = zero_derivs();
        let state = [3.0, 0.25];
        let control = MarketMakingControl::new(0.0, 0.0);

        let m1 = Heston {
            rho: 0.0,
            price_scale: 1.0,
            ..default_model()
        };
        let m10 = Heston {
            rho: 0.0,
            price_scale: 10.0,
            ..default_model()
        };

        let g1 = ControlProblem::generator(&m1, 0.0, &state, &control, &derivs);
        let g10 = ControlProblem::generator(&m10, 0.0, &state, &control, &derivs);

        let penalty1 = -0.5 * m1.gamma * m1.price_scale.powi(2) * state[1] * state[0].powi(2);
        let penalty10 = -0.5 * m10.gamma * m10.price_scale.powi(2) * state[1] * state[0].powi(2);
        let expected_delta = penalty10 - penalty1;
        let observed_delta = g10 - g1;

        assert!(
            (observed_delta - expected_delta).abs() < 1e-10,
            "Risk penalty scaling mismatch: observed_delta={}, expected_delta={}",
            observed_delta,
            expected_delta
        );
    }

    #[test]
    fn cross_term_scales_linearly_with_price_scale() {
        let mut derivs = zero_derivs();
        derivs.grad[1] = 1.0;
        let state = [2.0, 0.25];
        let control = MarketMakingControl::new(0.0, 0.0);
        let base = Heston {
            rho: -0.7,
            v_xi: 0.8,
            dv: 0.05,
            price_scale: 1.0,
            ..default_model()
        };
        let scaled = Heston {
            price_scale: 5.0,
            ..base
        };

        let g1 = ControlProblem::generator(&base, 0.0, &state, &control, &derivs);
        let g5 = ControlProblem::generator(&scaled, 0.0, &state, &control, &derivs);

        let diff1 =
            g1 - (-0.5 * base.gamma * base.price_scale.powi(2) * state[1] * state[0].powi(2));
        let diff5 =
            g5 - (-0.5 * scaled.gamma * scaled.price_scale.powi(2) * state[1] * state[0].powi(2));

        let ratio = diff5 / diff1;
        assert!(
            (ratio - 5.0).abs() < 1e-10,
            "Cross-term scaling should be linear in price_scale: ratio={}",
            ratio
        );
    }
}