solver/models/

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

// Re-export so existing `hjb_solver::models::bilateral_hawkes::TerminalCondition` paths compile.
pub use super::TerminalCondition;

/// Avellaneda-Stoikov model with a bilateral (2-D) Hawkes process.
///
/// # State Space  `[q, lambda_plus, lambda_minus]`
///
/// | dim | symbol   | description                                      |
/// |-----|----------|--------------------------------------------------|
/// | 0   | q        | Inventory level (integer, jump-controlled)       |
/// | 1   | lambda+  | Buy market-order intensity (ODE drift, no noise) |
/// | 2   | lambda-  | Sell market-order intensity (ODE drift, no noise)|
///
/// # Dynamics
///
/// `d(lambda+) = beta*(mu - lambda+) dt + alpha dN+`
/// `d(lambda-) = beta*(mu - lambda-) dt + alpha dN-`
///
/// `N+` counts buy market orders (rate `lambda+`), `N-` counts sell market orders (rate `lambda-`).
/// Each side self-excites independently — there is no cross-excitation.
///
/// # Market-making mapping
///
/// | event              | driven by | inventory effect |
/// |--------------------|-----------|-----------------|
/// | Sell MO hits bid   | lambda-   | q -> q + 1      |
/// | Buy  MO hits ask   | lambda+   | q -> q - 1      |
///
/// Effective fill rates:
/// * `Lambda_bid = lambda- * exp(-kappa * delta_bid)`  (sell MO hits MM bid)
/// * `Lambda_ask = lambda+ * exp(-kappa * delta_ask)`  (buy  MO hits MM ask)
///
/// # Hawkes approximation used in the FD solver
///
/// The lambda dimensions use the fluid (mean-field) approximation:
///
///   `d(lambda+)/dt ~ beta*(mu - lambda+) + alpha * lambda+`
///   `d(lambda-)/dt ~ beta*(mu - lambda-) + alpha * lambda-`
///
/// This is identical to the treatment in `AvellanedaHawkes` and ensures the
/// stationary mean of each intensity sits at `mu / (1 - alpha/beta)` in the
/// value-function PDE.  Using discrete fill-triggered jumps instead requires
/// `lambda_step << alpha` for accuracy; at the grid resolution affordable for
/// 4-D tables that condition is violated, causing the PDE to see a much lower
/// effective mean intensity than the actual simulation (`mu` vs `mu/(1-rho)`),
/// which corrupts the value function and the extracted spreads.
#[derive(Clone, Debug)]
pub struct BilateralHawkes {
    /// Risk-aversion parameter.
    pub gamma: f64,
    /// Price volatility.
    pub sigma: f64,
    /// Fill-rate decay (Avellaneda-Stoikov kappa).
    pub kappa: f64,
    // --- Hawkes parameters (shared for both sides) ---
    /// Jump size added to the excited side's intensity.
    pub alpha: f64,
    /// Mean-reversion speed.
    pub beta: f64,
    /// Baseline intensity (mean-reversion target).
    pub mu: f64,
    /// Grid step for both lambda+ and lambda- dimensions.
    pub lambda_step: f64,
    /// Grid step for the inventory dimension.
    pub dq: f64,
    /// Minimum inventory (lower hard boundary). Defaults to -infinity.
    pub q_min: f64,
    /// Maximum inventory (upper hard boundary). Defaults to +infinity.
    pub q_max: f64,
    /// Terminal condition for the value function at T. Defaults to Zero.
    pub terminal_condition: TerminalCondition,
    /// 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 BilateralHawkes {
    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,
            q_min: f64::NEG_INFINITY,
            q_max: f64::INFINITY,
            terminal_condition: TerminalCondition::Zero,
            terminal_liquidation_half_spread: None,
        }
    }

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

    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_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 bid/ask half-spreads from value-function gradients.
    ///
    /// Bid spread: driven by the inventory gradient in the +q direction (sell MO fills).
    /// Ask spread: driven by the inventory gradient in the -q direction (buy  MO fills).
    pub fn get_spreads(&self, derivs: &StateDerivatives<3>) -> (f64, f64) {
        let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa).ln();
        let delta_bid = (base_spread - derivs.fwd[0]).clamp(-10.0, 10.0);
        let delta_ask = (base_spread + derivs.bwd[0]).clamp(-10.0, 10.0);
        (delta_bid, delta_ask)
    }

    /// The base arrival intensity is state-dependent: returns the average of
    /// the buy/sell intensity states for intensity-to-spread conversion.
    pub fn fill_rate_base(&self, state: &[f64; 3]) -> f64 {
        ((state[1] + state[2]) * 0.5).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 BilateralHawkes {
    type Control = MarketMakingControl;

    fn optimize(&self, _t: f64, state: &[f64; 3], derivs: &StateDerivatives<3>) -> Self::Control {
        let q = state[0];
        let lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);

        let (d_bid, d_ask) = self.get_spreads(derivs);

        let max_mult = 20.0_f64;
        let lambda_bid_fill = lambda_minus * (-self.kappa * d_bid).exp().min(max_mult);
        let lambda_ask_fill = lambda_plus * (-self.kappa * d_ask).exp().min(max_mult);

        let lambda_bid_fill = if q >= self.q_max {
            0.0
        } else {
            lambda_bid_fill
        };
        let lambda_ask_fill = if q <= self.q_min {
            0.0
        } else {
            lambda_ask_fill
        };

        MarketMakingControl::new(lambda_bid_fill, lambda_ask_fill)
    }

    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 two Hawkes intensity drifts are 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.
        let q = state[0];
        let local = -0.5 * self.gamma * self.sigma.powi(2) * 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 lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);
        let net_drift_plus = self.beta * (self.mu - lambda_plus) + self.alpha * lambda_plus;
        let net_drift_minus = self.beta * (self.mu - lambda_minus) + self.alpha * lambda_minus;

        -0.5 * self.gamma * self.sigma.powi(2) * q.powi(2)
            + net_drift_plus * derivs.grad[1]
            + net_drift_minus * derivs.grad[2]
    }

    fn terminal(&self, state: &[f64; 3]) -> 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 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 lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);
        let drift_plus = self.beta * (self.mu - lambda_plus) + self.alpha * lambda_plus;
        let drift_minus = self.beta * (self.mu - lambda_minus) + self.alpha * lambda_minus;
        let mut next = [
            state[0],
            (lambda_plus + drift_plus * dt).max(0.0),
            (lambda_minus + drift_minus * dt).max(0.0),
        ];

        let u = normal_cdf(noise[0]);
        let lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);
        let base_spread = (1.0 / self.gamma) * (1.0 + self.gamma / self.kappa).ln();
        let lambda_bid_fill = lambda_minus * (-self.kappa * base_spread).exp();
        let lambda_ask_fill = lambda_plus * (-self.kappa * base_spread).exp();
        let p_bid = (lambda_bid_fill * dt).clamp(0.0, 1.0);
        let p_ask = (lambda_ask_fill * dt).clamp(0.0, 1.0);
        if u < p_bid {
            next[0] += 1.0;
            next[2] += self.alpha;
        } else if u > 1.0 - p_ask {
            next[0] -= 1.0;
            next[1] += self.alpha;
        }

        next[1] = next[1].max(0.0);
        next[2] = next[2].max(0.0);
        next
    }

    fn next_step_controlled(
        &self,
        _t: f64,
        state: &[f64; 3],
        control: &Self::Control,
        dt: f64,
        noise: &[f64; 3],
    ) -> [f64; 3] {
        let lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);
        let drift_plus = self.beta * (self.mu - lambda_plus) + self.alpha * lambda_plus;
        let drift_minus = self.beta * (self.mu - lambda_minus) + self.alpha * lambda_minus;
        let mut next = [
            state[0],
            (lambda_plus + drift_plus * dt).max(0.0),
            (lambda_minus + drift_minus * dt).max(0.0),
        ];

        // 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;
            next[2] += self.alpha;
        } else if u > 1.0 - p_ask {
            next[0] -= 1.0;
            next[1] += self.alpha;
        }

        next[1] = next[1].max(0.0);
        next[2] = next[2].max(0.0);
        next
    }

    fn is_diffusion_dimension(&self, _dim: usize) -> bool {
        false
    }

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

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

    fn transport(
        &self,
        _t: f64,
        state: &[f64; 3],
        control: &Self::Control,
        derivs: &StateDerivatives<3>,
    ) -> Transport<3> {
        let q = state[0];
        let lambda_plus = state[1].max(0.0);
        let lambda_minus = state[2].max(0.0);
        let net_drift_plus = self.beta * (self.mu - lambda_plus) + self.alpha * lambda_plus;
        let net_drift_minus = self.beta * (self.mu - lambda_minus) + self.alpha * lambda_minus;
        let step = self.lambda_step.abs().max(1e-12);

        let rate_plus = net_drift_plus.abs() / step;
        let rate_minus = net_drift_minus.abs() / step;
        let (lp_plus, lp_minus) = if net_drift_plus >= 0.0 {
            (rate_plus, 0.0)
        } else {
            (0.0, rate_plus)
        };
        let (lm_plus, lm_minus) = if net_drift_minus >= 0.0 {
            (rate_minus, 0.0)
        } else {
            (0.0, rate_minus)
        };

        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 * self.sigma.powi(2) * q.powi(2);

        Transport::new(
            [control.bid_intensity, lp_plus, lm_plus],
            [control.ask_intensity, lp_minus, lm_minus],
            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() -> BilateralHawkes {
        BilateralHawkes::new(0.5, 0.2, 1.5, 0.5, 2.0, 1.0)
            .with_lambda_step(0.1)
            .with_dq(1.0)
    }

    #[test]
    fn zero_gradient_gives_symmetric_base_spread() {
        let m = default_model();
        let (bid, ask) = m.get_spreads(&zero_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_intensity_transport() {
        let m = default_model();
        let state = [2.0, 1.0, 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 local = -0.5 * m.gamma * m.sigma.powi(2) * 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),
            "full driver adds intensity transport, bsde driver must not"
        );
    }

    #[test]
    fn next_step_controlled_uses_optimal_intensities() {
        let m = default_model();
        let state = [1.0, 1.0, 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 bid_driven_by_lambda_minus_ask_by_lambda_plus() {
        let m = default_model();
        let d = zero_derivs();
        let ctrl_high_minus = ControlProblem::optimize(&m, 0.0, &[0.0, 1.0, 5.0], &d);
        let ctrl_low_minus = ControlProblem::optimize(&m, 0.0, &[0.0, 1.0, 0.5], &d);
        assert!(
            ctrl_high_minus.bid_intensity > ctrl_low_minus.bid_intensity,
            "bid fill should scale with lambda-"
        );

        let ctrl_high_plus = ControlProblem::optimize(&m, 0.0, &[0.0, 5.0, 1.0], &d);
        let ctrl_low_plus = ControlProblem::optimize(&m, 0.0, &[0.0, 0.5, 1.0], &d);
        assert!(
            ctrl_high_plus.ask_intensity > ctrl_low_plus.ask_intensity,
            "ask fill should scale with lambda+"
        );
    }

    #[test]
    fn symmetric_state_gives_equal_bid_ask() {
        let m = default_model();
        let d = zero_derivs();
        let ctrl = ControlProblem::optimize(&m, 0.0, &[0.0, 1.0, 1.0], &d);
        assert!(
            (ctrl.bid_intensity - ctrl.ask_intensity).abs() < 1e-12,
            "bid fill == ask fill when lambda+ == lambda- and q == 0"
        );
    }

    #[test]
    fn generator_encodes_intensity_upwinding() {
        let m = default_model();
        let control = MarketMakingControl::new(1.0, 1.0);
        let state = [0.0, 0.1, 1.0];

        let zero = StateDerivatives::new([0.0; 3], [0.0; 3]);
        let mut plus_grad = zero;
        plus_grad.grad[1] = 1.0;

        let base = ControlProblem::generator(&m, 0.0, &state, &control, &zero);
        let upwind = ControlProblem::generator(&m, 0.0, &state, &control, &plus_grad);
        // At lambda+ = 0.1 the net drift is positive, so a positive gradient in
        // the lambda+ dimension increases the generator.
        assert!(upwind > base);
    }

    #[test]
    fn next_step_clamped_to_nonnegative() {
        let m = default_model();
        let state = [0.0, 0.0, 0.0];
        let next = ControlProblem::next_step(&m, 0.0, &state, 1.0, &[0.0; 3]);
        assert!(next[1] >= 0.0);
        assert!(next[2] >= 0.0);
    }

    #[test]
    fn terminal_zero_condition_returns_zero() {
        let m = default_model().with_terminal_condition(TerminalCondition::Zero);
        let v = ControlProblem::terminal(&m, &[3.0, 1.0, 1.0]);
        assert!(v.abs() < 1e-12);
    }

    #[test]
    fn terminal_liquidation_cost_matches_formula() {
        let m = default_model().with_terminal_condition(TerminalCondition::LiquidationCost);
        let q: f64 = -4.0;
        let base = (1.0 / m.gamma) * (1.0 + m.gamma / m.kappa).ln();
        let expected = -q.abs() * base;
        let got = ControlProblem::terminal(&m, &[q, 1.2, 0.8]);
        assert!((got - expected).abs() < 1e-12);
    }
}