solver/models/

heston.rs

use super::traits::{ControlOutput, Gradients, Model};

use market_model::process::HestonProcess;

// 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
    }

    /// Constructs a `HestonProcess` matching this model's variance parameters and
    /// correlation. Uses `price_scale` as the initial spot and `v_theta` as the initial
    /// variance so that the process starts in its stationary regime.
    pub fn build_process(&self) -> HestonProcess {
        HestonProcess::new(
            0.0,
            self.v_kappa,
            self.v_theta,
            self.v_xi,
            self.rho,
            self.price_scale,
            self.v_theta,
        )
    }

    pub fn get_spreads(&self, _q: f64, _v: f64, grads: &Gradients<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 = grads.fwd[0];
        let dv_dq_sell = grads.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)
    }
}

impl Model<2> for Heston {
    type Process = market_model::process::HestonProcess;

    fn process(&self) -> market_model::process::HestonProcess {
        self.build_process()
    }

    fn optimize(&self, state: &[f64; 2], grads: &Gradients<2>) -> ControlOutput<2> {
        let q = state[0];
        let v = state[1];

        // --- Inventory (Dim 0) ---
        let (d_bid, d_ask) = self.get_spreads(q, v, grads);

        let lambda_bid = self.a * (-self.kappa_m * d_bid).exp();
        let lambda_ask = self.a * (-self.kappa_m * d_ask).exp();

        // Shut off the forbidden quoting side at the inventory boundary.
        // At q_max the agent cannot receive more inventory, so no bid is posted.
        // At q_min the agent cannot shed more inventory, so no ask is posted.
        let lambda_bid = if q >= self.q_max { 0.0 } else { lambda_bid };
        let lambda_ask = if q <= self.q_min { 0.0 } else { lambda_ask };

        // Hamiltonian: GLT certainty-equivalent formula (same as AvellanedaStoikov).
        // H_GLT = (lambda_bid + lambda_ask) / (gamma + K)
        let hamiltonian_val = (lambda_bid + lambda_ask) / (self.gamma + self.kappa_m);

        // Risk penalty with stochastic volatility. For multiplicative price dynamics
        // dS = S0*sqrt(v)dW, instantaneous variance is S0^2 * v.
        let risk_penalty = -0.5 * self.gamma * self.price_scale.powi(2) * v * q.powi(2);

        // Cancel solver's inventory jump terms (lambda_bid*(V(q+1)-V(q)) + lambda_ask*(V(q-1)-V(q)))
        // to avoid double-counting: the GLT Hamiltonian already encodes the jump contribution.
        let drift_correction = self.dq * (lambda_bid * grads.fwd[0] - lambda_ask * grads.bwd[0]);

        // --- Variance (Dim 1) ---
        // Physical CIR drift plus the HJB cross-variation correction.
        // From the certainty-equivalent HJB, the joint quadratic variation
        // d[q*dS, dv] = price_scale * rho * xi * v * q * dt produces an effective drift term:
        //   -gamma * price_scale * rho * xi * v * q
        // that modifies how the variance dimension is valued in the Bellman equation.
        // Positive rho with long inventory reduces effective drift (variance decays faster from
        // the HJB perspective); negative rho increases it.
        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 diff_term = sigma2_v / (2.0 * self.dv.powi(2));
        let drift_term_abs = mu_v.abs() / self.dv;

        // Upwinding for Variance Drift
        let mut lambda_v_plus = diff_term;
        let mut lambda_v_minus = diff_term;

        if mu_v > 0.0 {
            lambda_v_plus += drift_term_abs;
        } else {
            lambda_v_minus += drift_term_abs;
        }

        ControlOutput {
            lambda_plus: [lambda_bid, lambda_v_plus],
            lambda_minus: [lambda_ask, lambda_v_minus],
            flow: hamiltonian_val + risk_penalty - drift_correction,
        }
    }

    fn terminal(&self, state: &[f64; 2]) -> f64 {
        match self.terminal_condition {
            TerminalCondition::Zero => 0.0,
            TerminalCondition::LiquidationCost => {
                // At expiry the agent flattens its book by crossing the spread with market
                // orders, paying one base half-spread per unit of remaining inventory.
                let q = state[0];
                let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa_m).ln();
                -q.abs() * base_spread
            }
        }
    }

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

    fn next_step(&self, current_state: &[f64; 2], dt: f64, noise: &[f64; 2]) -> [f64; 2] {
        let q = current_state[0];
        let v = current_state[1];
        let sqrt_dt = dt.sqrt();

        // Inventory is a jump-controlled dimension in this reduced model.
        // Keep it fixed in the forward diffusion step.
        let next_q = q;

        // Variance: risk-adjusted Heston drift for BSDE forward simulation.
        // The term -gamma*price_scale*rho*xi*v*q arises from the cross-variation
        //   d<q*dS, dv>_t = q * price_scale * rho * xi * v * dt
        // in the HJB. Simulating forward paths with this adjusted drift lets the
        // BSDE regression capture the rho-inventory interaction without requiring
        // an explicit driver correction for the variance diffusion dimension.
        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 * q;
        let drift = mu_v * dt;
        let diffusion = self.v_xi * v_pos.sqrt() * sqrt_dt * noise[1];

        let mut next_v = v_pos + drift + diffusion;

        if next_v < 0.0 {
            next_v = -next_v;
        }

        [next_q, next_v]
    }

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

    fn is_integer_dimension(&self, dim: usize) -> bool {
        dim == 0 // q is discrete integer; variance is continuous
    }

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

    fn fill_rate_base(&self, _state: &[f64; 2]) -> f64 {
        self.a
    }

    fn fill_rate_decay(&self) -> f64 {
        self.kappa_m
    }
}

#[cfg(test)]
mod tests {
    use super::super::traits::{Gradients, Model};
    use super::*;
    use market_model::process::StochasticProcess;
    use market_model::process::heston::HestonProcess;
    use ndarray::array;

    fn zero_grads() -> Gradients<2> {
        Gradients {
            fwd: [0.0; 2],
            bwd: [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 grads = zero_grads();
        let (bid, ask) = m.get_spreads(0.0, 0.25, &grads);
        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 base_spread_independent_of_variance_at_zero_gradient() {
        let m = default_model();
        let grads = zero_grads();
        let (bid_low, ask_low) = m.get_spreads(0.0, 0.05, &grads);
        let (bid_high, ask_high) = m.get_spreads(0.0, 0.80, &grads);
        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 grads = zero_grads();
        let ctrl_low = m.optimize(&[2.0, 0.05], &grads);
        let ctrl_high = m.optimize(&[2.0, 0.50], &grads);
        assert!(
            ctrl_high.flow < ctrl_low.flow,
            "Higher variance should decrease flow: low={}, high={}",
            ctrl_low.flow,
            ctrl_high.flow
        );
    }

    #[test]
    fn variance_transport_uses_upwinding() {
        let m = default_model();
        let grads = zero_grads();
        let ctrl = m.optimize(&[0.0, 0.05], &grads);
        assert!(
            ctrl.lambda_plus[1] > ctrl.lambda_minus[1],
            "Upwind: positive drift should give lambda_plus > lambda_minus"
        );

        let ctrl = m.optimize(&[0.0, 0.80], &grads);
        assert!(
            ctrl.lambda_minus[1] > ctrl.lambda_plus[1],
            "Upwind: negative drift should give lambda_minus > lambda_plus"
        );
    }

    #[test]
    fn rho_affects_variance_drift() {
        let grads = zero_grads();
        let state = [2.0, 0.25]; // q=2, v=theta

        let ctrl0 = Heston {
            rho: 0.0,
            ..default_model()
        }
        .optimize(&state, &grads);
        let ctrl_pos = Heston {
            rho: 0.5,
            ..default_model()
        }
        .optimize(&state, &grads);
        let ctrl_neg = Heston {
            rho: -0.5,
            ..default_model()
        }
        .optimize(&state, &grads);

        assert!(
            ctrl_pos.lambda_minus[1] > ctrl0.lambda_minus[1],
            "Positive rho with positive q should increase downward variance transport"
        );
        assert!(
            ctrl_neg.lambda_plus[1] > ctrl0.lambda_plus[1],
            "Negative rho with positive q should increase upward variance transport"
        );
    }

    #[test]
    fn rho_has_no_effect_at_zero_inventory() {
        let grads = zero_grads();
        let state = [0.0, 0.25]; // q=0

        let ctrl0 = Heston {
            rho: 0.0,
            ..default_model()
        }
        .optimize(&state, &grads);
        let ctrl_pos = Heston {
            rho: 0.8,
            ..default_model()
        }
        .optimize(&state, &grads);

        assert!((ctrl0.lambda_plus[1] - ctrl_pos.lambda_plus[1]).abs() < 1e-12);
        assert!((ctrl0.lambda_minus[1] - ctrl_pos.lambda_minus[1]).abs() < 1e-12);
    }

    #[test]
    fn cir_next_step_stays_positive() {
        let m = default_model();
        let state = [0.0, 0.001];
        let next = m.next_step(&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();
        // At v=0.05 << theta=0.25, drift is positive so variance should increase on average
        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 = m.next_step(&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 grads = zero_grads();
        let state = [3.0, 0.25];

        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 c1 = m1.optimize(&state, &grads);
        let c10 = m10.optimize(&state, &grads);

        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 = c10.flow - c1.flow;

        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 grads = zero_grads();
        let state = [2.0, 0.25];
        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 c1 = base.optimize(&state, &grads);
        let c5 = scaled.optimize(&state, &grads);

        let diff1 = c1.lambda_plus[1] - c1.lambda_minus[1];
        let diff5 = c5.lambda_plus[1] - c5.lambda_minus[1];

        assert!(
            diff1 > 0.0 && diff5 > 0.0,
            "Expected positive variance-drift asymmetry from cross term, got diff1={}, diff5={}",
            diff1,
            diff5
        );

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

    #[test]
    fn process_diffusion_covariance_matches_multiplicative_form() {
        let process = HestonProcess::new(0.0, 1.0, 0.25, 0.8, -0.7, 100.0, 0.36);
        let s0 = 100.0;
        let v = 0.36;
        let x = array![s0, v];

        let l = process.diffusion(&x, 0.0);
        let cov = l.dot(&l.t());

        let var_s = cov[[0, 0]];
        let cov_sv = cov[[0, 1]];

        let expected_var_s = s0.powi(2) * v;
        let expected_cov_sv = s0 * process.rho * process.sigma * v;

        assert!(
            (var_s - expected_var_s).abs() < 1e-10,
            "Var(dS)/dt mismatch: observed={}, expected={}",
            var_s,
            expected_var_s
        );
        assert!(
            (cov_sv - expected_cov_sv).abs() < 1e-10,
            "Cov(dS,dv)/dt mismatch: observed={}, expected={}",
            cov_sv,
            expected_cov_sv
        );
    }
}