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::{
6    OdeEquationsImplicit, OdeEquationsImplicitAdjoint, OdeEquationsImplicitSens,
7};
8use diffsol::{
9    matrix::MatrixRef, DefaultDenseMatrix, DenseMatrix, LinearSolver, Matrix, OdeSolverProblem,
10    OdeSolverState, VectorRef, VectorView,
11};
12use diffsol::{
13    ode_solver_error, AdjointOdeSolverMethod, CheckpointingPath, DefaultSolver, OdeSolverMethod,
14    OdeSolverStopReason, SensitivitiesOdeSolverMethod, Solution, StepVjp,
15};
16use schemars::JsonSchema;
17use serde::{Deserialize, Serialize};
18
19use crate::adjoint_checkpoint::{AdjointCheckpoint, AdjointCheckpointData};
20use crate::ode_solver_tag::{BdfTag, Esdirk34Tag, OdeSolverMethodTag, TrBdf2Tag, Tsit45Tag};
21use crate::scalar_type::Scalar;
22use crate::vjp_checkpoint::{VjpCheckpoint, VjpCheckpointData};
23use crate::{
24    linear_solver_type::LinearSolverType,
25    valid_linear_solver::{KluValidator, LuValidator},
26};
27use diffsol::Vector;
28
29/// Enumerates the possible ODE solver methods for diffsol. See the solver descriptions in the diffsol documentation (<https://github.com/martinjrobins/diffsol>) for more details.
30///
31/// :attr bdf: Backward Differentiation Formula (BDF) method for stiff ODEs and singular mass matrices
32/// :attr esdirk34: Explicit Singly Diagonally Implicit Runge-Kutta (ESDIRK) method for moderately stiff ODEs and singular mass matrices.
33/// :attr tr_bdf2: Trapezoidal Backward Differentiation Formula of order 2 (TR-BDF2) method for moderately stiff ODEs and singular mass matrices.
34/// :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.
35#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
36#[serde(rename_all = "snake_case")]
37pub enum OdeSolverType {
38    Bdf,
39    Esdirk34,
40    TrBdf2,
41    Tsit45,
42}
43
44fn solve_with_tag<M, Eqn, LS, Tag>(
45    problem: &mut OdeSolverProblem<Eqn>,
46    mut soln: Solution<M::V>,
47) -> Result<Solution<M::V>, DiffsolError>
48where
49    M: Matrix<T: Scalar> + DefaultSolver,
50    M::V: DefaultDenseMatrix,
51    LS: LinearSolver<M>,
52    Tag: OdeSolverMethodTag<M, Eqn>,
53    Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C>,
54    for<'b> &'b M::V: VectorRef<M::V>,
55    for<'b> &'b M: MatrixRef<M>,
56{
57    let mut solver = Tag::solver::<LS>(problem)?;
58    while !soln.is_complete() {
59        solver = solver.solve_soln(&mut soln)?;
60        let root_idx = match soln.stop_reason {
61            Some(OdeSolverStopReason::RootFound(_, root_idx)) if !soln.is_complete() => root_idx,
62            _ => continue,
63        };
64        if problem.eqn.reset().is_none() {
65            soln.truncate(problem, solver.state())?;
66            return Ok(soln);
67        }
68        let mut state = solver.into_state();
69        problem.eqn.set_model_index(root_idx);
70        state.as_mut().apply_reset_with_mass::<M::LS, _>(problem)?;
71        solver = Tag::solver_with_state::<LS>(problem, state)?;
72    }
73    Ok(soln)
74}
75
76fn solve_fwd_sens_with_tag<M, Eqn, LS, Tag>(
77    problem: &mut OdeSolverProblem<Eqn>,
78    t_eval: &[M::T],
79) -> Result<Solution<M::V>, DiffsolError>
80where
81    M: Matrix<T: Scalar> + DefaultSolver,
82    M::V: DefaultDenseMatrix,
83    LS: LinearSolver<M>,
84    Tag: OdeSolverMethodTag<M, Eqn>,
85    Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>,
86    for<'b> &'b M::V: VectorRef<M::V>,
87    for<'b> &'b M: MatrixRef<M>,
88{
89    let mut soln = Solution::new_dense(t_eval.to_vec())?;
90    let mut solver = Tag::solver_sens::<LS>(problem)?;
91    while !soln.is_complete() {
92        solver = solver.solve_soln_sensitivities(&mut soln)?;
93        let root_idx = match soln.stop_reason {
94            Some(OdeSolverStopReason::RootFound(_, root_idx)) if !soln.is_complete() => root_idx,
95            _ => continue,
96        };
97        if problem.eqn.reset().is_none() {
98            soln.truncate_sens(problem, solver.state())?;
99            return Ok(soln);
100        }
101        let mut state = solver.into_state();
102        problem.eqn.set_model_index(root_idx);
103        state
104            .as_mut()
105            .apply_reset_with_sens_mass::<M::LS, _>(problem, root_idx)?;
106        solver = Tag::solver_sens_with_state::<LS>(problem, state)?;
107    }
108    Ok(soln)
109}
110
111#[allow(clippy::type_complexity)]
112fn solve_with_checkpointing_with_tag<M, Eqn, LS, Tag>(
113    problem: &mut OdeSolverProblem<Eqn>,
114    mut soln: Solution<M::V>,
115) -> Result<(Solution<M::V>, CheckpointingPath<Eqn, Tag::State>), DiffsolError>
116where
117    M: Matrix<T: Scalar> + DefaultSolver,
118    M::V: DefaultDenseMatrix,
119    LS: LinearSolver<M>,
120    Tag: OdeSolverMethodTag<M, Eqn>,
121    Eqn: 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, Eqn, LS, Tag>(
160    problem: &mut OdeSolverProblem<Eqn>,
161    t_eval: &[M::T],
162    method: OdeSolverType,
163    linear_solver: LinearSolverType,
164) -> Result<(Solution<M::V>, Box<dyn AdjointCheckpoint>), DiffsolError>
165where
166    M: Matrix<T: Scalar> + DefaultSolver + 'static,
167    M::V: DefaultDenseMatrix,
168    LS: LinearSolver<M>,
169    Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C> + 'static,
170    Tag: OdeSolverMethodTag<M, Eqn> + 'static,
171    for<'b> &'b M::V: VectorRef<M::V>,
172    for<'b> &'b M: MatrixRef<M>,
173{
174    // keep the parameters of all the lanes, so the backward pass can use them again
175    let mut params = M::V::zeros(problem.eqn.nparams(), problem.context().clone());
176    problem.eqn.get_params(&mut params);
177    let soln = Solution::new_dense(t_eval.to_vec())?;
178    let (soln, checkpointing) =
179        solve_with_checkpointing_with_tag::<M, Eqn, LS, Tag>(problem, soln)?;
180    Ok((
181        soln,
182        Box::new(AdjointCheckpointData::<M, Eqn, Tag>::new(
183            checkpointing,
184            params,
185            method,
186            linear_solver,
187        )),
188    ))
189}
190
191fn solve_continuous_adjoint_with_tag<M, Eqn, LS, Tag>(
192    problem: &mut OdeSolverProblem<Eqn>,
193    final_time: M::T,
194    method: OdeSolverType,
195) -> Result<(M::V, M::V), DiffsolError>
196where
197    M: Matrix<T: Scalar> + DefaultSolver + 'static,
198    M::V: DefaultDenseMatrix,
199    LS: LinearSolver<M>,
200    Tag: OdeSolverMethodTag<M, Eqn> + 'static,
201    Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
202        + 'static
203        + diffsol::OdeEquationsImplicitAdjoint<M = M, T = M::T, V = M::V, C = M::C>,
204    for<'b> &'b M::V: VectorRef<M::V>,
205    for<'b> &'b M: MatrixRef<M>,
206{
207    let soln = Solution::new(final_time);
208    let (soln, checkpointing) =
209        solve_with_checkpointing_with_tag::<M, Eqn, LS, Tag>(problem, soln)?;
210    let integral = integral_from_soln(&soln)?;
211    let sg = match method {
212        OdeSolverType::Bdf => solve_adjoint_bkwds_with_fwd_bkwd_tag::<M, Eqn, LS, LS, Tag, BdfTag>(
213            problem,
214            &soln,
215            checkpointing,
216            &[],
217            None,
218        ),
219        OdeSolverType::Esdirk34 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
220            M,
221            Eqn,
222            LS,
223            LS,
224            Tag,
225            Esdirk34Tag,
226        >(problem, &soln, checkpointing, &[], None),
227        OdeSolverType::TrBdf2 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
228            M,
229            Eqn,
230            LS,
231            LS,
232            Tag,
233            TrBdf2Tag,
234        >(problem, &soln, checkpointing, &[], None),
235        OdeSolverType::Tsit45 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
236            M,
237            Eqn,
238            LS,
239            LS,
240            Tag,
241            Tsit45Tag,
242        >(problem, &soln, checkpointing, &[], None),
243    }?;
244    Ok((integral, sg))
245}
246
247fn solve_adjoint_bkwds_with_fwd_tag<M, Eqn, FwdLS, BwdLS, Tag>(
248    problem: &mut OdeSolverProblem<Eqn>,
249    checkpoint: &AdjointCheckpointData<M, Eqn, Tag>,
250    backwards_method: OdeSolverType,
251    dgdu_eval: &<M::V as DefaultDenseMatrix>::M,
252    t_eval: &[M::T],
253) -> Result<M::V, DiffsolError>
254where
255    M: Matrix<T: Scalar> + DefaultSolver + 'static,
256    M::V: DefaultDenseMatrix,
257    FwdLS: LinearSolver<M>,
258    BwdLS: LinearSolver<M>,
259    Tag: OdeSolverMethodTag<M, Eqn> + 'static,
260    Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
261        + 'static
262        + diffsol::OdeEquationsImplicitAdjoint<M = M, T = M::T, V = M::V, C = M::C>,
263    for<'b> &'b M::V: VectorRef<M::V>,
264    for<'b> &'b M: MatrixRef<M>,
265{
266    // use the parameters of the forward pass
267    problem.eqn.set_params(&checkpoint.params);
268    // TODO: can we avoid cloning here? Adjoint equations require ownership of the checkpointing segments so maybe not
269    // unless we can change the adjoint equations to take references to the checkpointing segments instead
270    let checkpointing = checkpoint.checkpointing.clone();
271    let soln = Solution::new_dense(t_eval.to_vec())?;
272
273    // we will only consider a single output g for now, so nout_override is 1
274    let dgdu_eval = [dgdu_eval];
275    match backwards_method {
276        OdeSolverType::Bdf => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
277            M,
278            Eqn,
279            FwdLS,
280            BwdLS,
281            Tag,
282            BdfTag,
283        >(problem, &soln, checkpointing, &dgdu_eval, Some(1)),
284        OdeSolverType::Esdirk34 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
285            M,
286            Eqn,
287            FwdLS,
288            BwdLS,
289            Tag,
290            Esdirk34Tag,
291        >(problem, &soln, checkpointing, &dgdu_eval, Some(1)),
292        OdeSolverType::TrBdf2 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
293            M,
294            Eqn,
295            FwdLS,
296            BwdLS,
297            Tag,
298            TrBdf2Tag,
299        >(problem, &soln, checkpointing, &dgdu_eval, Some(1)),
300        OdeSolverType::Tsit45 => solve_adjoint_bkwds_with_fwd_bkwd_tag::<
301            M,
302            Eqn,
303            FwdLS,
304            BwdLS,
305            Tag,
306            Tsit45Tag,
307        >(problem, &soln, checkpointing, &dgdu_eval, Some(1)),
308    }
309}
310
311fn solve_adjoint_bkwds_with_fwd_bkwd_tag<'solver, M, Eqn, FwdLS, BwdLS, FwdTag, BwdTag>(
312    problem: &'solver mut OdeSolverProblem<Eqn>,
313    soln: &Solution<M::V>,
314    mut checkpointing: CheckpointingPath<Eqn, FwdTag::State>,
315    dgdu_eval: &[&<M::V as DefaultDenseMatrix>::M],
316    nout_override: Option<usize>,
317) -> Result<M::V, DiffsolError>
318where
319    M: Matrix<T: Scalar> + DefaultSolver + 'solver,
320    M::V: DefaultDenseMatrix,
321    FwdLS: LinearSolver<M>,
322    BwdLS: LinearSolver<M>,
323    FwdTag: OdeSolverMethodTag<M, Eqn>,
324    BwdTag: OdeSolverMethodTag<M, Eqn>,
325    Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
326        + 'solver
327        + diffsol::OdeEquationsImplicitAdjoint<M = M, T = M::T, V = M::V, C = M::C>,
328    for<'b> &'b M::V: VectorRef<M::V>,
329    for<'b> &'b M: MatrixRef<M>,
330{
331    let checkpointing_len = checkpointing.len();
332    if checkpointing_len == 0 {
333        return Err(ode_solver_error!(
334            Other,
335            "Adjoint backward pass requires at least one checkpointing segment"
336        ));
337    }
338
339    let t_eval = if dgdu_eval.is_empty() {
340        &[]
341    } else {
342        soln.ts.as_slice()
343    };
344
345    let current_checkpointing = checkpointing
346        .pop()
347        .ok_or_else(|| ode_solver_error!(Other, "Adjoint backward pass returned no state"))?;
348    let model_index = checkpointing
349        .last()
350        .map(|segment| {
351            segment
352                .terminal_reset_root_idx()
353                .expect("Missing reset root index")
354        })
355        .unwrap_or(0);
356    problem.eqn_mut().set_model_index(model_index);
357    let fwd_solver = FwdTag::uninitialised_solver::<FwdLS>(&*problem)?;
358    let mut adjoint = BwdTag::solver_adjoint::<BwdLS, _>(
359        &*problem,
360        vec![current_checkpointing],
361        Some(fwd_solver),
362        nout_override,
363    )?;
364    loop {
365        let (mut state, adjoint_checkpointing) =
366            adjoint.solve_adjoint_backwards_pass(t_eval, dgdu_eval)?;
367        let Some(previous_checkpointing) = checkpointing.pop() else {
368            // the gradient has one lane for each (batch lane, output channel) pair
369            return Ok(state.into_common().sg);
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
406#[allow(clippy::type_complexity)]
407fn solve_vjp_fwd_with_tag<M, Eqn, LS, Tag>(
408    problem: &mut OdeSolverProblem<Eqn>,
409    t_eval: &[M::T],
410    method: OdeSolverType,
411    linear_solver: LinearSolverType,
412) -> Result<(Solution<M::V>, Box<dyn VjpCheckpoint>), DiffsolError>
413where
414    M: Matrix<T: Scalar> + DefaultSolver + 'static,
415    M::V: DefaultDenseMatrix + Send,
416    LS: LinearSolver<M>,
417    Eqn: OdeEquationsImplicitAdjoint<M = M, T = M::T, V = M::V, C = M::C> + 'static,
418    Tag: OdeSolverMethodTag<M, Eqn> + 'static,
419    Tag::State: Send,
420{
421    // keep the parameters of all the lanes, so the backward pass can use them again
422    let mut params = M::V::zeros(problem.eqn.nparams(), problem.context().clone());
423    problem.eqn.get_params(&mut params);
424    let (ys, tape) = Tag::solver_vjp::<LS>(problem)?.solve_dense_vjp_fwd(t_eval)?;
425    let mut soln = Solution::new_dense(t_eval.to_vec())?;
426    soln.ys = ys;
427    soln.stop_reason = Some(OdeSolverStopReason::TstopReached);
428    Ok((
429        soln,
430        Box::new(VjpCheckpointData::<M, Eqn, Tag>::new(
431            tape,
432            params,
433            method,
434            linear_solver,
435        )),
436    ))
437}
438
439fn solve_vjp_bkwd_with_tag<M, Eqn, LS, Tag>(
440    problem: &mut OdeSolverProblem<Eqn>,
441    checkpoint: &dyn VjpCheckpoint,
442    dgdu_eval: &<M::V as DefaultDenseMatrix>::M,
443) -> Result<M::V, DiffsolError>
444where
445    M: Matrix<T: Scalar> + DefaultSolver + 'static,
446    M::V: DefaultDenseMatrix,
447    LS: LinearSolver<M>,
448    Eqn: OdeEquationsImplicitAdjoint<M = M, T = M::T, V = M::V, C = M::C> + 'static,
449    Tag: OdeSolverMethodTag<M, Eqn> + 'static,
450{
451    let data = checkpoint.data::<M, Eqn, Tag>()?;
452    // use the parameters of the forward pass
453    problem.eqn.set_params(&data.params);
454    // a solver at the end of the forward pass, where the backward pass starts, so that its
455    // Newton matrix is a good preconditioner for the first transposed solves
456    let mut solver = Tag::solver_vjp_with_state::<LS>(problem, data.tape.end().clone())?;
457    solver.solve_dense_vjp_bkwd(&data.tape, dgdu_eval)
458}
459
460impl OdeSolverType {
461    pub(crate) fn solve<M, Eqn, LS>(
462        &self,
463        problem: &mut OdeSolverProblem<Eqn>,
464        final_time: M::T,
465    ) -> Result<Solution<M::V>, DiffsolError>
466    where
467        M: Matrix<T: Scalar> + DefaultSolver,
468        M::V: DefaultDenseMatrix,
469        LS: LinearSolver<M>,
470        for<'b> &'b M::V: VectorRef<M::V>,
471        for<'b> &'b M: MatrixRef<M>,
472        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>,
473    {
474        match self {
475            OdeSolverType::Bdf => {
476                solve_with_tag::<M, Eqn, LS, BdfTag>(problem, Solution::new(final_time))
477            }
478            OdeSolverType::Esdirk34 => {
479                solve_with_tag::<M, Eqn, LS, Esdirk34Tag>(problem, Solution::new(final_time))
480            }
481            OdeSolverType::TrBdf2 => {
482                solve_with_tag::<M, Eqn, LS, TrBdf2Tag>(problem, Solution::new(final_time))
483            }
484            OdeSolverType::Tsit45 => {
485                solve_with_tag::<M, Eqn, LS, Tsit45Tag>(problem, Solution::new(final_time))
486            }
487        }
488    }
489
490    pub(crate) fn solve_dense<M, Eqn, LS>(
491        &self,
492        problem: &mut OdeSolverProblem<Eqn>,
493        t_eval: &[M::T],
494    ) -> Result<Solution<M::V>, DiffsolError>
495    where
496        M: Matrix<T: Scalar> + DefaultSolver,
497        M::V: DefaultDenseMatrix,
498        LS: LinearSolver<M>,
499        for<'b> &'b M::V: VectorRef<M::V>,
500        for<'b> &'b M: MatrixRef<M>,
501        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>,
502    {
503        match self {
504            OdeSolverType::Bdf => {
505                solve_with_tag::<M, Eqn, LS, BdfTag>(problem, Solution::new_dense(t_eval.to_vec())?)
506            }
507            OdeSolverType::Esdirk34 => solve_with_tag::<M, Eqn, LS, Esdirk34Tag>(
508                problem,
509                Solution::new_dense(t_eval.to_vec())?,
510            ),
511            OdeSolverType::TrBdf2 => solve_with_tag::<M, Eqn, LS, TrBdf2Tag>(
512                problem,
513                Solution::new_dense(t_eval.to_vec())?,
514            ),
515            OdeSolverType::Tsit45 => solve_with_tag::<M, Eqn, LS, Tsit45Tag>(
516                problem,
517                Solution::new_dense(t_eval.to_vec())?,
518            ),
519        }
520    }
521
522    #[allow(clippy::type_complexity)]
523    pub(crate) fn solve_fwd_sens<M, Eqn, LS>(
524        &self,
525        problem: &mut OdeSolverProblem<Eqn>,
526        t_eval: &[M::T],
527    ) -> Result<Solution<M::V>, DiffsolError>
528    where
529        M: Matrix<T: Scalar> + DefaultSolver,
530        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>,
531        M::V: DefaultDenseMatrix,
532        LS: LinearSolver<M>,
533        for<'b> &'b M::V: VectorRef<M::V>,
534        for<'b> &'b M: MatrixRef<M>,
535    {
536        match self {
537            OdeSolverType::Bdf => solve_fwd_sens_with_tag::<M, Eqn, LS, BdfTag>(problem, t_eval),
538            OdeSolverType::Esdirk34 => {
539                solve_fwd_sens_with_tag::<M, Eqn, LS, Esdirk34Tag>(problem, t_eval)
540            }
541            OdeSolverType::TrBdf2 => {
542                solve_fwd_sens_with_tag::<M, Eqn, LS, TrBdf2Tag>(problem, t_eval)
543            }
544            OdeSolverType::Tsit45 => {
545                solve_fwd_sens_with_tag::<M, Eqn, LS, Tsit45Tag>(problem, t_eval)
546            }
547        }
548    }
549
550    #[allow(clippy::type_complexity)]
551    pub(crate) fn solve_adjoint_fwd<M, Eqn, LS>(
552        &self,
553        problem: &mut OdeSolverProblem<Eqn>,
554        t_eval: &[M::T],
555        linear_solver: LinearSolverType,
556    ) -> Result<(Solution<M::V>, Box<dyn AdjointCheckpoint>), DiffsolError>
557    where
558        M: Matrix<T: Scalar> + DefaultSolver + 'static,
559        M::V: DefaultDenseMatrix,
560        LS: LinearSolver<M>,
561        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C> + 'static,
562        for<'b> &'b M::V: VectorRef<M::V>,
563        for<'b> &'b M: MatrixRef<M>,
564    {
565        match self {
566            OdeSolverType::Bdf => solve_adjoint_fwd_with_tag::<M, Eqn, LS, BdfTag>(
567                problem,
568                t_eval,
569                *self,
570                linear_solver,
571            ),
572            OdeSolverType::Esdirk34 => solve_adjoint_fwd_with_tag::<M, Eqn, LS, Esdirk34Tag>(
573                problem,
574                t_eval,
575                *self,
576                linear_solver,
577            ),
578            OdeSolverType::TrBdf2 => solve_adjoint_fwd_with_tag::<M, Eqn, LS, TrBdf2Tag>(
579                problem,
580                t_eval,
581                *self,
582                linear_solver,
583            ),
584            OdeSolverType::Tsit45 => solve_adjoint_fwd_with_tag::<M, Eqn, LS, Tsit45Tag>(
585                problem,
586                t_eval,
587                *self,
588                linear_solver,
589            ),
590        }
591    }
592
593    pub(crate) fn solve_continuous_adjoint<M, Eqn, LS>(
594        &self,
595        problem: &mut OdeSolverProblem<Eqn>,
596        final_time: M::T,
597    ) -> Result<(M::V, M::V), DiffsolError>
598    where
599        M: Matrix<T: Scalar> + DefaultSolver + 'static,
600        M::V: DefaultDenseMatrix,
601        LS: LinearSolver<M>,
602        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
603            + 'static
604            + diffsol::OdeEquationsImplicitAdjoint,
605        for<'b> &'b M::V: VectorRef<M::V>,
606        for<'b> &'b M: MatrixRef<M>,
607    {
608        match self {
609            OdeSolverType::Bdf => {
610                solve_continuous_adjoint_with_tag::<M, Eqn, LS, BdfTag>(problem, final_time, *self)
611            }
612            OdeSolverType::Esdirk34 => {
613                solve_continuous_adjoint_with_tag::<M, Eqn, LS, Esdirk34Tag>(
614                    problem, final_time, *self,
615                )
616            }
617            OdeSolverType::TrBdf2 => solve_continuous_adjoint_with_tag::<M, Eqn, LS, TrBdf2Tag>(
618                problem, final_time, *self,
619            ),
620            OdeSolverType::Tsit45 => solve_continuous_adjoint_with_tag::<M, Eqn, LS, Tsit45Tag>(
621                problem, final_time, *self,
622            ),
623        }
624    }
625
626    pub(crate) fn solve_adjoint_bkwd<M, Eqn, BwdLS>(
627        &self,
628        problem: &mut OdeSolverProblem<Eqn>,
629        checkpoint: &dyn AdjointCheckpoint,
630        dgdu_eval: &<M::V as DefaultDenseMatrix>::M,
631        t_eval: &[M::T],
632    ) -> Result<M::V, DiffsolError>
633    where
634        M: Matrix<T: Scalar> + DefaultSolver + LuValidator<M> + KluValidator<M> + 'static,
635        M::V: DefaultDenseMatrix,
636        BwdLS: LinearSolver<M>,
637        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
638            + 'static
639            + diffsol::OdeEquationsImplicitAdjoint,
640        for<'b> &'b M::V: VectorRef<M::V>,
641        for<'b> &'b M: MatrixRef<M>,
642    {
643        match checkpoint.method() {
644            OdeSolverType::Bdf => {
645                let data = checkpoint.data::<M, Eqn, BdfTag>()?;
646                match data.linear_solver() {
647                    LinearSolverType::Default => {
648                        solve_adjoint_bkwds_with_fwd_tag::<
649                            M,
650                            Eqn,
651                            <M as DefaultSolver>::LS,
652                            BwdLS,
653                            BdfTag,
654                        >(problem, data, *self, dgdu_eval, t_eval)
655                    }
656                    LinearSolverType::Lu => {
657                        solve_adjoint_bkwds_with_fwd_tag::<
658                            M,
659                            Eqn,
660                            <M as LuValidator<M>>::LS,
661                            BwdLS,
662                            BdfTag,
663                        >(problem, data, *self, dgdu_eval, t_eval)
664                    }
665                    LinearSolverType::Klu => {
666                        solve_adjoint_bkwds_with_fwd_tag::<
667                            M,
668                            Eqn,
669                            <M as KluValidator<M>>::LS,
670                            BwdLS,
671                            BdfTag,
672                        >(problem, data, *self, dgdu_eval, t_eval)
673                    }
674                }
675            }
676            OdeSolverType::Esdirk34 => {
677                let data = checkpoint.data::<M, Eqn, Esdirk34Tag>()?;
678                match data.linear_solver() {
679                    LinearSolverType::Default => {
680                        solve_adjoint_bkwds_with_fwd_tag::<
681                            M,
682                            Eqn,
683                            <M as DefaultSolver>::LS,
684                            BwdLS,
685                            Esdirk34Tag,
686                        >(problem, data, *self, dgdu_eval, t_eval)
687                    }
688                    LinearSolverType::Lu => {
689                        solve_adjoint_bkwds_with_fwd_tag::<
690                            M,
691                            Eqn,
692                            <M as LuValidator<M>>::LS,
693                            BwdLS,
694                            Esdirk34Tag,
695                        >(problem, data, *self, dgdu_eval, t_eval)
696                    }
697                    LinearSolverType::Klu => {
698                        solve_adjoint_bkwds_with_fwd_tag::<
699                            M,
700                            Eqn,
701                            <M as KluValidator<M>>::LS,
702                            BwdLS,
703                            Esdirk34Tag,
704                        >(problem, data, *self, dgdu_eval, t_eval)
705                    }
706                }
707            }
708            OdeSolverType::TrBdf2 => {
709                let data = checkpoint.data::<M, Eqn, TrBdf2Tag>()?;
710                match data.linear_solver() {
711                    LinearSolverType::Default => {
712                        solve_adjoint_bkwds_with_fwd_tag::<
713                            M,
714                            Eqn,
715                            <M as DefaultSolver>::LS,
716                            BwdLS,
717                            TrBdf2Tag,
718                        >(problem, data, *self, dgdu_eval, t_eval)
719                    }
720                    LinearSolverType::Lu => {
721                        solve_adjoint_bkwds_with_fwd_tag::<
722                            M,
723                            Eqn,
724                            <M as LuValidator<M>>::LS,
725                            BwdLS,
726                            TrBdf2Tag,
727                        >(problem, data, *self, dgdu_eval, t_eval)
728                    }
729                    LinearSolverType::Klu => {
730                        solve_adjoint_bkwds_with_fwd_tag::<
731                            M,
732                            Eqn,
733                            <M as KluValidator<M>>::LS,
734                            BwdLS,
735                            TrBdf2Tag,
736                        >(problem, data, *self, dgdu_eval, t_eval)
737                    }
738                }
739            }
740            OdeSolverType::Tsit45 => {
741                let data = checkpoint.data::<M, Eqn, Tsit45Tag>()?;
742                match data.linear_solver() {
743                    LinearSolverType::Default => {
744                        solve_adjoint_bkwds_with_fwd_tag::<
745                            M,
746                            Eqn,
747                            <M as DefaultSolver>::LS,
748                            BwdLS,
749                            Tsit45Tag,
750                        >(problem, data, *self, dgdu_eval, t_eval)
751                    }
752                    LinearSolverType::Lu => {
753                        solve_adjoint_bkwds_with_fwd_tag::<
754                            M,
755                            Eqn,
756                            <M as LuValidator<M>>::LS,
757                            BwdLS,
758                            Tsit45Tag,
759                        >(problem, data, *self, dgdu_eval, t_eval)
760                    }
761                    LinearSolverType::Klu => {
762                        solve_adjoint_bkwds_with_fwd_tag::<
763                            M,
764                            Eqn,
765                            <M as KluValidator<M>>::LS,
766                            BwdLS,
767                            Tsit45Tag,
768                        >(problem, data, *self, dgdu_eval, t_eval)
769                    }
770                }
771            }
772        }
773    }
774
775    #[allow(clippy::type_complexity)]
776    pub(crate) fn solve_vjp_fwd<M, Eqn, LS>(
777        &self,
778        problem: &mut OdeSolverProblem<Eqn>,
779        t_eval: &[M::T],
780        linear_solver: LinearSolverType,
781    ) -> Result<(Solution<M::V>, Box<dyn VjpCheckpoint>), DiffsolError>
782    where
783        M: Matrix<T: Scalar> + DefaultSolver + 'static,
784        M::V: DefaultDenseMatrix + Send,
785        <M::V as DefaultDenseMatrix>::M: Send,
786        LS: LinearSolver<M>,
787        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
788            + OdeEquationsImplicitAdjoint
789            + 'static,
790        for<'b> &'b M::V: VectorRef<M::V>,
791        for<'b> &'b M: MatrixRef<M>,
792    {
793        match self {
794            OdeSolverType::Bdf => {
795                solve_vjp_fwd_with_tag::<M, Eqn, LS, BdfTag>(problem, t_eval, *self, linear_solver)
796            }
797            OdeSolverType::Esdirk34 => solve_vjp_fwd_with_tag::<M, Eqn, LS, Esdirk34Tag>(
798                problem,
799                t_eval,
800                *self,
801                linear_solver,
802            ),
803            OdeSolverType::TrBdf2 => solve_vjp_fwd_with_tag::<M, Eqn, LS, TrBdf2Tag>(
804                problem,
805                t_eval,
806                *self,
807                linear_solver,
808            ),
809            OdeSolverType::Tsit45 => solve_vjp_fwd_with_tag::<M, Eqn, LS, Tsit45Tag>(
810                problem,
811                t_eval,
812                *self,
813                linear_solver,
814            ),
815        }
816    }
817
818    /// The backward pass of a VJP. It uses the method and the linear solver of the forward
819    /// pass that made `checkpoint`, because the tape holds states of that solver.
820    pub(crate) fn solve_vjp_bkwd<M, Eqn>(
821        problem: &mut OdeSolverProblem<Eqn>,
822        checkpoint: &dyn VjpCheckpoint,
823        dgdu_eval: &<M::V as DefaultDenseMatrix>::M,
824    ) -> Result<M::V, DiffsolError>
825    where
826        M: Matrix<T: Scalar> + DefaultSolver + LuValidator<M> + KluValidator<M> + 'static,
827        M::V: DefaultDenseMatrix,
828        Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
829            + OdeEquationsImplicitAdjoint
830            + 'static,
831        for<'b> &'b M::V: VectorRef<M::V>,
832        for<'b> &'b M: MatrixRef<M>,
833    {
834        macro_rules! with_ls {
835            ($tag:ty) => {
836                match checkpoint.linear_solver() {
837                    LinearSolverType::Default => {
838                        solve_vjp_bkwd_with_tag::<M, Eqn, <M as DefaultSolver>::LS, $tag>(
839                            problem, checkpoint, dgdu_eval,
840                        )
841                    }
842                    LinearSolverType::Lu => {
843                        solve_vjp_bkwd_with_tag::<M, Eqn, <M as LuValidator<M>>::LS, $tag>(
844                            problem, checkpoint, dgdu_eval,
845                        )
846                    }
847                    LinearSolverType::Klu => {
848                        solve_vjp_bkwd_with_tag::<M, Eqn, <M as KluValidator<M>>::LS, $tag>(
849                            problem, checkpoint, dgdu_eval,
850                        )
851                    }
852                }
853            };
854        }
855        match checkpoint.method() {
856            OdeSolverType::Bdf => with_ls!(BdfTag),
857            OdeSolverType::Esdirk34 => with_ls!(Esdirk34Tag),
858            OdeSolverType::TrBdf2 => with_ls!(TrBdf2Tag),
859            OdeSolverType::Tsit45 => with_ls!(Tsit45Tag),
860        }
861    }
862}
863
864#[cfg(all(test, any(feature = "diffsl-cranelift", feature = "diffsl-llvm")))]
865mod tests {
866    use diffsol::{
867        CodegenModuleCompile, CodegenModuleJit, DefaultSolver, DenseMatrix, OdeBuilder,
868        OdeSolverProblem,
869    };
870    #[cfg(feature = "diffsl-llvm")]
871    use diffsol::{DefaultDenseMatrix, Matrix, MatrixCommon, Op, Vector};
872
873    #[cfg(feature = "diffsl-llvm")]
874    use crate::linear_solver_type::LinearSolverType;
875    use crate::test_support::{
876        assert_close, hybrid_logistic_diffsl_code, hybrid_logistic_state, logistic_diffsl_code,
877        logistic_state, LOGISTIC_X0,
878    };
879    #[cfg(feature = "diffsl-llvm")]
880    use crate::test_support::{hybrid_logistic_state_dr, logistic_state_dr};
881    use crate::valid_linear_solver::LuValidator;
882
883    use super::OdeSolverType;
884
885    type M = diffsol::NalgebraMat<f64>;
886
887    fn build_problem<CG>(code: &str) -> OdeSolverProblem<diffsol::DiffSl<M, CG>>
888    where
889        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
890    {
891        OdeBuilder::<M>::new()
892            .p([2.0])
893            .rtol(1e-6)
894            .atol([1e-6])
895            .sens_rtol(1e-6)
896            .sens_atol([1e-6])
897            .build_from_diffsl::<CG>(code)
898            .unwrap()
899    }
900
901    fn assert_dense_solution_matches_expected(
902        soln: &diffsol::Solution<diffsol::NalgebraVec<f64>>,
903        t_eval: &[f64],
904        expected: impl Fn(f64) -> f64,
905    ) {
906        assert_eq!(soln.ts, t_eval);
907        for (i, &t) in t_eval.iter().enumerate() {
908            assert_close(
909                soln.ys.get_index(0, i),
910                expected(t),
911                5e-4,
912                &format!("solution[{i}]"),
913            );
914        }
915    }
916
917    fn test_all_solver_variants<CG>()
918    where
919        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
920    {
921        let t_eval = [0.25, 0.5, 1.0];
922        for method in [
923            OdeSolverType::Bdf,
924            OdeSolverType::Esdirk34,
925            OdeSolverType::TrBdf2,
926            OdeSolverType::Tsit45,
927        ] {
928            let mut problem = build_problem::<CG>(logistic_diffsl_code());
929            let soln = method
930                .solve::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(&mut problem, 1.0)
931                .unwrap();
932            assert_close(*soln.ts.last().unwrap(), 1.0, 5e-4, "solve final time");
933            assert_close(
934                soln.ys.get_index(0, soln.ts.len() - 1),
935                logistic_state(LOGISTIC_X0, 2.0, 1.0),
936                5e-4,
937                "solve final value",
938            );
939
940            let mut problem = build_problem::<CG>(logistic_diffsl_code());
941            let soln = method
942                .solve_dense::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(
943                    &mut problem,
944                    &t_eval,
945                )
946                .unwrap();
947            assert_dense_solution_matches_expected(&soln, &t_eval, |t| {
948                logistic_state(LOGISTIC_X0, 2.0, t)
949            });
950        }
951    }
952
953    fn test_all_hybrid_solver_variants<CG>()
954    where
955        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
956    {
957        let t_eval = [0.5, 1.0, 1.25, 1.5, 2.0];
958        for method in [
959            OdeSolverType::Bdf,
960            OdeSolverType::Esdirk34,
961            OdeSolverType::TrBdf2,
962            OdeSolverType::Tsit45,
963        ] {
964            let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
965            let soln = method
966                .solve::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(&mut problem, 2.0)
967                .unwrap();
968            assert_close(*soln.ts.last().unwrap(), 2.0, 5e-4, "hybrid final time");
969            assert_close(
970                soln.ys.get_index(0, soln.ts.len() - 1),
971                hybrid_logistic_state(2.0, 2.0),
972                5e-4,
973                "hybrid final value",
974            );
975
976            let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
977            let soln = method
978                .solve_dense::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(
979                    &mut problem,
980                    &t_eval,
981                )
982                .unwrap();
983            assert_dense_solution_matches_expected(&soln, &t_eval, |t| {
984                hybrid_logistic_state(2.0, t)
985            });
986        }
987    }
988
989    fn test_all_solver_variants_with_lu<CG>()
990    where
991        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
992    {
993        let t_eval = [0.25, 0.5, 1.0];
994        for method in [
995            OdeSolverType::Bdf,
996            OdeSolverType::Esdirk34,
997            OdeSolverType::TrBdf2,
998            OdeSolverType::Tsit45,
999        ] {
1000            let mut problem = build_problem::<CG>(logistic_diffsl_code());
1001            let soln = method
1002                .solve::<M, diffsol::DiffSl<M, CG>, <M as LuValidator<M>>::LS>(&mut problem, 1.0)
1003                .unwrap();
1004            assert_close(*soln.ts.last().unwrap(), 1.0, 5e-4, "lu solve final time");
1005
1006            let mut problem = build_problem::<CG>(logistic_diffsl_code());
1007            let soln = method
1008                .solve_dense::<M, diffsol::DiffSl<M, CG>, <M as LuValidator<M>>::LS>(
1009                    &mut problem,
1010                    &t_eval,
1011                )
1012                .unwrap();
1013            assert_dense_solution_matches_expected(&soln, &t_eval, |t| {
1014                logistic_state(LOGISTIC_X0, 2.0, t)
1015            });
1016        }
1017    }
1018
1019    fn test_all_hybrid_solver_variants_with_lu<CG>()
1020    where
1021        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
1022    {
1023        let t_eval = [0.5, 1.0, 1.25, 1.5, 2.0];
1024        for method in [
1025            OdeSolverType::Bdf,
1026            OdeSolverType::Esdirk34,
1027            OdeSolverType::TrBdf2,
1028            OdeSolverType::Tsit45,
1029        ] {
1030            let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
1031            let soln = method
1032                .solve::<M, diffsol::DiffSl<M, CG>, <M as LuValidator<M>>::LS>(&mut problem, 2.0)
1033                .unwrap();
1034            assert_close(*soln.ts.last().unwrap(), 2.0, 5e-4, "lu hybrid final time");
1035
1036            let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
1037            let soln = method
1038                .solve_dense::<M, diffsol::DiffSl<M, CG>, <M as LuValidator<M>>::LS>(
1039                    &mut problem,
1040                    &t_eval,
1041                )
1042                .unwrap();
1043            assert_dense_solution_matches_expected(&soln, &t_eval, |t| {
1044                hybrid_logistic_state(2.0, t)
1045            });
1046        }
1047    }
1048
1049    fn assert_direct_hybrid_restart_path_for_method<CG>(method: OdeSolverType)
1050    where
1051        CG: diffsol::CodegenModule + CodegenModuleJit + CodegenModuleCompile,
1052    {
1053        let t_eval = [0.5, 1.0, 1.25, 1.5, 2.0];
1054
1055        let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
1056        let soln = method
1057            .solve::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(&mut problem, 2.0)
1058            .unwrap();
1059        assert_close(
1060            *soln.ts.last().unwrap(),
1061            2.0,
1062            5e-4,
1063            "direct hybrid restart final time",
1064        );
1065        assert_close(
1066            soln.ys.get_index(0, soln.ts.len() - 1),
1067            hybrid_logistic_state(2.0, 2.0),
1068            5e-4,
1069            "direct hybrid restart final value",
1070        );
1071
1072        let mut problem = build_problem::<CG>(hybrid_logistic_diffsl_code());
1073        let soln = method
1074            .solve_dense::<M, diffsol::DiffSl<M, CG>, <M as DefaultSolver>::LS>(
1075                &mut problem,
1076                &t_eval,
1077            )
1078            .unwrap();
1079        assert_dense_solution_matches_expected(&soln, &t_eval, |t| hybrid_logistic_state(2.0, t));
1080    }
1081
1082    #[cfg(feature = "diffsl-llvm")]
1083    fn test_all_sensitivity_solver_variants() {
1084        let t_eval = [0.25, 0.5, 1.0];
1085        for method in [
1086            OdeSolverType::Bdf,
1087            OdeSolverType::Esdirk34,
1088            OdeSolverType::TrBdf2,
1089            OdeSolverType::Tsit45,
1090        ] {
1091            let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
1092            let soln = method
1093                .solve_fwd_sens::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1094                    &mut problem,
1095                    &t_eval,
1096                )
1097                .unwrap();
1098            for (i, &t) in t_eval.iter().enumerate() {
1099                assert_close(
1100                    soln.y_sens[0].get_index(0, i),
1101                    logistic_state_dr(LOGISTIC_X0, 2.0, t),
1102                    5e-4,
1103                    &format!("fwd_sens[{i}]"),
1104                );
1105            }
1106
1107            let mut problem = build_problem::<diffsol::LlvmModule>(hybrid_logistic_diffsl_code());
1108            let soln = method
1109                .solve_fwd_sens::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1110                    &mut problem,
1111                    &t_eval,
1112                )
1113                .unwrap();
1114            for (i, &t) in t_eval.iter().enumerate() {
1115                assert_close(
1116                    soln.y_sens[0].get_index(0, i),
1117                    hybrid_logistic_state_dr(2.0, t),
1118                    5e-4,
1119                    &format!("hybrid_fwd_sens[{i}]"),
1120                );
1121            }
1122        }
1123    }
1124
1125    #[cfg(feature = "diffsl-llvm")]
1126    fn test_lu_sensitivity_and_adjoint_solver_variants() {
1127        let t_eval = [0.25, 0.5, 1.0];
1128        for method in [
1129            OdeSolverType::Bdf,
1130            OdeSolverType::Esdirk34,
1131            OdeSolverType::TrBdf2,
1132            OdeSolverType::Tsit45,
1133        ] {
1134            let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
1135            let soln = method
1136                .solve_fwd_sens::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as LuValidator<M>>::LS>(
1137                    &mut problem,
1138                    &t_eval,
1139                )
1140                .unwrap();
1141            for (i, &t) in t_eval.iter().enumerate() {
1142                assert_close(
1143                    soln.y_sens[0].get_index(0, i),
1144                    logistic_state_dr(LOGISTIC_X0, 2.0, t),
1145                    5e-4,
1146                    &format!("lu fwd_sens[{i}]"),
1147                );
1148            }
1149        }
1150
1151        let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
1152        let adjoint_t_eval = [0.0, 0.25, 0.5, 1.0];
1153        let (soln, checkpoint) = OdeSolverType::Bdf
1154            .solve_adjoint_fwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as LuValidator<M>>::LS>(
1155                &mut problem,
1156                &adjoint_t_eval,
1157                LinearSolverType::Lu,
1158            )
1159            .unwrap();
1160        let dgdu = <<M as MatrixCommon>::V as DefaultDenseMatrix>::M::zeros(
1161            problem.eqn.nout(),
1162            soln.ts.len(),
1163            problem.context().to_owned(),
1164        );
1165        let gradient = OdeSolverType::TrBdf2
1166            .solve_adjoint_bkwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as LuValidator<M>>::LS>(
1167                &mut problem,
1168                checkpoint.as_ref(),
1169                &dgdu,
1170                &soln.ts,
1171            )
1172            .unwrap();
1173        assert_eq!(gradient.len(), 1);
1174        assert!(gradient.get_index(0).is_finite());
1175    }
1176
1177    #[cfg(feature = "diffsl-llvm")]
1178    fn test_direct_hybrid_sensitivity_restart_paths() {
1179        let t_eval = [0.5, 1.0, 2.5, 3.0, 4.5];
1180        for method in [
1181            OdeSolverType::Esdirk34,
1182            OdeSolverType::TrBdf2,
1183            OdeSolverType::Tsit45,
1184        ] {
1185            let mut problem = build_problem::<diffsol::LlvmModule>(hybrid_logistic_diffsl_code());
1186            let soln = method
1187                .solve_fwd_sens::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1188                    &mut problem,
1189                    &t_eval,
1190                )
1191                .unwrap();
1192            for (i, &t) in t_eval.iter().enumerate() {
1193                assert_close(
1194                    soln.ys.get_index(0, i),
1195                    hybrid_logistic_state(2.0, t),
1196                    5e-4,
1197                    &format!("direct hybrid value[{i}]"),
1198                );
1199                assert_close(
1200                    soln.y_sens[0].get_index(0, i),
1201                    hybrid_logistic_state_dr(2.0, t),
1202                    5e-4,
1203                    &format!("direct hybrid fwd sens[{i}]"),
1204                );
1205            }
1206        }
1207    }
1208
1209    #[cfg(feature = "diffsl-llvm")]
1210    fn test_adjoint_backwards_methods_and_klu_branch() {
1211        for backwards_method in [OdeSolverType::Esdirk34, OdeSolverType::TrBdf2] {
1212            let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
1213            let t_eval = [0.0, 0.25, 0.5, 1.0];
1214            let (soln, checkpoint) = OdeSolverType::Bdf
1215                .solve_adjoint_fwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1216                    &mut problem,
1217                    &t_eval,
1218                        LinearSolverType::Default,
1219                )
1220                .unwrap();
1221            let dgdu = <<M as MatrixCommon>::V as DefaultDenseMatrix>::M::zeros(
1222                problem.eqn.nout(),
1223                soln.ts.len(),
1224                problem.context().to_owned(),
1225            );
1226            let gradient = backwards_method
1227                .solve_adjoint_bkwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as crate::valid_linear_solver::KluValidator<M>>::LS>(
1228                    &mut problem,
1229                    checkpoint.as_ref(),
1230                    &dgdu,
1231                    &soln.ts,
1232                )
1233                .unwrap();
1234            assert_eq!(gradient.len(), 1);
1235            assert!(gradient.get_index(0).is_finite());
1236        }
1237
1238        let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
1239        let t_eval = [0.0, 0.25, 0.5, 1.0];
1240        let (soln, checkpoint) = OdeSolverType::Tsit45
1241            .solve_adjoint_fwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1242                &mut problem,
1243                &t_eval,
1244                LinearSolverType::Default,
1245            )
1246            .unwrap();
1247        let dgdu = <<M as MatrixCommon>::V as DefaultDenseMatrix>::M::zeros(
1248            problem.eqn.nout(),
1249            soln.ts.len(),
1250            problem.context().to_owned(),
1251        );
1252        let gradient = OdeSolverType::Bdf
1253            .solve_adjoint_bkwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1254                &mut problem,
1255                checkpoint.as_ref(),
1256                &dgdu,
1257                &soln.ts,
1258            )
1259            .unwrap();
1260        assert_eq!(gradient.len(), 1);
1261        assert!(gradient.get_index(0).is_finite());
1262    }
1263
1264    /// A VJP with `dg/du = 1` at each time gives `sum_i dy/dr(t_i)`, for each method and
1265    /// linear solver.
1266    #[cfg(feature = "diffsl-llvm")]
1267    fn test_all_vjp_solver_variants() {
1268        type Eqn = diffsol::DiffSl<M, diffsol::LlvmModule>;
1269        let t_eval = [0.0, 0.25, 0.5, 1.0];
1270        let expected: f64 = t_eval
1271            .iter()
1272            .map(|&t| logistic_state_dr(LOGISTIC_X0, 2.0, t))
1273            .sum();
1274        for method in [
1275            OdeSolverType::Bdf,
1276            OdeSolverType::Esdirk34,
1277            OdeSolverType::TrBdf2,
1278            OdeSolverType::Tsit45,
1279        ] {
1280            for linear_solver in [LinearSolverType::Default, LinearSolverType::Lu] {
1281                let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
1282                let (soln, checkpoint) = match linear_solver {
1283                    LinearSolverType::Lu => method
1284                        .solve_vjp_fwd::<M, Eqn, <M as LuValidator<M>>::LS>(
1285                            &mut problem,
1286                            &t_eval,
1287                            linear_solver,
1288                        ),
1289                    _ => method.solve_vjp_fwd::<M, Eqn, <M as DefaultSolver>::LS>(
1290                        &mut problem,
1291                        &t_eval,
1292                        linear_solver,
1293                    ),
1294                }
1295                .unwrap();
1296                for (i, &t) in t_eval.iter().enumerate() {
1297                    assert_close(
1298                        soln.ys.get_index(0, i),
1299                        logistic_state(LOGISTIC_X0, 2.0, t),
1300                        5e-4,
1301                        &format!("{method:?}/{linear_solver:?} y[{i}]"),
1302                    );
1303                }
1304                let nout = problem.eqn.nout();
1305                let dgdu = <<M as MatrixCommon>::V as DefaultDenseMatrix>::M::from_vec(
1306                    nout,
1307                    t_eval.len(),
1308                    vec![1.0; nout * t_eval.len()],
1309                    problem.context().to_owned(),
1310                );
1311                let gradient = OdeSolverType::solve_vjp_bkwd::<M, Eqn>(
1312                    &mut problem,
1313                    checkpoint.as_ref(),
1314                    &dgdu,
1315                )
1316                .unwrap();
1317                assert_eq!(gradient.len(), 1);
1318                assert_close(
1319                    gradient.get_index(0),
1320                    expected,
1321                    5e-4,
1322                    &format!("{method:?}/{linear_solver:?} VJP gradient"),
1323                );
1324            }
1325        }
1326    }
1327
1328    #[cfg(feature = "diffsl-llvm")]
1329    fn test_all_adjoint_solver_variants() {
1330        let t_eval = [0.0, 0.25, 0.5, 1.0];
1331        for method in [
1332            OdeSolverType::Bdf,
1333            OdeSolverType::Esdirk34,
1334            OdeSolverType::TrBdf2,
1335        ] {
1336            let mut problem = build_problem::<diffsol::LlvmModule>(logistic_diffsl_code());
1337            let (soln, checkpoint) = method
1338                .solve_adjoint_fwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1339                    &mut problem,
1340                    &t_eval,
1341                        LinearSolverType::Default,
1342                )
1343                .unwrap();
1344            let dgdu = <<M as MatrixCommon>::V as DefaultDenseMatrix>::M::zeros(
1345                problem.eqn.nout(),
1346                soln.ts.len(),
1347                problem.context().to_owned(),
1348            );
1349            let gradient = OdeSolverType::Bdf
1350                .solve_adjoint_bkwd::<M, diffsol::DiffSl<M, diffsol::LlvmModule>, <M as DefaultSolver>::LS>(
1351                    &mut problem,
1352                    checkpoint.as_ref(),
1353                    &dgdu,
1354                    &soln.ts,
1355                )
1356                .unwrap();
1357            assert_eq!(gradient.len(), 1);
1358            assert!(gradient.get_index(0).is_finite());
1359        }
1360    }
1361
1362    #[cfg(feature = "diffsl-cranelift")]
1363    #[test]
1364    fn runtime_dispatch_solves_all_variants_for_cranelift() {
1365        test_all_solver_variants::<diffsol::CraneliftJitModule>();
1366        test_all_solver_variants_with_lu::<diffsol::CraneliftJitModule>();
1367    }
1368
1369    #[cfg(feature = "diffsl-cranelift")]
1370    #[test]
1371    fn runtime_dispatch_solves_all_hybrid_variants_for_cranelift() {
1372        test_all_hybrid_solver_variants::<diffsol::CraneliftJitModule>();
1373        test_all_hybrid_solver_variants_with_lu::<diffsol::CraneliftJitModule>();
1374        assert_direct_hybrid_restart_path_for_method::<diffsol::CraneliftJitModule>(
1375            OdeSolverType::Esdirk34,
1376        );
1377        assert_direct_hybrid_restart_path_for_method::<diffsol::CraneliftJitModule>(
1378            OdeSolverType::TrBdf2,
1379        );
1380        assert_direct_hybrid_restart_path_for_method::<diffsol::CraneliftJitModule>(
1381            OdeSolverType::Tsit45,
1382        );
1383    }
1384
1385    #[cfg(feature = "diffsl-llvm")]
1386    #[test]
1387    fn runtime_dispatch_solves_all_variants_for_llvm() {
1388        test_all_solver_variants::<diffsol::LlvmModule>();
1389        test_all_solver_variants_with_lu::<diffsol::LlvmModule>();
1390    }
1391
1392    #[cfg(feature = "diffsl-llvm")]
1393    #[test]
1394    fn runtime_dispatch_solves_all_hybrid_variants_for_llvm() {
1395        test_all_hybrid_solver_variants::<diffsol::LlvmModule>();
1396        test_all_hybrid_solver_variants_with_lu::<diffsol::LlvmModule>();
1397        assert_direct_hybrid_restart_path_for_method::<diffsol::LlvmModule>(
1398            OdeSolverType::Esdirk34,
1399        );
1400        assert_direct_hybrid_restart_path_for_method::<diffsol::LlvmModule>(OdeSolverType::TrBdf2);
1401        assert_direct_hybrid_restart_path_for_method::<diffsol::LlvmModule>(OdeSolverType::Tsit45);
1402    }
1403
1404    #[cfg(feature = "diffsl-llvm")]
1405    #[test]
1406    fn runtime_dispatch_solves_all_forward_sensitivity_variants_for_llvm() {
1407        test_all_sensitivity_solver_variants();
1408        test_lu_sensitivity_and_adjoint_solver_variants();
1409        test_direct_hybrid_sensitivity_restart_paths();
1410    }
1411
1412    #[cfg(feature = "diffsl-llvm")]
1413    #[test]
1414    fn runtime_dispatch_solves_all_adjoint_variants_for_llvm() {
1415        test_all_adjoint_solver_variants();
1416        test_adjoint_backwards_methods_and_klu_branch();
1417    }
1418
1419    #[cfg(feature = "diffsl-llvm")]
1420    #[test]
1421    fn runtime_dispatch_solves_all_vjp_variants_for_llvm() {
1422        test_all_vjp_solver_variants();
1423    }
1424}
diffsol_c/ode_solver_type.rs

diffsol_c/
ode_solver_type.rs
