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

diffsol_c/
ode_solver_type.rs
