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

diffsol_c/
ode.rs
