solver/models/

heston_hawkes.rs

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

/// Combined Heston-Hawkes Model for Market Making (N=3).
///
/// Merges stochastic volatility (Heston) with self-exciting order flow (Hawkes)
/// into a single three-dimensional model, demonstrating seamless N=3 extensibility.
///
/// # State Space
/// `[q, v, lambda]` where:
/// - `q` (dim 0): Inventory level (discrete, jump-controlled)
/// - `v` (dim 1): Stochastic variance following CIR dynamics
/// - `lambda` (dim 2): Hawkes order arrival intensity (mean-reverting with self-excitation)
///
/// # Dynamics
/// * Inventory: $dq = +1$ (buy fill) or $-1$ (sell fill)
/// * Variance: $dv_t = \kappa_v (\theta - v_t) dt + \xi \sqrt{v_t} dW^v_t$
/// * Intensity: $d\lambda_t = \beta (\mu_\lambda - \lambda_t) dt + \alpha \lambda_t dt$
///
/// The optimal spreads incorporate both volatility risk premia (from Heston)
/// and state-dependent arrival rates (from Hawkes).
#[derive(Clone, Debug)]
pub struct HestonHawkes {
    // -- Market-making parameters --
    /// Risk aversion for inventory holding.
    pub gamma: f64,
    /// Order filling intensity decay.
    pub kappa: f64,

    // -- Heston parameters --
    /// Mean reversion speed for variance.
    pub v_kappa: f64,
    /// Long-run mean variance.
    pub v_theta: f64,
    /// Volatility of variance.
    pub v_xi: f64,
    /// Correlation between price and variance Brownian motions ($\rho$).
    /// Enters the HJB through the cross-variation drift: $-\gamma \rho \xi v q$.
    pub rho: f64,
    /// Grid step size for variance dimension.
    pub dv: f64,

    // -- Hawkes parameters --
    /// Self-excitation jump size.
    pub alpha: f64,
    /// Mean-reversion speed of intensity.
    pub beta: f64,
    /// Baseline intensity level.
    pub mu_lambda: f64,
    /// Grid step for intensity dimension.
    pub lambda_step: f64,
    /// Grid step for inventory dimension.
    pub dq: f64,
}

impl HestonHawkes {
    pub fn new(gamma: f64, kappa: f64) -> Self {
        Self {
            gamma,
            kappa,
            v_kappa: 2.0,
            v_theta: 0.25,
            v_xi: 0.3,
            rho: 0.0,
            dv: 1.0,
            alpha: 0.5,
            beta: 2.0,
            mu_lambda: 1.0,
            lambda_step: 1.0,
            dq: 1.0,
        }
    }

    pub fn with_heston_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_hawkes_params(mut self, alpha: f64, beta: f64, mu_lambda: f64) -> Self {
        self.alpha = alpha;
        self.beta = beta;
        self.mu_lambda = mu_lambda;
        self
    }

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

    /// Compute optimal bid/ask spreads given the current state and value gradients.
    ///
    /// Combines the Heston volatility risk premium with the Hawkes intensity dependence.
    pub fn get_spreads(&self, _q: f64, v: f64, grads: &Gradients<3>) -> (f64, f64) {
        // Pure HJB-derived formula: all variance risk is encoded in V(q, v, lambda).
        // Removing ad-hoc premia ensures the degenerate limit (kappa_v -> inf, alpha -> 0)
        // collapses cleanly to the base AS model.
        let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa).ln();

        let _v_pos = v.max(0.0);

        // Value gradient w.r.t. inventory
        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;

        delta_bid = delta_bid.clamp(-5.0, 10.0);
        delta_ask = delta_ask.clamp(-5.0, 10.0);

        (delta_bid, delta_ask)
    }
}

impl Model<3> for HestonHawkes {
    type Process = ();

    fn process(&self) {}

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

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

        // Arrival rate is modulated by Hawkes intensity lambda
        // Cap fill-rate multiplier to prevent blow-up at extreme inventories
        let max_fill_mult = 20.0_f64;
        let lambda_bid = lambda * (-self.kappa * d_bid).exp().min(max_fill_mult);
        let lambda_ask = lambda * (-self.kappa * d_ask).exp().min(max_fill_mult);

        let hamiltonian_inv = (lambda_bid + lambda_ask) / (self.gamma + self.kappa);
        let risk_penalty = -0.5 * self.gamma * v * q.powi(2);

        // Cancel solver's q-dimension transport by dq factor to match matrix contribution
        let drift_correction_q = self.dq * (lambda_bid * grads.fwd[0] - lambda_ask * grads.bwd[0]);

        // -- Dimension 1: Heston variance (CIR diffusion) --
        // Physical CIR drift plus the HJB cross-variation correction.
        // d[q*dS, dv] = rho * xi * v * q * dt produces -gamma * rho * xi * v * q.
        let v_pos = v.max(1e-5);
        let mu_v =
            self.v_kappa * (self.v_theta - v_pos) - self.gamma * self.rho * self.v_xi * v_pos * q;
        let sigma2_v = self.v_xi.powi(2) * v_pos;

        let diff_term = sigma2_v / (2.0 * self.dv.powi(2));
        let drift_term_abs = mu_v.abs() / self.dv;

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

        // -- Dimension 2: Hawkes intensity (properly in matrix for implicit stability) --
        let net_drift_lambda = self.beta * (self.mu_lambda - lambda) + self.alpha * lambda;
        let drift_abs_lambda = net_drift_lambda.abs() / self.lambda_step;

        let (lambda_lam_plus, lambda_lam_minus) = if net_drift_lambda >= 0.0 {
            (drift_abs_lambda, 0.0)
        } else {
            (0.0, drift_abs_lambda)
        };

        ControlOutput {
            lambda_plus: [lambda_bid, lambda_v_plus, lambda_lam_plus],
            lambda_minus: [lambda_ask, lambda_v_minus, lambda_lam_minus],
            flow: hamiltonian_inv + risk_penalty - drift_correction_q,
        }
    }

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

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

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

        // Dim 0: Inventory held fixed in forward diffusion
        let next_q = q;

        // Dim 1: Variance — CIR with reflection
        let v_pos = v.max(1e-5);
        let drift_v = self.v_kappa * (self.v_theta - v_pos) * dt;
        let diff_v = self.v_xi * v_pos.sqrt() * sqrt_dt * noise[1];
        let mut next_v = v_pos + drift_v + diff_v;
        if next_v < 0.0 {
            next_v = -next_v;
        }

        // Dim 2: Hawkes intensity — first-order ODE (no diffusion)
        let net_drift = self.beta * (self.mu_lambda - lambda) + self.alpha * lambda;
        let next_lambda = (lambda + net_drift * dt).max(0.0);

        [next_q, next_v, next_lambda]
    }

    fn is_diffusion_dimension(&self, dim: usize) -> bool {
        // dim 1 (variance) is a true diffusion
        // dim 2 (Hawkes) is treated as deterministic drift
        dim == 1
    }

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

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

    /// The base arrival intensity is state-dependent: returns the current
    /// intensity `state[2]` (clamped above 0) for intensity-to-spread conversion.
    fn fill_rate_base(&self, state: &[f64; 3]) -> f64 {
        state[2].max(1e-10)
    }

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

#[cfg(test)]
mod tests {
    use super::super::traits::{Gradients, Model};
    use super::*;

    fn zero_grads() -> Gradients<3> {
        Gradients {
            fwd: [0.0; 3],
            bwd: [0.0; 3],
        }
    }

    fn default_model() -> HestonHawkes {
        HestonHawkes::new(0.5, 1.5)
            .with_heston_params(2.0, 0.25, 0.3)
            .with_rho(0.0)
            .with_hawkes_params(0.5, 2.0, 1.0)
            .with_grid_steps(&[1.0, 0.02, 0.1])
    }

    #[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).ln();
        assert!((bid - base).abs() < 1e-12);
        assert!((ask - base).abs() < 1e-12);
    }

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

        let ctrl0 = HestonHawkes::new(0.5, 1.5)
            .with_heston_params(2.0, 0.25, 0.3)
            .with_rho(0.0)
            .with_hawkes_params(0.5, 2.0, 1.0)
            .with_grid_steps(&[1.0, 0.02, 0.1])
            .optimize(&state, &grads);

        let ctrl_pos = HestonHawkes::new(0.5, 1.5)
            .with_heston_params(2.0, 0.25, 0.3)
            .with_rho(0.5)
            .with_hawkes_params(0.5, 2.0, 1.0)
            .with_grid_steps(&[1.0, 0.02, 0.1])
            .optimize(&state, &grads);

        assert!(
            ctrl_pos.lambda_minus[1] > ctrl0.lambda_minus[1],
            "Positive rho should increase downward variance transport in 3D model"
        );
    }

    #[test]
    fn hawkes_intensity_dimension_uses_upwinding() {
        let m = default_model();
        let grads = zero_grads();
        // Low lambda: drift positive -> lambda_plus[2] > 0
        let ctrl = m.optimize(&[0.0, 0.25, 0.1], &grads);
        assert!(ctrl.lambda_plus[2] > 0.0);
        assert!(ctrl.lambda_minus[2] < 1e-12);
    }

    #[test]
    fn builder_produces_correct_params() {
        let m = HestonHawkes::new(0.3, 2.0)
            .with_heston_params(3.0, 0.1, 0.5)
            .with_rho(-0.2)
            .with_hawkes_params(0.8, 4.0, 2.0)
            .with_grid_steps(&[2.0, 0.05, 0.2]);
        assert_eq!(m.gamma, 0.3);
        assert_eq!(m.kappa, 2.0);
        assert_eq!(m.v_kappa, 3.0);
        assert_eq!(m.rho, -0.2);
        assert_eq!(m.alpha, 0.8);
        assert_eq!(m.dq, 2.0);
        assert_eq!(m.dv, 0.05);
        assert_eq!(m.lambda_step, 0.2);
    }

    #[test]
    fn all_three_dimensions_have_transport() {
        let m = default_model();
        let grads = zero_grads();
        let ctrl = m.optimize(&[1.0, 0.1, 0.5], &grads);
        assert!(ctrl.lambda_plus[0] > 0.0);
        assert!(ctrl.lambda_minus[0] > 0.0);
        assert!(ctrl.lambda_plus[1] > 0.0 || ctrl.lambda_minus[1] > 0.0);
        assert!(ctrl.lambda_plus[2] > 0.0 || ctrl.lambda_minus[2] > 0.0);
    }

    #[test]
    fn is_diffusion_dimension_correct() {
        let m = default_model();
        assert!(
            !m.is_diffusion_dimension(0),
            "Inventory is controlled, not diffusion"
        );
        assert!(m.is_diffusion_dimension(1), "Variance is diffusion");
        assert!(
            !m.is_diffusion_dimension(2),
            "Hawkes intensity is deterministic drift"
        );
    }
}