Skip to main content

diffsol_c/
ode_options.rs

1use std::sync::{Arc, Mutex};
2
3use serde::{Deserialize, Serialize, Serializer};
4
5use crate::{error::DiffsolRtError, ode::Ode, solve::Solve};
6
7#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
8pub struct OdeSolverOptionsSnapshot {
9    pub max_nonlinear_solver_iterations: usize,
10    pub max_error_test_failures: usize,
11    pub update_jacobian_after_steps: usize,
12    pub update_rhs_jacobian_after_steps: usize,
13    pub threshold_to_update_jacobian: f64,
14    pub threshold_to_update_rhs_jacobian: f64,
15    pub min_timestep: f64,
16    #[serde(default = "default_derivative_solve_tolerance")]
17    pub derivative_solve_tolerance: f64,
18    #[serde(default = "default_derivative_solve_max_iter")]
19    pub derivative_solve_max_iter: usize,
20}
21
22fn default_derivative_solve_tolerance() -> f64 {
23    diffsol::OdeSolverOptions::<f64>::default().derivative_solve_tolerance
24}
25
26fn default_derivative_solve_max_iter() -> usize {
27    diffsol::OdeSolverOptions::<f64>::default().derivative_solve_max_iter
28}
29
30impl OdeSolverOptionsSnapshot {
31    pub(crate) fn from_solve(solve: &dyn Solve) -> Self {
32        Self {
33            max_nonlinear_solver_iterations: solve.ode_max_nonlinear_solver_iterations(),
34            max_error_test_failures: solve.ode_max_error_test_failures(),
35            update_jacobian_after_steps: solve.ode_update_jacobian_after_steps(),
36            update_rhs_jacobian_after_steps: solve.ode_update_rhs_jacobian_after_steps(),
37            threshold_to_update_jacobian: solve.ode_threshold_to_update_jacobian(),
38            threshold_to_update_rhs_jacobian: solve.ode_threshold_to_update_rhs_jacobian(),
39            min_timestep: solve.ode_min_timestep(),
40            derivative_solve_tolerance: solve.ode_derivative_solve_tolerance(),
41            derivative_solve_max_iter: solve.ode_derivative_solve_max_iter(),
42        }
43    }
44
45    pub(crate) fn apply_to_solve(&self, solve: &mut dyn Solve) {
46        solve.set_ode_max_nonlinear_solver_iterations(self.max_nonlinear_solver_iterations);
47        solve.set_ode_max_error_test_failures(self.max_error_test_failures);
48        solve.set_ode_update_jacobian_after_steps(self.update_jacobian_after_steps);
49        solve.set_ode_update_rhs_jacobian_after_steps(self.update_rhs_jacobian_after_steps);
50        solve.set_ode_threshold_to_update_jacobian(self.threshold_to_update_jacobian);
51        solve.set_ode_threshold_to_update_rhs_jacobian(self.threshold_to_update_rhs_jacobian);
52        solve.set_ode_min_timestep(self.min_timestep);
53        solve.set_ode_derivative_solve_tolerance(self.derivative_solve_tolerance);
54        solve.set_ode_derivative_solve_max_iter(self.derivative_solve_max_iter);
55    }
56}
57
58#[derive(Clone)]
59pub struct OdeSolverOptions {
60    ode: Arc<Mutex<Ode>>,
61}
62impl OdeSolverOptions {
63    pub(crate) fn new(ode: Arc<Mutex<Ode>>) -> Self {
64        Self { ode }
65    }
66    fn guard(&self) -> Result<std::sync::MutexGuard<'_, Ode>, DiffsolRtError> {
67        self.ode.lock().map_err(|_| {
68            DiffsolRtError::from(diffsol::error::DiffsolError::Other(
69                "Failed to acquire lock on Ode object".to_string(),
70            ))
71        })
72    }
73}
74
75impl OdeSolverOptions {
76    pub fn get_max_nonlinear_solver_iterations(&self) -> Result<usize, DiffsolRtError> {
77        Ok(self.guard()?.solve.ode_max_nonlinear_solver_iterations())
78    }
79    pub fn set_max_nonlinear_solver_iterations(&self, value: usize) -> Result<(), DiffsolRtError> {
80        self.guard()?
81            .solve
82            .set_ode_max_nonlinear_solver_iterations(value);
83        Ok(())
84    }
85    pub fn get_max_error_test_failures(&self) -> Result<usize, DiffsolRtError> {
86        Ok(self.guard()?.solve.ode_max_error_test_failures())
87    }
88    pub fn set_max_error_test_failures(&self, value: usize) -> Result<(), DiffsolRtError> {
89        self.guard()?.solve.set_ode_max_error_test_failures(value);
90        Ok(())
91    }
92    pub fn get_update_jacobian_after_steps(&self) -> Result<usize, DiffsolRtError> {
93        Ok(self.guard()?.solve.ode_update_jacobian_after_steps())
94    }
95    pub fn set_update_jacobian_after_steps(&self, value: usize) -> Result<(), DiffsolRtError> {
96        self.guard()?
97            .solve
98            .set_ode_update_jacobian_after_steps(value);
99        Ok(())
100    }
101    pub fn get_update_rhs_jacobian_after_steps(&self) -> Result<usize, DiffsolRtError> {
102        Ok(self.guard()?.solve.ode_update_rhs_jacobian_after_steps())
103    }
104    pub fn set_update_rhs_jacobian_after_steps(&self, value: usize) -> Result<(), DiffsolRtError> {
105        self.guard()?
106            .solve
107            .set_ode_update_rhs_jacobian_after_steps(value);
108        Ok(())
109    }
110    pub fn get_threshold_to_update_jacobian(&self) -> Result<f64, DiffsolRtError> {
111        Ok(self.guard()?.solve.ode_threshold_to_update_jacobian())
112    }
113    pub fn set_threshold_to_update_jacobian(&self, value: f64) -> Result<(), DiffsolRtError> {
114        self.guard()?
115            .solve
116            .set_ode_threshold_to_update_jacobian(value);
117        Ok(())
118    }
119    pub fn get_threshold_to_update_rhs_jacobian(&self) -> Result<f64, DiffsolRtError> {
120        Ok(self.guard()?.solve.ode_threshold_to_update_rhs_jacobian())
121    }
122    pub fn set_threshold_to_update_rhs_jacobian(&self, value: f64) -> Result<(), DiffsolRtError> {
123        self.guard()?
124            .solve
125            .set_ode_threshold_to_update_rhs_jacobian(value);
126        Ok(())
127    }
128    pub fn get_min_timestep(&self) -> Result<f64, DiffsolRtError> {
129        Ok(self.guard()?.solve.ode_min_timestep())
130    }
131    pub fn set_min_timestep(&self, value: f64) -> Result<(), DiffsolRtError> {
132        self.guard()?.solve.set_ode_min_timestep(value);
133        Ok(())
134    }
135    /// The tolerance of the BDF derivative solves (forward sensitivities, step VJPs).
136    pub fn get_derivative_solve_tolerance(&self) -> Result<f64, DiffsolRtError> {
137        Ok(self.guard()?.solve.ode_derivative_solve_tolerance())
138    }
139    pub fn set_derivative_solve_tolerance(&self, value: f64) -> Result<(), DiffsolRtError> {
140        self.guard()?
141            .solve
142            .set_ode_derivative_solve_tolerance(value);
143        Ok(())
144    }
145    /// The iteration limit of the BDF derivative solves.
146    pub fn get_derivative_solve_max_iter(&self) -> Result<usize, DiffsolRtError> {
147        Ok(self.guard()?.solve.ode_derivative_solve_max_iter())
148    }
149    pub fn set_derivative_solve_max_iter(&self, value: usize) -> Result<(), DiffsolRtError> {
150        self.guard()?.solve.set_ode_derivative_solve_max_iter(value);
151        Ok(())
152    }
153}
154
155impl Serialize for OdeSolverOptions {
156    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
157    where
158        S: Serializer,
159    {
160        let guard = self.guard().map_err(serde::ser::Error::custom)?;
161        OdeSolverOptionsSnapshot::from_solve(guard.solve.as_ref()).serialize(serializer)
162    }
163}
164
165#[cfg(all(test, any(feature = "diffsl-cranelift", feature = "diffsl-llvm")))]
166mod tests {
167    use crate::{
168        jit::JitBackendType,
169        linear_solver_type::LinearSolverType,
170        matrix_type::MatrixType,
171        ode::OdeWrapper,
172        ode_solver_type::OdeSolverType,
173        scalar_type::ScalarType,
174        test_support::{available_jit_backends, logistic_diffsl_code},
175    };
176
177    use super::OdeSolverOptions;
178
179    fn make_options(jit_backend: JitBackendType) -> OdeSolverOptions {
180        OdeWrapper::new_jit(
181            logistic_diffsl_code(),
182            jit_backend,
183            ScalarType::F64,
184            MatrixType::NalgebraDense,
185            LinearSolverType::Default,
186            OdeSolverType::Bdf,
187        )
188        .unwrap()
189        .get_options()
190    }
191
192    #[test]
193    fn ode_solver_options_roundtrip_and_serialize() {
194        for jit_backend in available_jit_backends() {
195            let options = make_options(jit_backend);
196            options.set_max_nonlinear_solver_iterations(17).unwrap();
197            options.set_max_error_test_failures(19).unwrap();
198            options.set_update_jacobian_after_steps(23).unwrap();
199            options.set_update_rhs_jacobian_after_steps(29).unwrap();
200            options.set_threshold_to_update_jacobian(1e-3).unwrap();
201            options.set_threshold_to_update_rhs_jacobian(2e-3).unwrap();
202            options.set_min_timestep(1e-4).unwrap();
203            // the derivative solves' settings are their own, with the Newton solve's defaults
204            assert_eq!(options.get_derivative_solve_max_iter().unwrap(), 10);
205            assert_eq!(options.get_derivative_solve_tolerance().unwrap(), 0.2);
206            options.set_derivative_solve_tolerance(1e-6).unwrap();
207            options.set_derivative_solve_max_iter(31).unwrap();
208
209            assert_eq!(options.get_max_nonlinear_solver_iterations().unwrap(), 17);
210            assert_eq!(options.get_max_error_test_failures().unwrap(), 19);
211            assert_eq!(options.get_update_jacobian_after_steps().unwrap(), 23);
212            assert_eq!(options.get_update_rhs_jacobian_after_steps().unwrap(), 29);
213            assert_eq!(options.get_threshold_to_update_jacobian().unwrap(), 1e-3);
214            assert_eq!(
215                options.get_threshold_to_update_rhs_jacobian().unwrap(),
216                2e-3
217            );
218            assert_eq!(options.get_min_timestep().unwrap(), 1e-4);
219            assert_eq!(options.get_derivative_solve_tolerance().unwrap(), 1e-6);
220            assert_eq!(options.get_derivative_solve_max_iter().unwrap(), 31);
221
222            let value = serde_json::to_value(&options).unwrap();
223            assert_eq!(value["max_nonlinear_solver_iterations"], 17);
224            assert_eq!(value["max_error_test_failures"], 19);
225            assert_eq!(value["update_jacobian_after_steps"], 23);
226            assert_eq!(value["update_rhs_jacobian_after_steps"], 29);
227            assert_eq!(value["threshold_to_update_jacobian"], 1e-3);
228            assert_eq!(value["threshold_to_update_rhs_jacobian"], 2e-3);
229            assert_eq!(value["min_timestep"], 1e-4);
230            assert_eq!(value["derivative_solve_tolerance"], 1e-6);
231            assert_eq!(value["derivative_solve_max_iter"], 31);
232        }
233    }
234}