diffsol_c/
ode.rs

1use std::sync::{Arc, Mutex};
2
3use serde::{de::Error as DeError, Deserialize, Deserializer, Serialize, Serializer};
4
5use crate::jit::JitBackendType;
6use crate::{
7    adjoint_checkpoint::AdjointCheckpointWrapper,
8    error::DiffsolRtError,
9    host_array::{FromHostArray, HostArray},
10    initial_condition_options::{
11        InitialConditionSolverOptions, InitialConditionSolverOptionsSnapshot,
12    },
13    linear_solver_type::LinearSolverType,
14    matrix_type::MatrixType,
15    ode_options::{OdeSolverOptions, OdeSolverOptionsSnapshot},
16    ode_solver_type::OdeSolverType,
17    scalar_type::ScalarType,
18    solution_wrapper::SolutionWrapper,
19    solve::Solve,
20};
21
22pub struct Ode {
23    pub(crate) solve: Box<dyn Solve>,
24    code: String,
25    scalar_type: ScalarType,
26    jit_backend: Option<JitBackendType>,
27    linear_solver: LinearSolverType,
28    ode_solver: OdeSolverType,
29}
30
31#[derive(Clone)]
32pub struct OdeWrapper(Arc<Mutex<Ode>>);
33
34// `OdeWrapper` must be `Send` and `Sync`
35const _: () = {
36    const fn assert_send_sync<T: Send + Sync>() {}
37    assert_send_sync::<OdeWrapper>();
38};
39
40#[derive(Clone, Debug, Serialize, Deserialize)]
41struct OdeWrapperSnapshot {
42    code: String,
43    equation: Vec<u8>,
44    jit_backend: JitBackendType,
45    scalar_type: ScalarType,
46    matrix_type: MatrixType,
47    linear_solver: LinearSolverType,
48    ode_solver: OdeSolverType,
49    rtol: f64,
50    atol: f64,
51    t0: f64,
52    h0: f64,
53    integrate_out: bool,
54    sens_rtol: Option<f64>,
55    sens_atol: Option<f64>,
56    out_rtol: Option<f64>,
57    out_atol: Option<f64>,
58    param_rtol: Option<f64>,
59    param_atol: Option<f64>,
60    ic_options: InitialConditionSolverOptionsSnapshot,
61    ode_options: OdeSolverOptionsSnapshot,
62}
63
64impl OdeWrapper {
65    fn guard(&self) -> Result<std::sync::MutexGuard<'_, Ode>, DiffsolRtError> {
66        self.0.lock().map_err(|_| {
67            DiffsolRtError::from(diffsol::error::DiffsolError::Other(
68                "Failed to acquire lock on ODE solver".to_string(),
69            ))
70        })
71    }
72}
73
74impl OdeWrapper {
75    fn snapshot(&self) -> Result<OdeWrapperSnapshot, DiffsolRtError> {
76        let ode = self.guard()?;
77        let jit_backend = ode.jit_backend.ok_or_else(|| {
78            DiffsolRtError::from(diffsol::error::DiffsolError::Other(
79                "OdeWrapper serialization is only supported for JIT-backed solvers".to_string(),
80            ))
81        })?;
82        Ok(OdeWrapperSnapshot {
83            code: ode.code.clone(),
84            equation: ode.solve.serialized_diffsl()?,
85            jit_backend,
86            scalar_type: ode.scalar_type,
87            matrix_type: ode.solve.matrix_type(),
88            linear_solver: ode.linear_solver,
89            ode_solver: ode.ode_solver,
90            rtol: ode.solve.rtol(),
91            atol: ode.solve.atol(),
92            t0: ode.solve.t0(),
93            h0: ode.solve.h0(),
94            integrate_out: ode.solve.integrate_out(),
95            sens_rtol: ode.solve.sens_rtol(),
96            sens_atol: ode.solve.sens_atol(),
97            out_rtol: ode.solve.out_rtol(),
98            out_atol: ode.solve.out_atol(),
99            param_rtol: ode.solve.param_rtol(),
100            param_atol: ode.solve.param_atol(),
101            ic_options: InitialConditionSolverOptionsSnapshot::from_solve(ode.solve.as_ref()),
102            ode_options: OdeSolverOptionsSnapshot::from_solve(ode.solve.as_ref()),
103        })
104    }
105
106    pub(crate) fn build(
107        code: String,
108        scalar_type: ScalarType,
109        solve: Box<dyn Solve>,
110        jit_backend: Option<JitBackendType>,
111        linear_solver: LinearSolverType,
112        ode_solver: OdeSolverType,
113    ) -> Result<Self, DiffsolRtError> {
114        solve.check(linear_solver)?;
115        Ok(OdeWrapper(Arc::new(Mutex::new(Ode {
116            code,
117            scalar_type,
118            solve,
119            jit_backend,
120            linear_solver,
121            ode_solver,
122        }))))
123    }
124
125    fn from_snapshot(snapshot: OdeWrapperSnapshot) -> Result<Self, DiffsolRtError> {
126        let solve = crate::ode_diffsl::solve_factory_from_serialized_diffsl(
127            snapshot.equation.as_slice(),
128            snapshot.matrix_type,
129            snapshot.scalar_type,
130        )?;
131        let wrapper = Self::build(
132            snapshot.code,
133            snapshot.scalar_type,
134            solve,
135            Some(snapshot.jit_backend),
136            snapshot.linear_solver,
137            snapshot.ode_solver,
138        )?;
139        {
140            let mut ode = wrapper.guard()?;
141            ode.solve.set_rtol(snapshot.rtol);
142            ode.solve.set_atol(snapshot.atol);
143            ode.solve.set_t0(snapshot.t0);
144            ode.solve.set_h0(snapshot.h0);
145            ode.solve.set_integrate_out(snapshot.integrate_out);
146            ode.solve.set_sens_rtol(snapshot.sens_rtol);
147            ode.solve.set_sens_atol(snapshot.sens_atol);
148            ode.solve.set_out_rtol(snapshot.out_rtol);
149            ode.solve.set_out_atol(snapshot.out_atol);
150            ode.solve.set_param_rtol(snapshot.param_rtol);
151            ode.solve.set_param_atol(snapshot.param_atol);
152            snapshot.ic_options.apply_to_solve(ode.solve.as_mut());
153            snapshot.ode_options.apply_to_solve(ode.solve.as_mut());
154        }
155        Ok(wrapper)
156    }
157
158    /// Matrix type used in the ODE solver. This is fixed after construction.
159    pub fn get_matrix_type(&self) -> Result<MatrixType, DiffsolRtError> {
160        Ok(self.guard()?.solve.matrix_type())
161    }
162
163    pub fn get_nstates(&self) -> Result<usize, DiffsolRtError> {
164        Ok(self.guard()?.solve.nstates())
165    }
166
167    pub fn get_nparams(&self) -> Result<usize, DiffsolRtError> {
168        Ok(self.guard()?.solve.nparams())
169    }
170
171    pub fn get_nout(&self) -> Result<usize, DiffsolRtError> {
172        Ok(self.guard()?.solve.nout())
173    }
174
175    pub fn has_stop(&self) -> Result<bool, DiffsolRtError> {
176        Ok(self.guard()?.solve.has_stop())
177    }
178
179    /// Ode solver method, default Bdf (backward differentiation formula).
180    pub fn get_ode_solver(&self) -> Result<OdeSolverType, DiffsolRtError> {
181        Ok(self.guard()?.ode_solver)
182    }
183
184    pub fn set_ode_solver(&self, value: OdeSolverType) -> Result<(), DiffsolRtError> {
185        self.guard()?.ode_solver = value;
186        Ok(())
187    }
188
189    /// Linear solver type used in the ODE solver. Set to default to use the
190    /// solver's default choice, which is typically an LU solver.
191    pub fn get_linear_solver(&self) -> Result<LinearSolverType, DiffsolRtError> {
192        Ok(self.guard()?.linear_solver)
193    }
194
195    pub fn set_linear_solver(&self, value: LinearSolverType) -> Result<(), DiffsolRtError> {
196        self.guard()?.solve.check(value)?;
197        self.guard()?.linear_solver = value;
198        Ok(())
199    }
200
201    /// Relative tolerance for the solver, default 1e-6. Governs the error relative to the solution size.
202    pub fn get_rtol(&self) -> Result<f64, DiffsolRtError> {
203        Ok(self.guard()?.solve.rtol())
204    }
205
206    pub fn set_rtol(&self, value: f64) -> Result<(), DiffsolRtError> {
207        self.guard()?.solve.set_rtol(value);
208        Ok(())
209    }
210
211    /// Absolute tolerance for the solver, default 1e-6. Governs the error as the solution goes to zero.
212    pub fn get_atol(&self) -> Result<f64, DiffsolRtError> {
213        Ok(self.guard()?.solve.atol())
214    }
215
216    pub fn set_atol(&self, value: f64) -> Result<(), DiffsolRtError> {
217        self.guard()?.solve.set_atol(value);
218        Ok(())
219    }
220
221    /// Initial time for the ODE solve, default 0.0.
222    pub fn get_t0(&self) -> Result<f64, DiffsolRtError> {
223        Ok(self.guard()?.solve.t0())
224    }
225
226    pub fn set_t0(&self, value: f64) -> Result<(), DiffsolRtError> {
227        self.guard()?.solve.set_t0(value);
228        Ok(())
229    }
230
231    /// Initial step size for the ODE solver, default 1.0.
232    pub fn get_h0(&self) -> Result<f64, DiffsolRtError> {
233        Ok(self.guard()?.solve.h0())
234    }
235
236    pub fn set_h0(&self, value: f64) -> Result<(), DiffsolRtError> {
237        self.guard()?.solve.set_h0(value);
238        Ok(())
239    }
240
241    /// Whether to integrate output equations alongside state equations.
242    pub fn get_integrate_out(&self) -> Result<bool, DiffsolRtError> {
243        Ok(self.guard()?.solve.integrate_out())
244    }
245
246    pub fn set_integrate_out(&self, value: bool) -> Result<(), DiffsolRtError> {
247        self.guard()?.solve.set_integrate_out(value);
248        Ok(())
249    }
250
251    /// Relative tolerance for forward sensitivity or adjoint equations.
252    pub fn get_sens_rtol(&self) -> Result<Option<f64>, DiffsolRtError> {
253        Ok(self.guard()?.solve.sens_rtol())
254    }
255
256    pub fn set_sens_rtol(&self, value: Option<f64>) -> Result<(), DiffsolRtError> {
257        self.guard()?.solve.set_sens_rtol(value);
258        Ok(())
259    }
260
261    /// Absolute tolerance for forward sensitivity or adjoint equations.
262    pub fn get_sens_atol(&self) -> Result<Option<f64>, DiffsolRtError> {
263        Ok(self.guard()?.solve.sens_atol())
264    }
265
266    pub fn set_sens_atol(&self, value: Option<f64>) -> Result<(), DiffsolRtError> {
267        self.guard()?.solve.set_sens_atol(value);
268        Ok(())
269    }
270
271    /// Relative tolerance for integrated output equations.
272    pub fn get_out_rtol(&self) -> Result<Option<f64>, DiffsolRtError> {
273        Ok(self.guard()?.solve.out_rtol())
274    }
275
276    pub fn set_out_rtol(&self, value: Option<f64>) -> Result<(), DiffsolRtError> {
277        self.guard()?.solve.set_out_rtol(value);
278        Ok(())
279    }
280
281    /// Absolute tolerance for integrated output equations.
282    pub fn get_out_atol(&self) -> Result<Option<f64>, DiffsolRtError> {
283        Ok(self.guard()?.solve.out_atol())
284    }
285
286    pub fn set_out_atol(&self, value: Option<f64>) -> Result<(), DiffsolRtError> {
287        self.guard()?.solve.set_out_atol(value);
288        Ok(())
289    }
290
291    /// Relative tolerance for adjoint parameter gradient equations.
292    pub fn get_param_rtol(&self) -> Result<Option<f64>, DiffsolRtError> {
293        Ok(self.guard()?.solve.param_rtol())
294    }
295
296    pub fn set_param_rtol(&self, value: Option<f64>) -> Result<(), DiffsolRtError> {
297        self.guard()?.solve.set_param_rtol(value);
298        Ok(())
299    }
300
301    /// Absolute tolerance for adjoint parameter gradient equations.
302    pub fn get_param_atol(&self) -> Result<Option<f64>, DiffsolRtError> {
303        Ok(self.guard()?.solve.param_atol())
304    }
305
306    pub fn set_param_atol(&self, value: Option<f64>) -> Result<(), DiffsolRtError> {
307        self.guard()?.solve.set_param_atol(value);
308        Ok(())
309    }
310
311    pub fn get_code(&self) -> Result<String, DiffsolRtError> {
312        Ok(self.guard()?.code.clone())
313    }
314
315    pub fn get_scalar_type(&self) -> Result<ScalarType, DiffsolRtError> {
316        Ok(self.guard()?.scalar_type)
317    }
318
319    pub fn get_jit_backend(&self) -> Result<Option<JitBackendType>, DiffsolRtError> {
320        Ok(self.guard()?.jit_backend)
321    }
322
323    pub fn get_ic_options(&self) -> InitialConditionSolverOptions {
324        InitialConditionSolverOptions::new(self.0.clone())
325    }
326
327    pub fn get_options(&self) -> OdeSolverOptions {
328        OdeSolverOptions::new(self.0.clone())
329    }
330
331    /// Get the initial condition vector y0 as a 1D numpy array.
332    pub fn y0(&self, params: HostArray) -> Result<HostArray, DiffsolRtError> {
333        let mut self_guard = self.guard()?;
334        self_guard.solve.y0(params.as_slice()?)
335    }
336
337    /// evaluate the right-hand side function at time `t` and state `y`.
338    pub fn rhs(
339        &self,
340        params: HostArray,
341        t: f64,
342        y: HostArray,
343    ) -> Result<HostArray, DiffsolRtError> {
344        let mut self_guard = self.guard()?;
345        self_guard.solve.rhs(params.as_slice()?, t, y.as_slice()?)
346    }
347
348    /// evaluate the right-hand side Jacobian-vector product `Jv`` at time `t` and state `y`.
349    pub fn rhs_jac_mul(
350        &self,
351        params: HostArray,
352        t: f64,
353        y: HostArray,
354        v: HostArray,
355    ) -> Result<HostArray, DiffsolRtError> {
356        let mut self_guard = self.guard()?;
357        self_guard
358            .solve
359            .rhs_jac_mul(params.as_slice()?, t, y.as_slice()?, v.as_slice()?)
360    }
361
362    /// Using the provided state, solve the problem up to time `final_time`.
363    ///
364    /// The number of params must match the expected params in the diffsl code.
365    /// If specified, the config can be used to override the solver method
366    /// (Bdf by default) and SolverType (Lu by default) along with other solver
367    /// params like `rtol`.
368    ///
369    /// :param params: 1D array of solver parameters
370    /// :type params: numpy.ndarray
371    /// :param final_time: end time of solver
372    /// :type final_time: float
373    /// :return: `(ys, ts)` tuple where `ys` is a 2D array of values at times
374    ///     `ts` chosen by the solver
375    /// :rtype: Tuple[numpy.ndarray, numpy.ndarray]
376    ///
377    /// Example:
378    ///     >>> print(ode.solve(np.array([]), 0.5))
379    #[allow(clippy::type_complexity)]
380    pub fn solve(
381        &self,
382        params: HostArray,
383        final_time: f64,
384    ) -> Result<SolutionWrapper, DiffsolRtError> {
385        let mut self_guard = self.guard()?;
386        let params = params.as_slice()?;
387        let linear_solver = self_guard.linear_solver;
388        let method = self_guard.ode_solver;
389        let solution = self_guard
390            .solve
391            .solve(method, linear_solver, params, final_time)?;
392        Ok(SolutionWrapper::new(solution))
393    }
394
395    /// Using the provided state, solve the problem up to time
396    /// `t_eval[t_eval.len()-1]`. Returns 2D array of solution values at
397    /// timepoints given by `t_eval`.
398    ///
399    /// The number of params must match the expected params in the diffsl code.
400    /// The config may be optionally specified to override solver settings.
401    ///
402    /// :param params: 1D array of solver parameters
403    /// :type params: numpy.ndarray
404    /// :param t_eval: 1D array of solver times
405    /// :type params: numpy.ndarray
406    /// :return: 2D array of values at times `t_eval`
407    /// :rtype: numpy.ndarray
408    pub fn solve_dense(
409        &self,
410        params: HostArray,
411        t_eval: HostArray,
412    ) -> Result<SolutionWrapper, DiffsolRtError> {
413        let mut self_guard = self.guard()?;
414        let params = params.as_slice()?;
415        let t_eval = t_eval.as_slice()?;
416        let linear_solver = self_guard.linear_solver;
417        let method = self_guard.ode_solver;
418        let solution = self_guard
419            .solve
420            .solve_dense(method, linear_solver, params, t_eval)?;
421        Ok(SolutionWrapper::new(solution))
422    }
423
424    /// Using the provided state, solve the problem up to time `t_eval[t_eval.len()-1]`.
425    /// Returns 2D array of solution values at timepoints given by `t_eval`.
426    /// Also returns a list of 2D arrays of sensitivities at the same timepoints
427    /// as the solution.
428    /// The number of params must match the expected params in the diffsl code.
429    /// The config may be optionally specified to override solver settings.
430    /// :param params: 1D array of solver parameters
431    /// :type params: numpy.ndarray
432    /// :param t_eval: 1D array of solver times
433    /// :type params: numpy.ndarray
434    /// :return: 2D array of values at times `t_eval` and a list of 2D arrays of sensitivities at the same timepoints
435    /// :rtype: (numpy.ndarray, List[numpy.ndarray])
436    #[allow(clippy::type_complexity)]
437    pub fn solve_fwd_sens(
438        &self,
439        params: HostArray,
440        t_eval: HostArray,
441    ) -> Result<SolutionWrapper, DiffsolRtError> {
442        let mut self_guard = self.guard()?;
443        let params = params.as_slice()?;
444        let t_eval = t_eval.as_slice()?;
445        let linear_solver = self_guard.linear_solver;
446        let method = self_guard.ode_solver;
447        let solution = self_guard
448            .solve
449            .solve_fwd_sens(method, linear_solver, params, t_eval)?;
450        Ok(SolutionWrapper::new(solution))
451    }
452
453    /// Solve the continuous adjoint problem for the integral of the model output
454    /// from the initial time to `final_time`.
455    ///
456    /// Returns `(integral, gradient)`, where `integral` is a vector of length
457    /// `nout` and `gradient` is an `(nparams, nout)` matrix.
458    pub fn solve_continuous_adjoint(
459        &self,
460        params: HostArray,
461        final_time: f64,
462    ) -> Result<(HostArray, HostArray), DiffsolRtError> {
463        let mut self_guard = self.guard()?;
464        let linear_solver = self_guard.linear_solver;
465        let ode_solver = self_guard.ode_solver;
466        self_guard.solve.solve_continuous_adjoint(
467            ode_solver,
468            linear_solver,
469            params.as_slice()?,
470            final_time,
471        )
472    }
473
474    /// Solve the forward problem at `t_eval` and retain checkpoint data for a
475    /// later discrete adjoint backward pass.
476    pub fn solve_adjoint_fwd(
477        &self,
478        params: HostArray,
479        t_eval: HostArray,
480    ) -> Result<(SolutionWrapper, AdjointCheckpointWrapper), DiffsolRtError> {
481        let mut self_guard = self.guard()?;
482        let params = params.as_slice()?;
483        let t_eval = t_eval.as_slice()?;
484        let linear_solver = self_guard.linear_solver;
485        let method = self_guard.ode_solver;
486        let (solution, checkpoint) =
487            self_guard
488                .solve
489                .solve_adjoint_fwd(method, linear_solver, params, t_eval)?;
490        Ok((SolutionWrapper::new(solution), checkpoint))
491    }
492
493    /// Solve the discrete adjoint backward pass using a prior forward adjoint
494    /// checkpoint and the gradient of a scalar objective with respect to model
495    /// outputs at each saved evaluation time.
496    ///
497    /// Returns an `(nparams, 1)` gradient matrix.
498    pub fn solve_adjoint_bkwd(
499        &self,
500        solution: &SolutionWrapper,
501        checkpoint: &AdjointCheckpointWrapper,
502        dgdu_eval: HostArray,
503    ) -> Result<HostArray, DiffsolRtError> {
504        let t_eval_host = solution.get_ts()?;
505        let t_eval = Vec::<f64>::from_host_array(t_eval_host)?;
506        let mut self_guard = self.guard()?;
507        let linear_solver = self_guard.linear_solver;
508        let method = self_guard.ode_solver;
509        self_guard
510            .solve
511            .solve_adjoint_bkwd(method, linear_solver, checkpoint, &t_eval, dgdu_eval)
512    }
513}
514
515impl Serialize for OdeWrapper {
516    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
517    where
518        S: Serializer,
519    {
520        self.snapshot()
521            .map_err(serde::ser::Error::custom)?
522            .serialize(serializer)
523    }
524}
525
526impl<'de> Deserialize<'de> for OdeWrapper {
527    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
528    where
529        D: Deserializer<'de>,
530    {
531        let snapshot = OdeWrapperSnapshot::deserialize(deserializer)?;
532        Self::from_snapshot(snapshot).map_err(DeError::custom)
533    }
534}
535
536#[cfg(all(test, feature = "diffsl-external-f64"))]
537mod tests {
538    use super::*;
539    use crate::host_array::FromHostArray;
540    use crate::linear_solver_type::LinearSolverType;
541    use crate::scalar_type::ScalarType;
542    use crate::test_support::{
543        assert_close, assert_solution_tail, logistic_state, logistic_state_dr, mass_state_deps,
544        rhs_input_deps, rhs_state_deps, vector_host, ASSERT_TOL, LOGISTIC_X0,
545    };
546
547    fn all_ode_solvers() -> [OdeSolverType; 4] {
548        [
549            OdeSolverType::Bdf,
550            OdeSolverType::Esdirk34,
551            OdeSolverType::TrBdf2,
552            OdeSolverType::Tsit45,
553        ]
554    }
555
556    fn make_ode(matrix_type: MatrixType, ode_solver: OdeSolverType) -> OdeWrapper {
557        OdeWrapper::new_external(
558            rhs_state_deps(),
559            rhs_input_deps(),
560            mass_state_deps(),
561            ScalarType::F64,
562            matrix_type,
563            LinearSolverType::Default,
564            ode_solver,
565        )
566        .unwrap()
567    }
568
569    fn assert_runtime_dispatch(matrix_type: MatrixType) {
570        let ode = make_ode(matrix_type, OdeSolverType::Bdf);
571        assert_eq!(ode.get_matrix_type().unwrap(), matrix_type);
572        assert_eq!(ode.get_nstates().unwrap(), 1);
573        assert_eq!(ode.get_nparams().unwrap(), 1);
574        assert_eq!(ode.get_nout().unwrap(), 1);
575        assert!(ode.has_stop().unwrap());
576
577        let y0 = ode.y0(vector_host(&[2.0])).unwrap();
578        assert_eq!(Vec::<f64>::from_host_array(y0).unwrap(), vec![LOGISTIC_X0]);
579
580        let rhs = ode
581            .rhs(vector_host(&[2.0]), 0.0, vector_host(&[0.25]))
582            .unwrap();
583        assert_close(
584            Vec::<f64>::from_host_array(rhs).unwrap()[0],
585            0.375,
586            ASSERT_TOL,
587            "rhs(0.25)",
588        );
589
590        let rhs_jac_mul = ode
591            .rhs_jac_mul(
592                vector_host(&[2.0]),
593                0.0,
594                vector_host(&[0.25]),
595                vector_host(&[3.0]),
596            )
597            .unwrap();
598        assert_close(
599            Vec::<f64>::from_host_array(rhs_jac_mul).unwrap()[0],
600            3.0,
601            ASSERT_TOL,
602            "rhs_jac_mul(0.25, 3.0)",
603        );
604    }
605
606    fn assert_solver_dense_solution(matrix_type: MatrixType, ode_solver: OdeSolverType) {
607        let ode = make_ode(matrix_type, ode_solver);
608        ode.set_rtol(1e-8).unwrap();
609        ode.set_atol(1e-8).unwrap();
610
611        let t_eval = [0.25, 0.5, 1.0];
612        let solution = ode
613            .solve_dense(vector_host(&[2.0]), vector_host(&t_eval))
614            .unwrap();
615
616        assert_solution_tail(&solution, &t_eval, LOGISTIC_X0, 2.0, 5e-4);
617    }
618
619    fn hybrid_root_time() -> f64 {
620        0.5 * 9.0_f64.ln()
621    }
622
623    fn assert_hybrid_solution_applies_reset_after_root(ode_solver: OdeSolverType) {
624        let ode = make_ode(MatrixType::NalgebraDense, ode_solver);
625        ode.set_rtol(1e-8).unwrap();
626        ode.set_atol(1e-8).unwrap();
627
628        let final_time = 2.0;
629        let solution = ode.solve(vector_host(&[2.0]), final_time).unwrap();
630        let ys = solution.get_ys().unwrap();
631        let ys = ys.as_array::<f64>().unwrap();
632        let ts = Vec::<f64>::from_host_array(solution.get_ts().unwrap()).unwrap();
633        let root_time = hybrid_root_time();
634
635        assert_eq!(ys.nrows(), 1);
636        assert_eq!(ys.ncols(), ts.len());
637        assert!(!ts.is_empty(), "expected hybrid solve to produce output");
638        assert_close(
639            *ts.last().unwrap(),
640            final_time,
641            ASSERT_TOL,
642            "hybrid final time",
643        );
644        assert_close(ys[(0, ys.ncols() - 1)], 1.0, 5e-4, "hybrid final value");
645        assert!(
646            ts.iter().any(|&t| t < root_time),
647            "expected pre-root samples"
648        );
649        assert!(
650            ts.iter().any(|&t| t > root_time),
651            "expected post-root samples after reset"
652        );
653    }
654
655    fn assert_hybrid_dense_solution_continues_after_reset(ode_solver: OdeSolverType) {
656        let ode = make_ode(MatrixType::NalgebraDense, ode_solver);
657        ode.set_rtol(1e-8).unwrap();
658        ode.set_atol(1e-8).unwrap();
659
660        let t_eval = [0.5, 1.0, 1.25, 1.5, 2.0];
661        let solution = ode
662            .solve_dense(vector_host(&[2.0]), vector_host(&t_eval))
663            .unwrap();
664        let ys = solution.get_ys().unwrap();
665        let ys = ys.as_array::<f64>().unwrap();
666
667        assert_eq!(ys.nrows(), 1);
668        assert_eq!(ys.ncols(), t_eval.len());
669        assert_close(
670            ys[(0, 0)],
671            logistic_state(LOGISTIC_X0, 2.0, t_eval[0]),
672            5e-4,
673            "hybrid dense pre-root value",
674        );
675        assert_close(
676            ys[(0, 1)],
677            logistic_state(LOGISTIC_X0, 2.0, t_eval[1]),
678            5e-4,
679            "hybrid dense near-root value",
680        );
681        for col in 2..t_eval.len() {
682            assert_close(ys[(0, col)], 1.0, 5e-4, "hybrid dense post-root value");
683        }
684    }
685
686    fn assert_hybrid_forward_sensitivities_complete_across_reset(ode_solver: OdeSolverType) {
687        let ode = make_ode(MatrixType::NalgebraDense, ode_solver);
688        ode.set_rtol(1e-8).unwrap();
689        ode.set_atol(1e-8).unwrap();
690
691        let t_eval = [0.5, 1.0, 1.25, 1.5, 2.0];
692        let solution = ode
693            .solve_fwd_sens(vector_host(&[2.0]), vector_host(&t_eval))
694            .unwrap();
695        let ys = solution.get_ys().unwrap();
696        let ys = ys.as_array::<f64>().unwrap();
697        let sens = solution.get_sens().unwrap();
698
699        assert_eq!(ys.nrows(), 1);
700        assert_eq!(ys.ncols(), t_eval.len());
701        assert_eq!(sens.len(), 1);
702        let sens_values = sens[0].as_array::<f64>().unwrap();
703        assert_eq!(sens_values.nrows(), 1);
704        assert_eq!(sens_values.ncols(), t_eval.len());
705        assert_close(
706            ys[(0, 0)],
707            logistic_state(LOGISTIC_X0, 2.0, t_eval[0]),
708            5e-4,
709            "hybrid sens pre-root value",
710        );
711        for col in 2..t_eval.len() {
712            assert_close(ys[(0, col)], 1.0, 5e-4, "hybrid sens post-root value");
713            assert!(
714                sens_values[(0, col)].is_finite(),
715                "expected finite post-root sensitivity at column {col}"
716            );
717        }
718    }
719
720    #[test]
721    fn runtime_dispatch_matches_requested_matrix_type() {
722        for matrix_type in [
723            MatrixType::NalgebraDense,
724            MatrixType::FaerDense,
725            MatrixType::FaerSparse,
726        ] {
727            assert_runtime_dispatch(matrix_type);
728        }
729    }
730
731    #[test]
732    fn bdf_dense_solution_matches_logistic_solution() {
733        let ode = make_ode(MatrixType::NalgebraDense, OdeSolverType::Bdf);
734        ode.set_rtol(1e-8).unwrap();
735        ode.set_atol(1e-8).unwrap();
736
737        let t_eval = [0.25, 0.5, 1.0];
738        let solution = ode
739            .solve_dense(vector_host(&[2.0]), vector_host(&t_eval))
740            .unwrap();
741
742        assert_solution_tail(&solution, &t_eval, LOGISTIC_X0, 2.0, 5e-4);
743    }
744
745    #[test]
746    fn esdirk34_dense_solution_matches_logistic_solution() {
747        assert_solver_dense_solution(MatrixType::FaerDense, OdeSolverType::Esdirk34);
748    }
749
750    #[test]
751    fn tr_bdf2_sparse_solution_matches_logistic_solution() {
752        assert_solver_dense_solution(MatrixType::FaerSparse, OdeSolverType::TrBdf2);
753    }
754
755    #[test]
756    fn tsit45_dense_solution_matches_logistic_solution() {
757        assert_solver_dense_solution(MatrixType::NalgebraDense, OdeSolverType::Tsit45);
758    }
759
760    #[test]
761    fn bdf_forward_sensitivities_match_logistic_derivative() {
762        let ode = make_ode(MatrixType::NalgebraDense, OdeSolverType::Bdf);
763        ode.set_rtol(1e-8).unwrap();
764        ode.set_atol(1e-8).unwrap();
765        ode.set_sens_rtol(Some(1e-6)).unwrap();
766        ode.set_sens_atol(Some(1e-6)).unwrap();
767
768        let t_eval = [0.25, 0.5, 1.0];
769        let solution = ode
770            .solve_fwd_sens(vector_host(&[2.0]), vector_host(&t_eval))
771            .unwrap();
772
773        assert_solution_tail(&solution, &t_eval, LOGISTIC_X0, 2.0, 5e-4);
774        let sens = solution.get_sens().unwrap();
775        assert_eq!(sens.len(), 1);
776        let sens_values = sens[0].as_array::<f64>().unwrap();
777        assert_eq!(sens_values.nrows(), 1);
778        assert_eq!(sens_values.ncols(), t_eval.len());
779        for (i, &t) in t_eval.iter().enumerate() {
780            assert_close(
781                sens_values[(0, i)],
782                logistic_state_dr(LOGISTIC_X0, 2.0, t),
783                ASSERT_TOL,
784                &format!("sensitivity[{i}]"),
785            );
786        }
787    }
788
789    #[test]
790    fn hybrid_solution_applies_reset_after_root_for_all_solvers() {
791        for ode_solver in all_ode_solvers() {
792            assert_hybrid_solution_applies_reset_after_root(ode_solver);
793        }
794    }
795
796    #[test]
797    fn hybrid_dense_solution_continues_after_reset_for_all_solvers() {
798        for ode_solver in all_ode_solvers() {
799            assert_hybrid_dense_solution_continues_after_reset(ode_solver);
800        }
801    }
802
803    #[test]
804    fn hybrid_forward_sensitivities_complete_across_reset_for_all_solvers() {
805        for ode_solver in all_ode_solvers() {
806            assert_hybrid_forward_sensitivities_complete_across_reset(ode_solver);
807        }
808    }
809}
810
811#[cfg(all(test, feature = "diffsl-external-dynamic"))]
812mod dynamic_tests {
813    use crate::host_array::FromHostArray;
814    use crate::linear_solver_type::LinearSolverType;
815    use crate::scalar_type::ScalarType;
816    use crate::test_support::{
817        assert_close, assert_solution_tail, external_dynamic_fixture_path, mass_state_deps,
818        rhs_input_deps, rhs_state_deps, vector_host, ASSERT_TOL, LOGISTIC_X0,
819    };
820
821    use super::*;
822
823    fn make_ode(matrix_type: MatrixType, ode_solver: OdeSolverType) -> OdeWrapper {
824        OdeWrapper::new_external_dynamic(
825            external_dynamic_fixture_path(),
826            rhs_state_deps(),
827            rhs_input_deps(),
828            mass_state_deps(),
829            ScalarType::F64,
830            matrix_type,
831            LinearSolverType::Default,
832            ode_solver,
833        )
834        .unwrap()
835    }
836
837    #[test]
838    fn runtime_dispatch_matches_requested_matrix_type() {
839        for matrix_type in [
840            MatrixType::NalgebraDense,
841            MatrixType::FaerDense,
842            MatrixType::FaerSparse,
843        ] {
844            let ode = make_ode(matrix_type, OdeSolverType::Bdf);
845            assert_eq!(ode.get_matrix_type().unwrap(), matrix_type);
846            assert_eq!(ode.get_code().unwrap(), "");
847            assert_eq!(ode.get_jit_backend().unwrap(), None);
848            assert_eq!(ode.get_nstates().unwrap(), 1);
849            assert_eq!(ode.get_nparams().unwrap(), 1);
850            assert_eq!(ode.get_nout().unwrap(), 1);
851            assert!(ode.has_stop().unwrap());
852
853            let y0 = ode.y0(vector_host(&[2.0])).unwrap();
854            assert_eq!(Vec::<f64>::from_host_array(y0).unwrap(), vec![LOGISTIC_X0]);
855
856            let rhs = ode
857                .rhs(vector_host(&[2.0]), 0.0, vector_host(&[0.25]))
858                .unwrap();
859            assert_close(
860                Vec::<f64>::from_host_array(rhs).unwrap()[0],
861                0.375,
862                ASSERT_TOL,
863                "rhs(0.25)",
864            );
865
866            let rhs_jac_mul = ode
867                .rhs_jac_mul(
868                    vector_host(&[2.0]),
869                    0.0,
870                    vector_host(&[0.25]),
871                    vector_host(&[3.0]),
872                )
873                .unwrap();
874            assert_close(
875                Vec::<f64>::from_host_array(rhs_jac_mul).unwrap()[0],
876                3.0,
877                ASSERT_TOL,
878                "rhs_jac_mul(0.25, 3.0)",
879            );
880        }
881    }
882
883    #[test]
884    fn dense_solution_matches_logistic_solution() {
885        let ode = make_ode(MatrixType::NalgebraDense, OdeSolverType::Bdf);
886        ode.set_rtol(1e-8).unwrap();
887        ode.set_atol(1e-8).unwrap();
888
889        let t_eval = [0.25, 0.5, 1.0];
890        let solution = ode
891            .solve_dense(vector_host(&[2.0]), vector_host(&t_eval))
892            .unwrap();
893
894        assert_solution_tail(&solution, &t_eval, LOGISTIC_X0, 2.0, 5e-4);
895    }
896
897    #[test]
898    fn non_jit_serialization_is_rejected() {
899        let ode = make_ode(MatrixType::NalgebraDense, OdeSolverType::Bdf);
900        let err = serde_json::to_string(&ode).unwrap_err().to_string();
901        assert!(err.contains("JIT-backed"));
902    }
903}
diffsol_c/ode.rs

diffsol_c/
ode.rs
