solver/models/

heston_hawkes.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;

/// 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, derivs: &StateDerivatives<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();

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

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

        (delta_bid, delta_ask)
    }

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

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

impl ControlProblem<3> for HestonHawkes {
    type Control = MarketMakingControl;

    fn optimize(&self, _t: f64, state: &[f64; 3], derivs: &StateDerivatives<3>) -> Self::Control {
        let q = state[0];
        let v = state[1];
        let lambda = state[2];
        let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa).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 max_fill_mult = 20.0_f64;
        let lambda_bid = lambda * (-self.kappa * delta_bid).exp().min(max_fill_mult);
        let lambda_ask = lambda * (-self.kappa * delta_ask).exp().min(max_fill_mult);

        let _ = (q, v);
        MarketMakingControl::new(lambda_bid, lambda_ask)
    }

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

    fn bsde_driver(
        &self,
        _t: f64,
        state: &[f64; 3],
        control: &Self::Control,
        _derivs: &StateDerivatives<3>,
        dt: f64,
    ) -> f64 {
        // Full-value problem: the variance and Hawkes-intensity transport is
        // already simulated forward, so the backward driver adds only the
        // running reward plus the local inventory-risk source. The reward is
        // bounded at `1/dt` per fill side to match the forward fill-probability
        // clamp `lambda * dt <= 1`.
        let q = state[0];
        let v = state[1];
        let local = -0.5 * self.gamma * 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);
        reward + local
    }

    fn generator(
        &self,
        _t: f64,
        state: &[f64; 3],
        _control: &Self::Control,
        derivs: &StateDerivatives<3>,
    ) -> f64 {
        let q = state[0];
        let v = state[1];
        let lambda = state[2];

        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 net_drift_lambda = self.beta * (self.mu_lambda - lambda) + self.alpha * lambda;

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

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

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

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

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

        let mut next_q = q;
        let u = normal_cdf(noise[0]);
        let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa).ln();
        let lambda_bid = lambda * (-self.kappa * base_spread).exp();
        let lambda_ask = lambda * (-self.kappa * 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_q += 1.0;
        } else if u > 1.0 - p_ask {
            next_q -= 1.0;
        }

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

        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 next_step_controlled(
        &self,
        _t: f64,
        state: &[f64; 3],
        control: &Self::Control,
        dt: f64,
        noise: &[f64; 3],
    ) -> [f64; 3] {
        let v = state[1];
        let lambda = state[2];
        let sqrt_dt = dt.sqrt();

        let mut next_q = state[0];
        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_q += 1.0;
        } else if u > 1.0 - p_ask {
            next_q -= 1.0;
        }

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

        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
    }

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

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

    fn transport(
        &self,
        _t: f64,
        state: &[f64; 3],
        control: &Self::Control,
        derivs: &StateDerivatives<3>,
    ) -> Transport<3> {
        let q = state[0];
        let v = state[1].max(1e-5);
        let lambda = state[2].max(0.0);

        let mu_v = self.v_kappa * (self.v_theta - v) - self.gamma * 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 net_drift_lambda = self.beta * (self.mu_lambda - lambda) + self.alpha * lambda;
        let rate = net_drift_lambda.abs() / self.lambda_step.abs().max(1e-12);
        let (plus_lam, minus_lam) = if net_drift_lambda >= 0.0 {
            (rate, 0.0)
        } else {
            (0.0, rate)
        };

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

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

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

    fn zero_derivs() -> StateDerivatives<3> {
        StateDerivatives::new([0.0; 3], [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 derivs = zero_derivs();
        let (bid, ask) = m.get_spreads(&derivs);
        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 bsde_driver_excludes_variance_and_intensity_transport() {
        let m = default_model();
        let state = [2.0, 0.25, 1.0];
        let control = ControlProblem::optimize(&m, 0.0, &state, &zero_derivs());

        let mut derivs = zero_derivs();
        derivs.grad[1] = 1.0;
        derivs.grad[2] = 1.0;

        let dt = 0.01;
        let q = state[0];
        let v = state[1];
        let local = -0.5 * m.gamma * 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);
        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)
        );
    }

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

        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, 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, 0.0]);
        assert_eq!(next[0], 2.0);
    }

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

        let g0 = ControlProblem::generator(
            &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]),
            0.0,
            &state,
            &control,
            &derivs,
        );

        let g_pos = ControlProblem::generator(
            &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]),
            0.0,
            &state,
            &control,
            &derivs,
        );

        assert!(
            g0 > g_pos,
            "Positive rho lowers effective variance drift in the 3D generator"
        );
    }

    #[test]
    fn hawkes_intensity_dimension_uses_upwinding() {
        let m = default_model();
        let mut derivs = zero_derivs();
        derivs.grad[2] = 1.0;
        let control = MarketMakingControl::new(0.0, 0.0);
        let g = ControlProblem::generator(&m, 0.0, &[0.0, 0.25, 0.1], &control, &derivs);
        assert!(g > 0.0);
    }

    #[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 is_diffusion_dimension_correct() {
        let m = default_model();
        assert!(
            !ControlProblem::is_diffusion_dimension(&m, 0),
            "Inventory is controlled, not diffusion"
        );
        assert!(
            ControlProblem::is_diffusion_dimension(&m, 1),
            "Variance is diffusion"
        );
        assert!(
            !ControlProblem::is_diffusion_dimension(&m, 2),
            "Hawkes intensity is deterministic drift"
        );
    }

    #[test]
    fn dimension_kinds_match_state_space() {
        let m = default_model();
        assert_eq!(
            PdeProblem::dimension_kind(&m, 0),
            DimensionKind::DiscreteJump
        );
        assert_eq!(PdeProblem::dimension_kind(&m, 1), DimensionKind::Diffusion);
        assert_eq!(
            PdeProblem::dimension_kind(&m, 2),
            DimensionKind::DeterministicDrift
        );
    }
}