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