diffsol_c/
ode_solver_type.rs

1// Solver method Python enum. This is used to select the overarching solver
2// stragegy like bdf or esdirk34 in diffsol.
3
4use diffsol::error::DiffsolError;
5use diffsol::ode_equations::{OdeEquationsImplicit, OdeEquationsImplicitSens};
6use diffsol::{
7    matrix::MatrixRef, Context, DefaultDenseMatrix, DenseMatrix, LinearSolver, Matrix,
8    OdeSolverProblem, OdeSolverState, Vector, VectorRef, VectorView,
9};
10use diffsol::{
11    ode_solver_error, AdjointOdeSolverMethod, CheckpointingPath, DefaultSolver, OdeSolverMethod,
12    OdeSolverStopReason, SensitivitiesOdeSolverMethod, Solution,
13};
14use schemars::JsonSchema;
15use serde::{Deserialize, Serialize};
16
17use crate::adjoint_checkpoint::{AdjointCheckpoint, AdjointCheckpointData};
18use crate::ode_solver_tag::{BdfTag, Esdirk34Tag, OdeSolverMethodTag, TrBdf2Tag, Tsit45Tag};
19use crate::scalar_type::Scalar;
20use crate::{
21    linear_solver_type::LinearSolverType,
22    valid_linear_solver::{KluValidator, LuValidator},
23};
24
25/// Enumerates the possible ODE solver methods for diffsol. See the solver descriptions in the diffsol documentation (<https://github.com/martinjrobins/diffsol>) for more details.
26///
27/// :attr bdf: Backward Differentiation Formula (BDF) method for stiff ODEs and singular mass matrices
28/// :attr esdirk34: Explicit Singly Diagonally Implicit Runge-Kutta (ESDIRK) method for moderately stiff ODEs and singular mass matrices.
29/// :attr tr_bdf2: Trapezoidal Backward Differentiation Formula of order 2 (TR-BDF2) method for moderately stiff ODEs and singular mass matrices.
30/// :attr tsit45: Tsitouras 4/5th order Explicit Runge-Kutta (TSIT45) method for non-stiff ODEs. This is an explicit method, it cannot handle singular mass matrices and does not require a linear solver.
31#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
32#[serde(rename_all = "snake_case")]
33pub enum OdeSolverType {
34    Bdf,
35    Esdirk34,
36    TrBdf2,
37    Tsit45,
38}
39
40fn solve_with_tag<M, Eqn, LS, Tag>(
41    problem: &mut OdeSolverProblem<Eqn>,
42    mut soln: Solution<M::V>,
43) -> Result<Solution<M::V>, DiffsolError>
44where
45    M: Matrix<T: Scalar> + DefaultSolver,
46    M::V: DefaultDenseMatrix,
47    LS: LinearSolver<M>,
48    Tag: OdeSolverMethodTag<M, Eqn>,
49    Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C>,
50    for<'b> &'b M::V: VectorRef<M::V>,
51    for<'b> &'b M: MatrixRef<M>,
52{
53    let mut solver = Tag::solver::<LS>(problem)?;
54    while !soln.is_complete() {
55        solver = solver.solve_soln(&mut soln)?;
56        let root_idx = match soln.stop_reason {
57            Some(OdeSolverStopReason::RootFound(_, root_idx)) if !soln.is_complete() => root_idx,
58            _ => continue,
59        };
60        if problem.eqn.reset().is_none() {
61            soln.truncate(problem, solver.state())?;
62            return Ok(soln);
63        }
64        let mut state = solver.into_state();
65        problem.eqn.set_model_index(root_idx);
66        state.as_mut().apply_reset_with_mass::<M::LS, _>(problem)?;
67        solver = Tag::solver_with_state::<LS>(problem, state)?;
68    }
69    Ok(soln)
70}
71
72fn solve_fwd_sens_with_tag<M, Eqn, LS, Tag>(
73    problem: &mut OdeSolverProblem<Eqn>,
74    t_eval: &[M::T],
75) -> Result<Solution<M::V>, DiffsolError>
76where
77    M: Matrix<T: Scalar> + DefaultSolver,
78    M::V: DefaultDenseMatrix,
79    LS: LinearSolver<M>,
80    Tag: OdeSolverMethodTag<M, Eqn>,
81    Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>,
82    for<'b> &'b M::V: VectorRef<M::V>,
83    for<'b> &'b M: MatrixRef<M>,
84{
85    let mut soln = Solution::new_dense(t_eval.to_vec())?;
86    let mut solver = Tag::solver_sens::<LS>(problem)?;
87    while !soln.is_complete() {
88        solver = solver.solve_soln_sensitivities(&mut soln)?;
89        let root_idx = match soln.stop_reason {
90            Some(OdeSolverStopReason::RootFound(_, root_idx)) if !soln.is_complete() => root_idx,
91            _ => continue,
92        };
93        if problem.eqn.reset().is_none() {
94            soln.truncate_sens(problem, solver.state())?;
95            return Ok(soln);
96        }
97        let mut state = solver.into_state();
98        problem.eqn.set_model_index(root_idx);
99        state
100            .as_mut()
101            .apply_reset_with_sens_mass::<M::LS, _>(problem, root_idx)?;
102        solver = Tag::solver_sens_with_state::<LS>(problem, state)?;
103    }
104    Ok(soln)
105}
106
107#[allow(clippy::type_complexity)]
108fn solve_with_checkpointing_with_tag<M, Eqn, LS, Tag>(
109    problem: &mut OdeSolverProblem<Eqn>,
110    mut soln: Solution<M::V>,
111) -> Result<(Solution<M::V>, CheckpointingPath<Eqn, Tag::State>), DiffsolError>
112where
113    M: Matrix<T: Scalar> + DefaultSolver,
114    M::V: DefaultDenseMatrix,
115    LS: LinearSolver<M>,
116    Tag: OdeSolverMethodTag<M, Eqn>,
117    Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C>,
118    for<'b> &'b M::V: VectorRef<M::V>,
119    for<'b> &'b M: MatrixRef<M>,
120{
121    let mut solver = Tag::solver::<LS>(problem)?;
122    let mut checkpointing = Vec::new();
123    while !soln.is_complete() {
124        solver = solver.solve_soln_with_checkpointing(&mut soln, &mut checkpointing, None)?;
125        let root_idx = match soln.stop_reason {
126            Some(OdeSolverStopReason::RootFound(_, root_idx)) if !soln.is_complete() => root_idx,
127            _ => continue,
128        };
129        if problem.eqn.reset().is_none() {
130            soln.truncate(problem, solver.state())?;
131            return Ok((soln, checkpointing));
132        }
133        let mut state = solver.into_state();
134        problem.eqn.set_model_index(root_idx);
135        state.as_mut().apply_reset_with_mass::<M::LS, _>(problem)?;
136        solver = Tag::solver_with_state::<LS>(problem, state)?;
137    }
138    Ok((soln, checkpointing))
139}
140
141fn integral_from_soln<V>(soln: &Solution<V>) -> Result<V, DiffsolError>
142where
143    V: DefaultDenseMatrix,
144{
145    if soln.ts.is_empty() {
146        return Err(ode_solver_error!(
147            Other,
148            "Continuous adjoint solve returned no integral samples"
149        ));
150    }
151    Ok(soln.ys.column(soln.ts.len() - 1).into_owned())
152}
153
154#[allow(clippy::type_complexity)]
155fn solve_adjoint_fwd_with_tag<M, Eqn, LS, Tag>(
156    problem: &mut OdeSolverProblem<Eqn>,
157    t_eval: &[M::T],
158    params: &[f64],
159    method: OdeSolverType,
160    linear_solver: LinearSolverType,
161) -> Result<(Solution<M::V>, Box<dyn AdjointCheckpoint>), DiffsolError>
162where
163    M: Matrix<T: Scalar> + DefaultSolver + 'static,
164    M::V: DefaultDenseMatrix,
165    LS: LinearSolver<M>,
166    Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C> + 'static,
167    Tag: OdeSolverMethodTag<M, Eqn> + 'static,
168    for<'b> &'b M::V: VectorRef<M::V>,
169    for<'b> &'b M: MatrixRef<M>,
170{
171    let soln = Solution::new_dense(t_eval.to_vec())?;
172    let (soln, checkpointing) =
173        solve_with_checkpointing_with_tag::<M, Eqn, LS, Tag>(problem, soln)?;
174    Ok((
175        soln,
176        Box::new(AdjointCheckpointData::<M, Eqn, Tag>::new(
177            checkpointing,
178            params.to_vec(),
179            method,
180            linear_solver,
181        )),
182    ))
183}
184
185fn solve_continuous_adjoint_with_tag<M, Eqn, LS, Tag>(
186    problem: &mut OdeSolverProblem<Eqn>,
187    final_time: M::T,
188    method: OdeSolverType,
189) -> Result<(M::V, Vec<M::V>), DiffsolError>
190where
191    M: Matrix<T: Scalar> + DefaultSolver + 'static,
192    M::V: DefaultDenseMatrix,
193    LS: LinearSolver<M>,
194    Tag: OdeSolverMethodTag<M, Eqn> + 'static,
195    Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
196        + 'static
197        + diffsol::OdeEquationsImplicitAdjoint<M = M, T = M::T, V = M::V, C = M::C>,
198    for<'b> &'b M::V: VectorRef<M::V>,
199    for<'b> &'b M: MatrixRef<M>,
200{
201    let soln = Solution::new(final_time);
202    let (soln, checkpointing) =
203        solve_with_checkpointing_with_tag::<M, Eqn, LS, Tag>(problem, soln)?;
204    let integral = integral_from_soln(&soln)?;
205    let sg = match method {
206        OdeSolverType::Bdf => solve_adjoint_bkwds_with_fwd_bkwd_tag::<M, Eqn, LS, LS, Tag, BdfTag>(
207            problem,
208            &soln,
209            checkpointing,
210            &[],
211            None,
212        ),
213        OdeSolverType::Esdirk34 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
214            M,
215            Eqn,
216            LS,
217            LS,
218            Tag,
219            Esdirk34Tag,
220        >(problem, &soln, checkpointing, &[], None),
221        OdeSolverType::TrBdf2 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
222            M,
223            Eqn,
224            LS,
225            LS,
226            Tag,
227            TrBdf2Tag,
228        >(problem, &soln, checkpointing, &[], None),
229        OdeSolverType::Tsit45 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
230            M,
231            Eqn,
232            LS,
233            LS,
234            Tag,
235            Tsit45Tag,
236        >(problem, &soln, checkpointing, &[], None),
237    }?;
238    Ok((integral, sg))
239}
240
241fn solve_adjoint_bkwds_with_fwd_tag<M, Eqn, FwdLS, BwdLS, Tag>(
242    problem: &mut OdeSolverProblem<Eqn>,
243    checkpoint: &AdjointCheckpointData<M, Eqn, Tag>,
244    backwards_method: OdeSolverType,
245    dgdu_eval: &<M::V as DefaultDenseMatrix>::M,
246    t_eval: &[M::T],
247) -> Result<Vec<M::V>, DiffsolError>
248where
249    M: Matrix<T: Scalar> + DefaultSolver + 'static,
250    M::V: DefaultDenseMatrix,
251    FwdLS: LinearSolver<M>,
252    BwdLS: LinearSolver<M>,
253    Tag: OdeSolverMethodTag<M, Eqn> + 'static,
254    Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
255        + 'static
256        + diffsol::OdeEquationsImplicitAdjoint<M = M, T = M::T, V = M::V, C = M::C>,
257    for<'b> &'b M::V: VectorRef<M::V>,
258    for<'b> &'b M: MatrixRef<M>,
259{
260    // TODO: can we avoid cloning here? Adjoint equations require ownership of the checkpointing segments so maybe not
261    // unless we can change the adjoint equations to take references to the checkpointing segments instead
262    let checkpointing = checkpoint.checkpointing.clone();
263    let soln = Solution::new_dense(t_eval.to_vec())?;
264
265    // we will only consider a single output g for now, so nout_override is 1
266    let dgdu_eval = [dgdu_eval];
267    match backwards_method {
268        OdeSolverType::Bdf => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
269            M,
270            Eqn,
271            FwdLS,
272            BwdLS,
273            Tag,
274            BdfTag,
275        >(problem, &soln, checkpointing, &dgdu_eval, Some(1)),
276        OdeSolverType::Esdirk34 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
277            M,
278            Eqn,
279            FwdLS,
280            BwdLS,
281            Tag,
282            Esdirk34Tag,
283        >(problem, &soln, checkpointing, &dgdu_eval, Some(1)),
284        OdeSolverType::TrBdf2 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
285            M,
286            Eqn,
287            FwdLS,
288            BwdLS,
289            Tag,
290            TrBdf2Tag,
291        >(problem, &soln, checkpointing, &dgdu_eval, Some(1)),
292        OdeSolverType::Tsit45 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
293            M,
294            Eqn,
295            FwdLS,
296            BwdLS,
297            Tag,
298            Tsit45Tag,
299        >(problem, &soln, checkpointing, &dgdu_eval, Some(1)),
300    }
301}
302
303fn solve_adjoint_bkwds_with_fwd_bkwd_tag<'solver, M, Eqn, FwdLS, BwdLS, FwdTag, BwdTag>(
304    problem: &'solver mut OdeSolverProblem<Eqn>,
305    soln: &Solution<M::V>,
306    mut checkpointing: CheckpointingPath<Eqn, FwdTag::State>,
307    dgdu_eval: &[&<M::V as DefaultDenseMatrix>::M],
308    nout_override: Option<usize>,
309) -> Result<Vec<M::V>, DiffsolError>
310where
311    M: Matrix<T: Scalar> + DefaultSolver + 'solver,
312    M::V: DefaultDenseMatrix,
313    FwdLS: LinearSolver<M>,
314    BwdLS: LinearSolver<M>,
315    FwdTag: OdeSolverMethodTag<M, Eqn>,
316    BwdTag: OdeSolverMethodTag<M, Eqn>,
317    Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
318        + 'solver
319        + diffsol::OdeEquationsImplicitAdjoint<M = M, T = M::T, V = M::V, C = M::C>,
320    for<'b> &'b M::V: VectorRef<M::V>,
321    for<'b> &'b M: MatrixRef<M>,
322{
323    let checkpointing_len = checkpointing.len();
324    if checkpointing_len == 0 {
325        return Err(ode_solver_error!(
326            Other,
327            "Adjoint backward pass requires at least one checkpointing segment"
328        ));
329    }
330
331    let t_eval = if dgdu_eval.is_empty() {
332        &[]
333    } else {
334        soln.ts.as_slice()
335    };
336
337    let current_checkpointing = checkpointing
338        .pop()
339        .ok_or_else(|| ode_solver_error!(Other, "Adjoint backward pass returned no state"))?;
340    let model_index = checkpointing
341        .last()
342        .map(|segment| {
343            segment
344                .terminal_reset_root_idx()
345                .expect("Missing reset root index")
346        })
347        .unwrap_or(0);
348    problem.eqn_mut().set_model_index(model_index);
349    let fwd_solver = FwdTag::uninitialised_solver::<FwdLS>(&*problem)?;
350    let mut adjoint = BwdTag::solver_adjoint::<BwdLS, _>(
351        &*problem,
352        vec![current_checkpointing],
353        Some(fwd_solver),
354        nout_override,
355    )?;
356    loop {
357        let (mut state, adjoint_checkpointing) =
358            adjoint.solve_adjoint_backwards_pass(t_eval, dgdu_eval)?;
359        let Some(previous_checkpointing) = checkpointing.pop() else {
360            // the adjoint state holds one output channel per batch lane; the C API returns one
361            // gradient vector per channel
362            let sg = state.into_common().sg;
363            let nchannels = sg.context().nbatch();
364            return Ok((0..nchannels)
365                .map(|i| sg.get_batch(i).into_owned())
366                .collect());
367        };
368        let model_index = checkpointing
369            .last()
370            .map(|segment| {
371                segment
372                    .terminal_reset_root_idx()
373                    .expect("Missing reset root index")
374            })
375            .unwrap_or(0);
376        let fwd_state_minus = previous_checkpointing.last_checkpoint();
377        let fwd_state_plus = adjoint_checkpointing
378            .first()
379            .ok_or_else(|| {
380                ode_solver_error!(Other, "Adjoint backward pass returned no checkpointing")
381            })?
382            .first_checkpoint();
383        state.as_mut().apply_reset_with_adjoint(
384            problem.eqn(),
385            previous_checkpointing.terminal_reset_root_idx().unwrap(),
386            fwd_state_minus.as_ref(),
387            fwd_state_plus.as_ref(),
388            problem.integrate_out,
389        )?;
390        problem.eqn_mut().set_model_index(model_index);
391        let fwd_solver = FwdTag::uninitialised_solver::<FwdLS>(&*problem)?;
392        // TODO: remove clone here
393        let adjoint_eqn = problem.adjoint_equations(
394            vec![previous_checkpointing],
395            Some(fwd_solver),
396            nout_override,
397        )?;
398
399        adjoint = BwdTag::solver_adjoint_from_state::<BwdLS, _>(&*problem, state, adjoint_eqn)?;
400    }
401}
402
403impl OdeSolverType {
404    pub(crate) fn solve<M, Eqn, LS>(
405        &self,
406        problem: &mut OdeSolverProblem<Eqn>,
407        final_time: M::T,
408    ) -> Result<Solution<M::V>, DiffsolError>
409    where
410        M: Matrix<T: Scalar> + DefaultSolver,
411        M::V: DefaultDenseMatrix,
412        LS: LinearSolver<M>,
413        for<'b> &'b M::V: VectorRef<M::V>,
414        for<'b> &'b M: MatrixRef<M>,
415        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>,
416    {
417        match self {
418            OdeSolverType::Bdf => {
419                solve_with_tag::<M, Eqn, LS, BdfTag>(problem, Solution::new(final_time))
420            }
421            OdeSolverType::Esdirk34 => {
422                solve_with_tag::<M, Eqn, LS, Esdirk34Tag>(problem, Solution::new(final_time))
423            }
424            OdeSolverType::TrBdf2 => {
425                solve_with_tag::<M, Eqn, LS, TrBdf2Tag>(problem, Solution::new(final_time))
426            }
427            OdeSolverType::Tsit45 => {
428                solve_with_tag::<M, Eqn, LS, Tsit45Tag>(problem, Solution::new(final_time))
429            }
430        }
431    }
432
433    pub(crate) fn solve_dense<M, Eqn, LS>(
434        &self,
435        problem: &mut OdeSolverProblem<Eqn>,
436        t_eval: &[M::T],
437    ) -> Result<Solution<M::V>, DiffsolError>
438    where
439        M: Matrix<T: Scalar> + DefaultSolver,
440        M::V: DefaultDenseMatrix,
441        LS: LinearSolver<M>,
442        for<'b> &'b M::V: VectorRef<M::V>,
443        for<'b> &'b M: MatrixRef<M>,
444        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>,
445    {
446        match self {
447            OdeSolverType::Bdf => {
448                solve_with_tag::<M, Eqn, LS, BdfTag>(problem, Solution::new_dense(t_eval.to_vec())?)
449            }
450            OdeSolverType::Esdirk34 => solve_with_tag::<M, Eqn, LS, Esdirk34Tag>(
451                problem,
452                Solution::new_dense(t_eval.to_vec())?,
453            ),
454            OdeSolverType::TrBdf2 => solve_with_tag::<M, Eqn, LS, TrBdf2Tag>(
455                problem,
456                Solution::new_dense(t_eval.to_vec())?,
457            ),
458            OdeSolverType::Tsit45 => solve_with_tag::<M, Eqn, LS, Tsit45Tag>(
459                problem,
460                Solution::new_dense(t_eval.to_vec())?,
461            ),
462        }
463    }
464
465    #[allow(clippy::type_complexity)]
466    pub(crate) fn solve_fwd_sens<M, Eqn, LS>(
467        &self,
468        problem: &mut OdeSolverProblem<Eqn>,
469        t_eval: &[M::T],
470    ) -> Result<Solution<M::V>, DiffsolError>
471    where
472        M: Matrix<T: Scalar> + DefaultSolver,
473        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>,
474        M::V: DefaultDenseMatrix,
475        LS: LinearSolver<M>,
476        for<'b> &'b M::V: VectorRef<M::V>,
477        for<'b> &'b M: MatrixRef<M>,
478    {
479        match self {
480            OdeSolverType::Bdf => solve_fwd_sens_with_tag::<M, Eqn, LS, BdfTag>(problem, t_eval),
481            OdeSolverType::Esdirk34 => {
482                solve_fwd_sens_with_tag::<M, Eqn, LS, Esdirk34Tag>(problem, t_eval)
483            }
484            OdeSolverType::TrBdf2 => {
485                solve_fwd_sens_with_tag::<M, Eqn, LS, TrBdf2Tag>(problem, t_eval)
486            }
487            OdeSolverType::Tsit45 => {
488                solve_fwd_sens_with_tag::<M, Eqn, LS, Tsit45Tag>(problem, t_eval)
489            }
490        }
491    }
492
493    #[allow(clippy::type_complexity)]
494    pub(crate) fn solve_adjoint_fwd<M, Eqn, LS>(
495        &self,
496        problem: &mut OdeSolverProblem<Eqn>,
497        t_eval: &[M::T],
498        params: &[f64],
499        linear_solver: LinearSolverType,
500    ) -> Result<(Solution<M::V>, Box<dyn AdjointCheckpoint>), DiffsolError>
501    where
502        M: Matrix<T: Scalar> + DefaultSolver + 'static,
503        M::V: DefaultDenseMatrix,
504        LS: LinearSolver<M>,
505        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C> + 'static,
506        for<'b> &'b M::V: VectorRef<M::V>,
507        for<'b> &'b M: MatrixRef<M>,
508    {
509        match self {
510            OdeSolverType::Bdf => solve_adjoint_fwd_with_tag::<M, Eqn, LS, BdfTag>(
511                problem,
512                t_eval,
513                params,
514                *self,
515                linear_solver,
516            ),
517            OdeSolverType::Esdirk34 => solve_adjoint_fwd_with_tag::<M, Eqn, LS, Esdirk34Tag>(
518                problem,
519                t_eval,
520                params,
521                *self,
522                linear_solver,
523            ),
524            OdeSolverType::TrBdf2 => solve_adjoint_fwd_with_tag::<M, Eqn, LS, TrBdf2Tag>(
525                problem,
526                t_eval,
527                params,
528                *self,
529                linear_solver,
530            ),
531            OdeSolverType::Tsit45 => solve_adjoint_fwd_with_tag::<M, Eqn, LS, Tsit45Tag>(
532                problem,
533                t_eval,
534                params,
535                *self,
536                linear_solver,
537            ),
538        }
539    }
540
541    pub(crate) fn solve_continuous_adjoint<M, Eqn, LS>(
542        &self,
543        problem: &mut OdeSolverProblem<Eqn>,
544        final_time: M::T,
545    ) -> Result<(M::V, Vec<M::V>), DiffsolError>
546    where
547        M: Matrix<T: Scalar> + DefaultSolver + 'static,
548        M::V: DefaultDenseMatrix,
549        LS: LinearSolver<M>,
550        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
551            + 'static
552            + diffsol::OdeEquationsImplicitAdjoint,
553        for<'b> &'b M::V: VectorRef<M::V>,
554        for<'b> &'b M: MatrixRef<M>,
555    {
556        match self {
557            OdeSolverType::Bdf => {
558                solve_continuous_adjoint_with_tag::<M, Eqn, LS, BdfTag>(problem, final_time, *self)
559            }
560            OdeSolverType::Esdirk34 => {
561                solve_continuous_adjoint_with_tag::<M, Eqn, LS, Esdirk34Tag>(
562                    problem, final_time, *self,
563                )
564            }
565            OdeSolverType::TrBdf2 => solve_continuous_adjoint_with_tag::<M, Eqn, LS, TrBdf2Tag>(
566                problem, final_time, *self,
567            ),
568            OdeSolverType::Tsit45 => solve_continuous_adjoint_with_tag::<M, Eqn, LS, Tsit45Tag>(
569                problem, final_time, *self,
570            ),
571        }
572    }
573
574    pub(crate) fn solve_adjoint_bkwd<M, Eqn, BwdLS>(
575        &self,
576        problem: &mut OdeSolverProblem<Eqn>,
577        checkpoint: &dyn AdjointCheckpoint,
578        dgdu_eval: &<M::V as DefaultDenseMatrix>::M,
579        t_eval: &[M::T],
580    ) -> Result<Vec<M::V>, DiffsolError>
581    where
582        M: Matrix<T: Scalar> + DefaultSolver + LuValidator<M> + KluValidator<M> + 'static,
583        M::V: DefaultDenseMatrix,
584        BwdLS: LinearSolver<M>,
585        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
586            + 'static
587            + diffsol::OdeEquationsImplicitAdjoint,
588        for<'b> &'b M::V: VectorRef<M::V>,
589        for<'b> &'b M: MatrixRef<M>,
590    {
591        match checkpoint.method() {
592            OdeSolverType::Bdf => {
593                let data = checkpoint.data::<M, Eqn, BdfTag>()?;
594                match data.linear_solver() {
595                    LinearSolverType::Default => {
596                        solve_adjoint_bkwds_with_fwd_tag::<
597                            M,
598                            Eqn,
599                            <M as DefaultSolver>::LS,
600                            BwdLS,
601                            BdfTag,
602                        >(problem, data, *self, dgdu_eval, t_eval)
603                    }
604                    LinearSolverType::Lu => {
605                        solve_adjoint_bkwds_with_fwd_tag::<
606                            M,
607                            Eqn,
608                            <M as LuValidator<M>>::LS,
609                            BwdLS,
610                            BdfTag,
611                        >(problem, data, *self, dgdu_eval, t_eval)
612                    }
613                    LinearSolverType::Klu => {
614                        solve_adjoint_bkwds_with_fwd_tag::<
615                            M,
616                            Eqn,
617                            <M as KluValidator<M>>::LS,
618                            BwdLS,
619                            BdfTag,
620                        >(problem, data, *self, dgdu_eval, t_eval)
621                    }
622                }
623            }
624            OdeSolverType::Esdirk34 => {
625                let data = checkpoint.data::<M, Eqn, Esdirk34Tag>()?;
626                match data.linear_solver() {
627                    LinearSolverType::Default => {
628                        solve_adjoint_bkwds_with_fwd_tag::<
629                            M,
630                            Eqn,
631                            <M as DefaultSolver>::LS,
632                            BwdLS,
633                            Esdirk34Tag,
634                        >(problem, data, *self, dgdu_eval, t_eval)
635                    }
636                    LinearSolverType::Lu => {
637                        solve_adjoint_bkwds_with_fwd_tag::<
638                            M,
639                            Eqn,
640                            <M as LuValidator<M>>::LS,
641                            BwdLS,
642                            Esdirk34Tag,
643                        >(problem, data, *self, dgdu_eval, t_eval)
644                    }
645                    LinearSolverType::Klu => {
646                        solve_adjoint_bkwds_with_fwd_tag::<
647                            M,
648                            Eqn,
649                            <M as KluValidator<M>>::LS,
650                            BwdLS,
651                            Esdirk34Tag,
652                        >(problem, data, *self, dgdu_eval, t_eval)
653                    }
654                }
655            }
656            OdeSolverType::TrBdf2 => {
657                let data = checkpoint.data::<M, Eqn, TrBdf2Tag>()?;
658                match data.linear_solver() {
659                    LinearSolverType::Default => {
660                        solve_adjoint_bkwds_with_fwd_tag::<
661                            M,
662                            Eqn,
663                            <M as DefaultSolver>::LS,
664                            BwdLS,
665                            TrBdf2Tag,
666                        >(problem, data, *self, dgdu_eval, t_eval)
667                    }
668                    LinearSolverType::Lu => {
669                        solve_adjoint_bkwds_with_fwd_tag::<
670                            M,
671                            Eqn,
672                            <M as LuValidator<M>>::LS,
673                            BwdLS,
674                            TrBdf2Tag,
675                        >(problem, data, *self, dgdu_eval, t_eval)
676                    }
677                    LinearSolverType::Klu => {
678                        solve_adjoint_bkwds_with_fwd_tag::<
679                            M,
680                            Eqn,
681                            <M as KluValidator<M>>::LS,
682                            BwdLS,
683                            TrBdf2Tag,
684                        >(problem, data, *self, dgdu_eval, t_eval)
685                    }
686                }
687            }
688            OdeSolverType::Tsit45 => {
689                let data = checkpoint.data::<M, Eqn, Tsit45Tag>()?;
690                match data.linear_solver() {
691                    LinearSolverType::Default => {
692                        solve_adjoint_bkwds_with_fwd_tag::<
693                            M,
694                            Eqn,
695                            <M as DefaultSolver>::LS,
696                            BwdLS,
697                            Tsit45Tag,
698                        >(problem, data, *self, dgdu_eval, t_eval)
699                    }
700                    LinearSolverType::Lu => {
701                        solve_adjoint_bkwds_with_fwd_tag::<
702                            M,
703                            Eqn,
704                            <M as LuValidator<M>>::LS,
705                            BwdLS,
706                            Tsit45Tag,
707                        >(problem, data, *self, dgdu_eval, t_eval)
708                    }
709                    LinearSolverType::Klu => {
710                        solve_adjoint_bkwds_with_fwd_tag::<
711                            M,
712                            Eqn,
713                            <M as KluValidator<M>>::LS,
714                            BwdLS,
715                            Tsit45Tag,
716                        >(problem, data, *self, dgdu_eval, t_eval)
717                    }
718                }
719            }
720        }
721    }
722}
723
724#[cfg(all(test, any(feature = "diffsl-cranelift", feature = "diffsl-llvm")))]
725mod tests {
726    use diffsol::{
727        CodegenModuleCompile, CodegenModuleJit, DefaultSolver, DenseMatrix, OdeBuilder,
728        OdeSolverProblem,
729    };
730    #[cfg(feature = "diffsl-llvm")]
731    use diffsol::{DefaultDenseMatrix, Matrix, MatrixCommon, Op, Vector};
732
733    #[cfg(feature = "diffsl-llvm")]
734    use crate::linear_solver_type::LinearSolverType;
735    use crate::test_support::{
736        assert_close, hybrid_logistic_diffsl_code, hybrid_logistic_state, logistic_diffsl_code,
737        logistic_state, LOGISTIC_X0,
738    };
739    #[cfg(feature = "diffsl-llvm")]
740    use crate::test_support::{hybrid_logistic_state_dr, logistic_state_dr};
741    use crate::valid_linear_solver::LuValidator;
742
743    use super::OdeSolverType;
744
745    type M = diffsol::NalgebraMat<f64>;
746
747    fn build_problem<CG>(code: &str) -> OdeSolverProblem<diffsol::DiffSl<M, CG>>
748    where
749        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
750    {
751        OdeBuilder::<M>::new()
752            .p([2.0])
753            .rtol(1e-6)
754            .atol([1e-6])
755            .sens_rtol(1e-6)
756            .sens_atol([1e-6])
757            .build_from_diffsl::<CG>(code)
758            .unwrap()
759    }
760
761    fn assert_dense_solution_matches_expected(
762        soln: &diffsol::Solution<diffsol::NalgebraVec<f64>>,
763        t_eval: &[f64],
764        expected: impl Fn(f64) -> f64,
765    ) {
766        assert_eq!(soln.ts, t_eval);
767        for (i, &t) in t_eval.iter().enumerate() {
768            assert_close(
769                soln.ys.get_index(0, i),
770                expected(t),
771                5e-4,
772                &format!("solution[{i}]"),
773            );
774        }
775    }
776
777    fn test_all_solver_variants<CG>()
778    where
779        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
780    {
781        let t_eval = [0.25, 0.5, 1.0];
782        for method in [
783            OdeSolverType::Bdf,
784            OdeSolverType::Esdirk34,
785            OdeSolverType::TrBdf2,
786            OdeSolverType::Tsit45,
787        ] {
788            let mut problem = build_problem::<CG>(logistic_diffsl_code());
789            let soln = method
790                .solve::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(&mut problem, 1.0)
791                .unwrap();
792            assert_close(*soln.ts.last().unwrap(), 1.0, 5e-4, "solve final time");
793            assert_close(
794                soln.ys.get_index(0, soln.ts.len() - 1),
795                logistic_state(LOGISTIC_X0, 2.0, 1.0),
796                5e-4,
797                "solve final value",
798            );
799
800            let mut problem = build_problem::<CG>(logistic_diffsl_code());
801            let soln = method
802                .solve_dense::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(
803                    &mut problem,
804                    &t_eval,
805                )
806                .unwrap();
807            assert_dense_solution_matches_expected(&soln, &t_eval, |t| {
808                logistic_state(LOGISTIC_X0, 2.0, t)
809            });
810        }
811    }
812
813    fn test_all_hybrid_solver_variants<CG>()
814    where
815        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
816    {
817        let t_eval = [0.5, 1.0, 1.25, 1.5, 2.0];
818        for method in [
819            OdeSolverType::Bdf,
820            OdeSolverType::Esdirk34,
821            OdeSolverType::TrBdf2,
822            OdeSolverType::Tsit45,
823        ] {
824            let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
825            let soln = method
826                .solve::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(&mut problem, 2.0)
827                .unwrap();
828            assert_close(*soln.ts.last().unwrap(), 2.0, 5e-4, "hybrid final time");
829            assert_close(
830                soln.ys.get_index(0, soln.ts.len() - 1),
831                hybrid_logistic_state(2.0, 2.0),
832                5e-4,
833                "hybrid final value",
834            );
835
836            let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
837            let soln = method
838                .solve_dense::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(
839                    &mut problem,
840                    &t_eval,
841                )
842                .unwrap();
843            assert_dense_solution_matches_expected(&soln, &t_eval, |t| {
844                hybrid_logistic_state(2.0, t)
845            });
846        }
847    }
848
849    fn test_all_solver_variants_with_lu<CG>()
850    where
851        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
852    {
853        let t_eval = [0.25, 0.5, 1.0];
854        for method in [
855            OdeSolverType::Bdf,
856            OdeSolverType::Esdirk34,
857            OdeSolverType::TrBdf2,
858            OdeSolverType::Tsit45,
859        ] {
860            let mut problem = build_problem::<CG>(logistic_diffsl_code());
861            let soln = method
862                .solve::<M, diffsol::DiffSl<M, CG>, <M as LuValidator<M>>::LS>(&mut problem, 1.0)
863                .unwrap();
864            assert_close(*soln.ts.last().unwrap(), 1.0, 5e-4, "lu solve final time");
865
866            let mut problem = build_problem::<CG>(logistic_diffsl_code());
867            let soln = method
868                .solve_dense::<M, diffsol::DiffSl<M, CG>, <M as LuValidator<M>>::LS>(
869                    &mut problem,
870                    &t_eval,
871                )
872                .unwrap();
873            assert_dense_solution_matches_expected(&soln, &t_eval, |t| {
874                logistic_state(LOGISTIC_X0, 2.0, t)
875            });
876        }
877    }
878
879    fn test_all_hybrid_solver_variants_with_lu<CG>()
880    where
881        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
882    {
883        let t_eval = [0.5, 1.0, 1.25, 1.5, 2.0];
884        for method in [
885            OdeSolverType::Bdf,
886            OdeSolverType::Esdirk34,
887            OdeSolverType::TrBdf2,
888            OdeSolverType::Tsit45,
889        ] {
890            let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
891            let soln = method
892                .solve::<M, diffsol::DiffSl<M, CG>, <M as LuValidator<M>>::LS>(&mut problem, 2.0)
893                .unwrap();
894            assert_close(*soln.ts.last().unwrap(), 2.0, 5e-4, "lu hybrid final time");
895
896            let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
897            let soln = method
898                .solve_dense::<M, diffsol::DiffSl<M, CG>, <M as LuValidator<M>>::LS>(
899                    &mut problem,
900                    &t_eval,
901                )
902                .unwrap();
903            assert_dense_solution_matches_expected(&soln, &t_eval, |t| {
904                hybrid_logistic_state(2.0, t)
905            });
906        }
907    }
908
909    fn assert_direct_hybrid_restart_path_for_method<CG>(method: OdeSolverType)
910    where
911        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
912    {
913        let t_eval = [0.5, 1.0, 1.25, 1.5, 2.0];
914
915        let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
916        let soln = method
917            .solve::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(&mut problem, 2.0)
918            .unwrap();
919        assert_close(
920            *soln.ts.last().unwrap(),
921            2.0,
922            5e-4,
923            "direct hybrid restart final time",
924        );
925        assert_close(
926            soln.ys.get_index(0, soln.ts.len() - 1),
927            hybrid_logistic_state(2.0, 2.0),
928            5e-4,
929            "direct hybrid restart final value",
930        );
931
932        let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
933        let soln = method
934            .solve_dense::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(
935                &mut problem,
936                &t_eval,
937            )
938            .unwrap();
939        assert_dense_solution_matches_expected(&soln, &t_eval, |t| hybrid_logistic_state(2.0, t));
940    }
941
942    #[cfg(feature = "diffsl-llvm")]
943    fn test_all_sensitivity_solver_variants() {
944        let t_eval = [0.25, 0.5, 1.0];
945        for method in [
946            OdeSolverType::Bdf,
947            OdeSolverType::Esdirk34,
948            OdeSolverType::TrBdf2,
949            OdeSolverType::Tsit45,
950        ] {
951            let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
952            let soln = method
953                .solve_fwd_sens::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
954                    &mut problem,
955                    &t_eval,
956                )
957                .unwrap();
958            for (i, &t) in t_eval.iter().enumerate() {
959                assert_close(
960                    soln.y_sens[0].get_index(0, i),
961                    logistic_state_dr(LOGISTIC_X0, 2.0, t),
962                    5e-4,
963                    &format!("fwd_sens[{i}]"),
964                );
965            }
966
967            let mut problem = build_problem::<diffsol::LlvmModule>(hybrid_logistic_diffsl_code());
968            let soln = method
969                .solve_fwd_sens::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
970                    &mut problem,
971                    &t_eval,
972                )
973                .unwrap();
974            for (i, &t) in t_eval.iter().enumerate() {
975                assert_close(
976                    soln.y_sens[0].get_index(0, i),
977                    hybrid_logistic_state_dr(2.0, t),
978                    5e-4,
979                    &format!("hybrid_fwd_sens[{i}]"),
980                );
981            }
982        }
983    }
984
985    #[cfg(feature = "diffsl-llvm")]
986    fn test_lu_sensitivity_and_adjoint_solver_variants() {
987        let t_eval = [0.25, 0.5, 1.0];
988        for method in [
989            OdeSolverType::Bdf,
990            OdeSolverType::Esdirk34,
991            OdeSolverType::TrBdf2,
992            OdeSolverType::Tsit45,
993        ] {
994            let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
995            let soln = method
996                .solve_fwd_sens::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as LuValidator<M>>::LS>(
997                    &mut problem,
998                    &t_eval,
999                )
1000                .unwrap();
1001            for (i, &t) in t_eval.iter().enumerate() {
1002                assert_close(
1003                    soln.y_sens[0].get_index(0, i),
1004                    logistic_state_dr(LOGISTIC_X0, 2.0, t),
1005                    5e-4,
1006                    &format!("lu fwd_sens[{i}]"),
1007                );
1008            }
1009        }
1010
1011        let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
1012        let adjoint_t_eval = [0.0, 0.25, 0.5, 1.0];
1013        let (soln, checkpoint) = OdeSolverType::Bdf
1014            .solve_adjoint_fwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as LuValidator<M>>::LS>(
1015                &mut problem,
1016                &adjoint_t_eval,
1017                &[2.0],
1018                LinearSolverType::Lu,
1019            )
1020            .unwrap();
1021        let dgdu = <<M as MatrixCommon>::V as DefaultDenseMatrix>::M::zeros(
1022            problem.eqn.nout(),
1023            soln.ts.len(),
1024            problem.context().to_owned(),
1025        );
1026        let gradient = OdeSolverType::TrBdf2
1027            .solve_adjoint_bkwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as LuValidator<M>>::LS>(
1028                &mut problem,
1029                checkpoint.as_ref(),
1030                &dgdu,
1031                &soln.ts,
1032            )
1033            .unwrap();
1034        assert_eq!(gradient.len(), 1);
1035        assert!(gradient[0].get_index(0).is_finite());
1036    }
1037
1038    #[cfg(feature = "diffsl-llvm")]
1039    fn test_direct_hybrid_sensitivity_restart_paths() {
1040        let t_eval = [0.5, 1.0, 2.5, 3.0, 4.5];
1041        for method in [
1042            OdeSolverType::Esdirk34,
1043            OdeSolverType::TrBdf2,
1044            OdeSolverType::Tsit45,
1045        ] {
1046            let mut problem = build_problem::<diffsol::LlvmModule>(hybrid_logistic_diffsl_code());
1047            let soln = method
1048                .solve_fwd_sens::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1049                    &mut problem,
1050                    &t_eval,
1051                )
1052                .unwrap();
1053            for (i, &t) in t_eval.iter().enumerate() {
1054                assert_close(
1055                    soln.ys.get_index(0, i),
1056                    hybrid_logistic_state(2.0, t),
1057                    5e-4,
1058                    &format!("direct hybrid value[{i}]"),
1059                );
1060                assert_close(
1061                    soln.y_sens[0].get_index(0, i),
1062                    hybrid_logistic_state_dr(2.0, t),
1063                    5e-4,
1064                    &format!("direct hybrid fwd sens[{i}]"),
1065                );
1066            }
1067        }
1068    }
1069
1070    #[cfg(feature = "diffsl-llvm")]
1071    fn test_adjoint_backwards_methods_and_klu_branch() {
1072        for backwards_method in [OdeSolverType::Esdirk34, OdeSolverType::TrBdf2] {
1073            let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
1074            let t_eval = [0.0, 0.25, 0.5, 1.0];
1075            let (soln, checkpoint) = OdeSolverType::Bdf
1076                .solve_adjoint_fwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1077                    &mut problem,
1078                    &t_eval,
1079                    &[2.0],
1080                    LinearSolverType::Default,
1081                )
1082                .unwrap();
1083            let dgdu = <<M as MatrixCommon>::V as DefaultDenseMatrix>::M::zeros(
1084                problem.eqn.nout(),
1085                soln.ts.len(),
1086                problem.context().to_owned(),
1087            );
1088            let gradient = backwards_method
1089                .solve_adjoint_bkwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as crate::valid_linear_solver::KluValidator<M>>::LS>(
1090                    &mut problem,
1091                    checkpoint.as_ref(),
1092                    &dgdu,
1093                    &soln.ts,
1094                )
1095                .unwrap();
1096            assert_eq!(gradient.len(), 1);
1097            assert!(gradient[0].get_index(0).is_finite());
1098        }
1099
1100        let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
1101        let t_eval = [0.0, 0.25, 0.5, 1.0];
1102        let (soln, checkpoint) = OdeSolverType::Tsit45
1103            .solve_adjoint_fwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1104                &mut problem,
1105                &t_eval,
1106                &[2.0],
1107                LinearSolverType::Default,
1108            )
1109            .unwrap();
1110        let dgdu = <<M as MatrixCommon>::V as DefaultDenseMatrix>::M::zeros(
1111            problem.eqn.nout(),
1112            soln.ts.len(),
1113            problem.context().to_owned(),
1114        );
1115        let gradient = OdeSolverType::Bdf
1116            .solve_adjoint_bkwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1117                &mut problem,
1118                checkpoint.as_ref(),
1119                &dgdu,
1120                &soln.ts,
1121            )
1122            .unwrap();
1123        assert_eq!(gradient.len(), 1);
1124        assert!(gradient[0].get_index(0).is_finite());
1125    }
1126
1127    #[cfg(feature = "diffsl-llvm")]
1128    fn test_all_adjoint_solver_variants() {
1129        let t_eval = [0.0, 0.25, 0.5, 1.0];
1130        for method in [
1131            OdeSolverType::Bdf,
1132            OdeSolverType::Esdirk34,
1133            OdeSolverType::TrBdf2,
1134        ] {
1135            let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
1136            let (soln, checkpoint) = method
1137                .solve_adjoint_fwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1138                    &mut problem,
1139                    &t_eval,
1140                    &[2.0],
1141                    LinearSolverType::Default,
1142                )
1143                .unwrap();
1144            let dgdu = <<M as MatrixCommon>::V as DefaultDenseMatrix>::M::zeros(
1145                problem.eqn.nout(),
1146                soln.ts.len(),
1147                problem.context().to_owned(),
1148            );
1149            let gradient = OdeSolverType::Bdf
1150                .solve_adjoint_bkwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1151                    &mut problem,
1152                    checkpoint.as_ref(),
1153                    &dgdu,
1154                    &soln.ts,
1155                )
1156                .unwrap();
1157            assert_eq!(gradient.len(), 1);
1158            assert!(gradient[0].get_index(0).is_finite());
1159        }
1160    }
1161
1162    #[cfg(feature = "diffsl-cranelift")]
1163    #[test]
1164    fn runtime_dispatch_solves_all_variants_for_cranelift() {
1165        test_all_solver_variants::<diffsol::CraneliftJitModule>();
1166        test_all_solver_variants_with_lu::<diffsol::CraneliftJitModule>();
1167    }
1168
1169    #[cfg(feature = "diffsl-cranelift")]
1170    #[test]
1171    fn runtime_dispatch_solves_all_hybrid_variants_for_cranelift() {
1172        test_all_hybrid_solver_variants::<diffsol::CraneliftJitModule>();
1173        test_all_hybrid_solver_variants_with_lu::<diffsol::CraneliftJitModule>();
1174        assert_direct_hybrid_restart_path_for_method::<diffsol::CraneliftJitModule>(
1175            OdeSolverType::Esdirk34,
1176        );
1177        assert_direct_hybrid_restart_path_for_method::<diffsol::CraneliftJitModule>(
1178            OdeSolverType::TrBdf2,
1179        );
1180        assert_direct_hybrid_restart_path_for_method::<diffsol::CraneliftJitModule>(
1181            OdeSolverType::Tsit45,
1182        );
1183    }
1184
1185    #[cfg(feature = "diffsl-llvm")]
1186    #[test]
1187    fn runtime_dispatch_solves_all_variants_for_llvm() {
1188        test_all_solver_variants::<diffsol::LlvmModule>();
1189        test_all_solver_variants_with_lu::<diffsol::LlvmModule>();
1190    }
1191
1192    #[cfg(feature = "diffsl-llvm")]
1193    #[test]
1194    fn runtime_dispatch_solves_all_hybrid_variants_for_llvm() {
1195        test_all_hybrid_solver_variants::<diffsol::LlvmModule>();
1196        test_all_hybrid_solver_variants_with_lu::<diffsol::LlvmModule>();
1197        assert_direct_hybrid_restart_path_for_method::<diffsol::LlvmModule>(
1198            OdeSolverType::Esdirk34,
1199        );
1200        assert_direct_hybrid_restart_path_for_method::<diffsol::LlvmModule>(OdeSolverType::TrBdf2);
1201        assert_direct_hybrid_restart_path_for_method::<diffsol::LlvmModule>(OdeSolverType::Tsit45);
1202    }
1203
1204    #[cfg(feature = "diffsl-llvm")]
1205    #[test]
1206    fn runtime_dispatch_solves_all_forward_sensitivity_variants_for_llvm() {
1207        test_all_sensitivity_solver_variants();
1208        test_lu_sensitivity_and_adjoint_solver_variants();
1209        test_direct_hybrid_sensitivity_restart_paths();
1210    }
1211
1212    #[cfg(feature = "diffsl-llvm")]
1213    #[test]
1214    fn runtime_dispatch_solves_all_adjoint_variants_for_llvm() {
1215        test_all_adjoint_solver_variants();
1216        test_adjoint_backwards_methods_and_klu_branch();
1217    }
1218}
diffsol_c/ode_solver_type.rs

diffsol_c/
ode_solver_type.rs
