Skip to main content

solver/numeric/bsde/
solution.rs

1use super::config::BasisFunctionType;
2use crate::models::traits::{ControlOutput, Gradients, Model};
3use crate::numeric::basis::{Basis, PolynomialBasis, ScaledBasis};
4
5/// Complete BSDE solution storing regression coefficients for all time steps.
6///
7/// Enables point evaluation of the value function and optimal controls at any
8/// (time_index, state) pair without re-running the solver.
9pub struct BsdeSolution<const N: usize> {
10    /// Regression coefficients at each backward time step.
11    /// `coefficients[t]` contains the basis weights at time index `t`.
12    pub coefficients: Vec<Vec<f64>>,
13    /// Basis function configuration.
14    pub basis_type: BasisFunctionType,
15    /// Time step size.
16    pub dt: f64,
17    /// Number of backward time steps.
18    pub steps: usize,
19    /// Initial state used for centering.
20    pub x_init: [f64; N],
21    /// Feature scaling per dimension.
22    pub feature_scale: Vec<f64>,
23    /// Clamping range for basis inputs.
24    pub clamp_range: Option<(f64, f64)>,
25    /// Gradient step sizes per dimension.
26    pub gradient_steps: [f64; N],
27    /// Value function at t=0 and x_init.
28    pub v0: f64,
29    /// Optimal control at t=0 and x_init.
30    pub control_t0: ControlOutput<N>,
31}
32
33impl<const N: usize> BsdeSolution<N> {
34    /// Reconstruct the scaled basis used during solving.
35    fn make_basis(&self) -> ScaledBasis<N, PolynomialBasis> {
36        let mut inner = self.basis_type.to_basis::<N>();
37        if let Some((min, max)) = self.clamp_range {
38            inner.clamp_range = Some((min, max));
39        }
40        ScaledBasis::new(inner, self.feature_scale.clone()).with_center(self.x_init)
41    }
42
43    /// Evaluate the value function at a given time index and state.
44    ///
45    /// `time_index` ranges from 0 (start) to `self.steps` (terminal).
46    pub fn evaluate(&self, time_index: usize, state: &[f64; N]) -> f64 {
47        if time_index >= self.coefficients.len() || self.coefficients[time_index].is_empty() {
48            return 0.0;
49        }
50        let basis = self.make_basis();
51        basis.eval_dot(state, &self.coefficients[time_index])
52    }
53
54    /// Evaluate optimal control at a given time index and state.
55    pub fn evaluate_control<M: Model<N>>(
56        &self,
57        time_index: usize,
58        state: &[f64; N],
59        model: &M,
60    ) -> ControlOutput<N> {
61        if time_index >= self.coefficients.len() || self.coefficients[time_index].is_empty() {
62            return ControlOutput {
63                lambda_plus: [0.0; N],
64                lambda_minus: [0.0; N],
65                flow: 0.0,
66            };
67        }
68        let basis = self.make_basis();
69        let coeffs = &self.coefficients[time_index];
70        let (fwd, bwd) = compute_gradients(state, coeffs, &basis, &self.gradient_steps);
71        let grads = Gradients { fwd, bwd };
72        model.optimize(state, &grads)
73    }
74
75    /// Evaluate optimal spreads across all time steps for a given state.
76    ///
77    /// Returns vectors of (bid_spread, ask_spread) for each time step from 0 to steps-1.
78    pub fn evaluate_spread_trajectory<M: Model<N>>(
79        &self,
80        state: &[f64; N],
81        model: &M,
82        a: f64,
83        kappa: f64,
84    ) -> Vec<(f64, f64)> {
85        let basis = self.make_basis();
86        let mut result = Vec::with_capacity(self.steps);
87        for t in 0..self.steps {
88            if self.coefficients[t].is_empty() {
89                result.push((f64::NAN, f64::NAN));
90                continue;
91            }
92            let coeffs = &self.coefficients[t];
93            let (fwd, bwd) = compute_gradients(state, coeffs, &basis, &self.gradient_steps);
94            let grads = Gradients { fwd, bwd };
95            let ctrl = model.optimize(state, &grads);
96
97            let bid = -(ctrl.lambda_plus[0] / a).ln() / kappa;
98            let ask = -(ctrl.lambda_minus[0] / a).ln() / kappa;
99            result.push((bid, ask));
100        }
101        result
102    }
103
104    /// Convert a time value to the nearest time index.
105    /// `t` is time elapsed from 0 (start of horizon).
106    pub fn time_to_index(&self, t: f64) -> usize {
107        let idx = (t / self.dt).round() as usize;
108        idx.min(self.steps.saturating_sub(1))
109    }
110}
111
112/// Finite-difference gradient computation from basis coefficients.
113fn compute_gradients<const N: usize>(
114    state: &[f64; N],
115    coeffs: &[f64],
116    basis: &impl Basis<N>,
117    gradient_steps: &[f64; N],
118) -> ([f64; N], [f64; N]) {
119    let mut fwd = [0.0; N];
120    let mut bwd = [0.0; N];
121    let v_curr = basis.eval_dot(state, coeffs);
122
123    for i in 0..N {
124        let h = gradient_steps[i];
125        let h_inv = 1.0 / h;
126
127        let mut next_state = *state;
128        next_state[i] += h;
129        fwd[i] = (basis.eval_dot(&next_state, coeffs) - v_curr) * h_inv;
130
131        let mut prev_state = *state;
132        prev_state[i] -= h;
133        bwd[i] = (v_curr - basis.eval_dot(&prev_state, coeffs)) * h_inv;
134    }
135    (fwd, bwd)
136}