solver/models/

avellaneda_hawkes.rs

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

/// Avellaneda-Stoikov model with Hawkes Process for order arrival intensity.
///
/// State Space: [Inventory (q), Intensity (lambda)]
/// - q: Discrete inventory level
/// - lambda: Continuous order arrival intensity
///
/// Dynamics:
/// - q: Jumps ±1 upon fills.
/// - lambda: Mean-reverting jump-diffusion (approximated).
///   d(lambda) = beta * (mu - lambda) * dt + alpha * dN_market
///   where N_market is the market order counting process with intensity lambda.
///
/// The solver uses a grid-based approach. We approximate the Hawkes dynamics
/// on the lambda dimension using an upwind finite difference scheme mapped to
/// the solver's jump interface.
#[derive(Clone)]
pub struct AvellanedaHawkes {
    pub gamma: f64,
    pub sigma: f64,
    pub kappa: f64,
    // Hawkes Parameters
    pub alpha: f64, // Jump size of intensity
    pub beta: f64,  // Decay rate of intensity
    pub mu: f64,    // Baseline intensity level (mean reversion target)
    pub lambda_step: f64,
    /// Grid step size for inventory dimension (set from `grid.dx[0]` for FD).
    pub dq: f64,
    /// Terminal condition V(T, q, lambda).  Defaults to Zero.
    pub terminal_condition: TerminalCondition,
    /// 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,
    /// Optional terminal liquidation half-spread used when terminal_condition is
    /// LiquidationCost. If None, uses the AS base spread.
    pub terminal_liquidation_half_spread: Option<f64>,
}

impl AvellanedaHawkes {
    pub fn new(gamma: f64, sigma: f64, kappa: f64, alpha: f64, beta: f64, mu: f64) -> Self {
        Self {
            gamma,
            sigma,
            kappa,
            alpha,
            beta,
            mu,
            lambda_step: 1.0,
            dq: 1.0,
            terminal_condition: TerminalCondition::Zero,
            q_min: f64::NEG_INFINITY,
            q_max: f64::INFINITY,
            terminal_liquidation_half_spread: None,
        }
    }

    pub fn with_terminal_liquidation_half_spread(mut self, half_spread: f64) -> Self {
        self.terminal_liquidation_half_spread = Some(half_spread.max(0.0));
        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_terminal_condition(mut self, terminal_condition: TerminalCondition) -> Self {
        self.terminal_condition = terminal_condition;
        self
    }

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

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

    /// Computes optimal spreads based on current inventory and intensity.
    /// Note: The base intensity 'A' from standard Avellaneda is replaced by the state variable 'lambda'.
    pub fn get_spreads(&self, _lambda: f64, grads: &Gradients<2>) -> (f64, f64) {
        // Standard Avellaneda formula: delta = (1/gamma) * ln(1 + gamma/kappa) + ...
        // The 'A' parameter cancels out in the spread optimization, so it doesn't appear here explicitly.
        // However, the value function V depends on lambda, so grads will reflect that.

        let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa).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 = -10.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 AvellanedaHawkes {
    type Process = ();

    fn process(&self) {}

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

        // 1. Optimize Inventory Controls (Dim 0)
        let (d_bid, d_ask) = self.get_spreads(lambda, grads);

        // Fill probabilities (intensities)
        // 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);

        // Shut off the forbidden quoting side at the inventory boundary.
        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 part for Inventory
        let hamiltonian_inv = (lambda_bid + lambda_ask) / (self.gamma + self.kappa);

        // Risk Penalty
        let risk_penalty = -0.5 * self.gamma * self.sigma.powi(2) * q.powi(2);

        // Cancel solver's q-dimension transport. The solver adds:
        //   lambda_bid * (V[q+dq]-V[q]) + lambda_ask * (V[q-dq]-V[q])
        //   = dq * (lambda_bid * fwd[0] - lambda_ask * bwd[0])
        // We must subtract exactly that amount to replace it with the Hamiltonian.
        let drift_correction_q = self.dq * (lambda_bid * grads.fwd[0] - lambda_ask * grads.bwd[0]);

        // 2. Hawkes Dynamics (Dim 1) — properly encoded in matrix for implicit stability
        // Expected drift: D = beta*(mu - lambda) + lambda * alpha
        let net_drift_lambda = self.beta * (self.mu - lambda) + lambda * self.alpha;
        let drift_abs = net_drift_lambda.abs() / self.lambda_step;

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

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

    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).ln();
                let half_spread = self.terminal_liquidation_half_spread.unwrap_or(base_spread);
                -q.abs() * half_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 mut next = *current_state;
        let lambda = current_state[1];
        let net_drift_lambda = self.beta * (self.mu - lambda) + self.alpha * lambda;

        // Keep Hawkes intensity evolution consistent with the first-order drift treatment
        // used in the finite-difference HJB discretization.
        let _ = noise;
        next[1] += net_drift_lambda * dt;
        next[1] = next[1].max(0.0);

        next
    }

    fn next_step_controlled(
        &self,
        current_state: &[f64; 2],
        control: &ControlOutput<2>,
        dt: f64,
        noise: &[f64; 2],
    ) -> [f64; 2] {
        let mut next = *current_state;
        let lambda = current_state[1];

        // Dim 0 (q): Bernoulli fill simulation.
        // Transform noise[0] (N(0,1)) to uniform [0,1] via normal CDF, then gate against
        // fill probabilities.  Bid and ask fills are mutually exclusive within one step,
        // which is accurate to O(dt^2) via the small-dt Poisson approximation.
        let u = normal_cdf(noise[0]);
        let p_bid = (control.lambda_plus[0] * dt).clamp(0.0, 1.0);
        let p_ask = (control.lambda_minus[0] * dt).clamp(0.0, 1.0);
        if u < p_bid {
            next[0] += 1.0; // bid fill: MM sold, inventory increases
        } else if u > 1.0 - p_ask {
            next[0] -= 1.0; // ask fill: MM bought, inventory decreases
        }

        // Dim 1 (lambda): deterministic mean-reversion drift (same as next_step).
        let net_drift_lambda = self.beta * (self.mu - lambda) + self.alpha * lambda;
        let _ = noise[1];
        next[1] += net_drift_lambda * dt;
        next[1] = next[1].max(0.0);

        next
    }

    fn is_diffusion_dimension(&self, _dim: usize) -> bool {
        false // neither q nor lambda has Brownian noise
    }

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

    fn gradient_step(&self, dim: usize) -> f64 {
        if dim == 1 { self.lambda_step } else { 1.0 }
    }

    /// The base arrival intensity is state-dependent: returns the current
    /// intensity `state[1]` (clamped above 0) for intensity-to-spread conversion.
    fn fill_rate_base(&self, state: &[f64; 2]) -> f64 {
        state[1].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<2> {
        Gradients {
            fwd: [0.0; 2],
            bwd: [0.0; 2],
        }
    }

    fn default_model() -> AvellanedaHawkes {
        AvellanedaHawkes::new(0.5, 0.5, 1.5, 0.5, 2.0, 1.0)
            .with_lambda_step(0.1)
            .with_dq(1.0)
    }

    #[test]
    fn zero_gradient_gives_base_spread() {
        let m = default_model();
        let grads = zero_grads();
        let (bid, ask) = m.get_spreads(1.0, &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 intensity_scales_fill_rates() {
        let m = default_model();
        let grads = zero_grads();
        let ctrl_low = m.optimize(&[0.0, 0.5], &grads);
        let ctrl_high = m.optimize(&[0.0, 2.0], &grads);
        let total_low = ctrl_low.lambda_plus[0] + ctrl_low.lambda_minus[0];
        let total_high = ctrl_high.lambda_plus[0] + ctrl_high.lambda_minus[0];
        assert!(
            total_high > total_low,
            "Higher intensity should give higher fill rates: low={}, high={}",
            total_low,
            total_high
        );
    }

    #[test]
    fn subcritical_regime_mean_reverts() {
        let m = AvellanedaHawkes::new(0.5, 0.5, 1.5, 0.5, 2.0, 1.0);
        assert!(m.alpha < m.beta, "Subcritical: alpha < beta");

        // Equilibrium intensity = beta * mu / (beta - alpha)
        let eq = m.beta * m.mu / (m.beta - m.alpha);

        let state_high = [0.0, eq + 1.0];
        let next_high = m.next_step(&state_high, 0.01, &[0.0, 0.0]);
        assert!(
            next_high[1] < state_high[1],
            "Above equilibrium, intensity should decrease"
        );

        let state_low = [0.0, eq - 0.5];
        let next_low = m.next_step(&state_low, 0.01, &[0.0, 0.0]);
        assert!(
            next_low[1] > state_low[1],
            "Below equilibrium, intensity should increase"
        );
    }

    #[test]
    fn intensity_drift_uses_upwinding() {
        let m = default_model();
        let grads = zero_grads();
        let ctrl = m.optimize(&[0.0, 0.1], &grads);
        assert!(
            ctrl.lambda_plus[1] > 0.0,
            "Positive drift should give lambda_plus > 0"
        );
        assert!(
            ctrl.lambda_minus[1] < 1e-12,
            "Positive drift should give lambda_minus ~ 0"
        );
    }

    #[test]
    fn intensity_stays_non_negative() {
        let m = default_model();
        let state = [0.0, 0.001];
        let next = m.next_step(&state, 0.01, &[0.0, 0.0]);
        assert!(next[1] >= 0.0, "Intensity must stay non-negative");
    }

    #[test]
    fn alpha_zero_makes_drift_simple_mean_reversion() {
        let m = AvellanedaHawkes::new(0.5, 0.5, 1.5, 0.0, 2.0, 1.0).with_lambda_step(0.1);
        let grads = zero_grads();
        // With alpha=0, drift = beta*(mu - lambda). At lambda=mu, drift is zero.
        let ctrl = m.optimize(&[0.0, m.mu], &grads);
        assert!(
            ctrl.lambda_plus[1].abs() < 1e-10 && ctrl.lambda_minus[1].abs() < 1e-10,
            "At lambda=mu with alpha=0, net drift should be zero"
        );
    }
}