solver/numeric/bsde/
solution.rs1use super::config::BasisFunctionType;
2use crate::models::traits::{ControlOutput, Gradients, Model};
3use crate::numeric::basis::{Basis, PolynomialBasis, ScaledBasis};
4
5pub struct BsdeSolution<const N: usize> {
10 pub coefficients: Vec<Vec<f64>>,
13 pub basis_type: BasisFunctionType,
15 pub dt: f64,
17 pub steps: usize,
19 pub x_init: [f64; N],
21 pub feature_scale: Vec<f64>,
23 pub clamp_range: Option<(f64, f64)>,
25 pub gradient_steps: [f64; N],
27 pub v0: f64,
29 pub control_t0: ControlOutput<N>,
31}
32
33impl<const N: usize> BsdeSolution<N> {
34 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 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 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 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 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
112fn 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}