solver/numeric/bsde/

solution.rs

use super::config::BasisFunctionType;
use crate::models::control::{ControlProblem, StateDerivatives};
use crate::models::market_making::MarketMakingControl;
use crate::numeric::basis::{Basis, PolynomialBasis, ScaledBasis};

/// Complete BSDE solution storing regression coefficients for all time steps.
///
/// Enables point evaluation of the value function and optimal controls at any
/// `(time_index, state)` pair without re-running the solver.
pub struct BsdeSolution<const N: usize, C> {
    /// Regression coefficients at each backward time step.
    /// `coefficients[t]` contains the basis weights at time index `t`.
    pub coefficients: Vec<Vec<f64>>,
    /// Basis function configuration.
    pub basis_type: BasisFunctionType,
    /// Time step size.
    pub dt: f64,
    /// Number of backward time steps.
    pub steps: usize,
    /// Initial state used for centering.
    pub x_init: [f64; N],
    /// Feature scaling per dimension.
    pub feature_scale: Vec<f64>,
    /// Clamping range for basis inputs.
    pub clamp_range: Option<(f64, f64)>,
    /// Gradient step sizes per dimension.
    pub gradient_steps: [f64; N],
    /// Value function at t=0 and x_init.
    pub v0: f64,
    /// Optimal control at t=0 and x_init.
    pub control_t0: C,
}

impl<const N: usize, C> BsdeSolution<N, C> {
    /// Reconstruct the scaled basis used during solving.
    fn make_basis(&self) -> ScaledBasis<N, PolynomialBasis> {
        let mut inner = self.basis_type.to_basis::<N>();
        if let Some((min, max)) = self.clamp_range {
            inner.clamp_range = Some((min, max));
        }
        ScaledBasis::new(inner, self.feature_scale.clone()).with_center(self.x_init)
    }

    /// Evaluate the value function at a given time index and state.
    ///
    /// `time_index` ranges from 0 (start) to `self.steps` (terminal).
    pub fn evaluate(&self, time_index: usize, state: &[f64; N]) -> f64 {
        if time_index >= self.coefficients.len() || self.coefficients[time_index].is_empty() {
            return 0.0;
        }
        let basis = self.make_basis();
        basis.eval_dot(state, &self.coefficients[time_index])
    }

    /// Convert a time value to the nearest time index.
    ///
    /// `t` is time elapsed from 0 (start of horizon).
    pub fn time_to_index(&self, t: f64) -> usize {
        let idx = (t / self.dt).round() as usize;
        idx.min(self.steps.saturating_sub(1))
    }
}

impl<const N: usize> BsdeSolution<N, MarketMakingControl> {
    /// Evaluate the optimal market-making control at a time index and state.
    ///
    /// Returns `None` when no regression coefficients are available for the
    /// requested time index (the terminal step).
    pub fn evaluate_control<M: ControlProblem<N, Control = MarketMakingControl>>(
        &self,
        time_index: usize,
        state: &[f64; N],
        problem: &M,
    ) -> Option<MarketMakingControl> {
        if time_index >= self.coefficients.len() || self.coefficients[time_index].is_empty() {
            return None;
        }
        let basis = self.make_basis();
        let coeffs = &self.coefficients[time_index];
        let derivs = state_derivatives(state, coeffs, &basis, &self.gradient_steps);
        Some(problem.optimize(self.dt * time_index as f64, state, &derivs))
    }

    /// Evaluate optimal spreads across all time steps for a given state.
    ///
    /// Returns vectors of `(bid_spread, ask_spread)` for each time step from
    /// `0` to `steps - 1`. Time steps without coefficients yield `NaN`.
    pub fn evaluate_spread_trajectory<M: ControlProblem<N, Control = MarketMakingControl>>(
        &self,
        state: &[f64; N],
        problem: &M,
        a: f64,
        kappa: f64,
    ) -> Vec<(f64, f64)> {
        let basis = self.make_basis();
        let mut result = Vec::with_capacity(self.steps);
        for t in 0..self.steps {
            if self.coefficients[t].is_empty() {
                result.push((f64::NAN, f64::NAN));
                continue;
            }
            let coeffs = &self.coefficients[t];
            let derivs = state_derivatives(state, coeffs, &basis, &self.gradient_steps);
            let ctrl = problem.optimize(self.dt * t as f64, state, &derivs);
            let spreads = ctrl.to_spreads(a, kappa);
            result.push((spreads.bid_spread, spreads.ask_spread));
        }
        result
    }
}

/// Finite-difference derivatives of the regression model at a state.
pub(crate) fn state_derivatives<const N: usize>(
    state: &[f64; N],
    coeffs: &[f64],
    basis: &impl Basis<N>,
    gradient_steps: &[f64; N],
) -> StateDerivatives<N> {
    let v_curr = basis.eval_dot(state, coeffs);
    let mut grad = [0.0; N];
    let mut hessian = [0.0; N];
    let mut hessian_full = [[0.0; N]; N];
    let mut fwd = [0.0; N];
    let mut bwd = [0.0; N];

    // Diagonal entries and first derivatives.
    for i in 0..N {
        let h = gradient_steps[i];
        let h_inv = 1.0 / h;

        let mut plus = *state;
        plus[i] += h;
        let v_plus = basis.eval_dot(&plus, coeffs);

        let mut minus = *state;
        minus[i] -= h;
        let v_minus = basis.eval_dot(&minus, coeffs);

        grad[i] = (v_plus - v_minus) * 0.5 * h_inv;
        hessian[i] = (v_plus - 2.0 * v_curr + v_minus) * h_inv * h_inv;
        fwd[i] = (v_plus - v_curr) * h_inv;
        bwd[i] = (v_curr - v_minus) * h_inv;
        hessian_full[i][i] = hessian[i];
    }

    // Mixed second derivatives via the centered four-corner stencil
    // d^2V/dx_i dx_j = (V_{++} - V_{+-} - V_{-+} + V_{--}) / (4 h_i h_j).
    for i in 0..N {
        for j in (i + 1)..N {
            let hi = gradient_steps[i];
            let hj = gradient_steps[j];
            let denom = 4.0 * hi * hj;

            let mut pp = *state;
            pp[i] += hi;
            pp[j] += hj;

            let mut pm = *state;
            pm[i] += hi;
            pm[j] -= hj;

            let mut mp = *state;
            mp[i] -= hi;
            mp[j] += hj;

            let mut mm = *state;
            mm[i] -= hi;
            mm[j] -= hj;

            let mixed = (basis.eval_dot(&pp, coeffs)
                - basis.eval_dot(&pm, coeffs)
                - basis.eval_dot(&mp, coeffs)
                + basis.eval_dot(&mm, coeffs))
                / denom;
            hessian_full[i][j] = mixed;
            hessian_full[j][i] = mixed;
        }
    }

    StateDerivatives::with_full_hessian(grad, hessian_full, fwd, bwd)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::numeric::basis::Basis;

    /// Basis with a single feature phi(x) = x0 * x1, so the represented value
    /// is V(x) = x0 * x1 when the weight is 1.
    struct BilinearBasis;

    impl Basis<2> for BilinearBasis {
        fn num_features(&self) -> usize {
            1
        }

        fn evaluate(&self, state: &[f64; 2], buffer: &mut [f64]) {
            buffer[0] = state[0] * state[1];
        }
    }

    #[test]
    fn mixed_second_derivative_matches_known_quadratic() {
        // V(x0, x1) = x0 * x1 has exact mixed derivative d^2V/dx0 dx1 = 1 and
        // zero diagonal second derivatives. The centered four-corner stencil is
        // exact for this bilinear function, isolating the mixed-derivative
        // assembly rather than the finite-difference truncation error.
        let basis = BilinearBasis;
        let state = [0.3, -0.2];
        let steps = [1e-3, 1e-3];

        let d = state_derivatives(&state, &[1.0], &basis, &steps);

        assert!(
            (d.hessian_full[0][1] - 1.0).abs() < 1e-9,
            "mixed derivative {}, expected 1.0",
            d.hessian_full[0][1]
        );
        assert!(
            (d.hessian_full[1][0] - 1.0).abs() < 1e-9,
            "mixed derivative must be symmetric"
        );
        assert!((d.hessian[0]).abs() < 1e-9);
        assert!((d.hessian[1]).abs() < 1e-9);
    }
}