diffsol_c/
solve.rs

1// Delegate solver types selected at runtime in Host to concrete solver types
2// in Rust.
3
4use std::sync::Arc;
5
6use crate::adjoint_checkpoint::AdjointCheckpointWrapper;
7use crate::error::DiffsolRtError;
8use crate::scalar_type::Scalar;
9use crate::tensor::{column_major_values, FromTensor, Tensor, TensorLayout, TensorView, ToTensor};
10use crate::vjp_checkpoint::VjpCheckpointWrapper;
11use crate::{
12    generate_ic_option_accessors, generate_ode_option_accessors, generate_option_accessors,
13    option_value_from_store, option_value_to_store,
14};
15use crate::{generate_trait_ic_option_accessors, generate_trait_ode_option_accessors};
16use diffsol::OdeBuilder;
17use diffsol::{
18    error::DiffsolError,
19    matrix::{Matrix, MatrixRef},
20    ConstantOp, Context, DefaultDenseMatrix, DefaultSolver, NonLinearOp, NonLinearOpJacobian,
21    OdeEquations, OdeEquationsImplicitAdjoint, OdeEquationsImplicitSens, OdeSolverProblem, Vector,
22    VectorRef, VectorView,
23};
24use num_traits::{FromPrimitive, ToPrimitive}; // for from_f64 and to_f64
25use paste::paste;
26
27use crate::solve_eqn::{EqnRequirement, SolveEqn};
28use crate::{
29    linear_solver_type::LinearSolverType, matrix_type::MatrixType, ode_solver_type::OdeSolverType,
30};
31use crate::{
32    matrix_type::MatrixKind,
33    valid_linear_solver::{validate_linear_solver, KluValidator, LuValidator},
34};
35
36// Each matrix type implements PySolve as bridge between diffsol and Host
37
38use crate::solution::Solution;
39pub(crate) type SolveResult = Result<Arc<dyn Solution>, DiffsolRtError>;
40
41pub(crate) trait Solve: Send {
42    fn matrix_type(&self) -> MatrixType;
43    fn nstates(&self) -> usize;
44    fn nparams(&self) -> usize;
45    fn nout(&self) -> usize;
46    fn has_stop(&self) -> bool;
47    /// The number of lanes (independent copies of the system) that the solver integrates
48    /// together.
49    fn nbatch(&self) -> usize;
50
51    fn rhs(
52        &mut self,
53        params: TensorView<'_>,
54        t: f64,
55        y: TensorView<'_>,
56    ) -> Result<Tensor, DiffsolRtError>;
57
58    fn rhs_jac_mul(
59        &mut self,
60        params: TensorView<'_>,
61        t: f64,
62        y: TensorView<'_>,
63        v: TensorView<'_>,
64    ) -> Result<Tensor, DiffsolRtError>;
65
66    fn y0(&mut self, params: TensorView<'_>) -> Result<Tensor, DiffsolRtError>;
67
68    fn check(&self, linear_solver: LinearSolverType) -> Result<(), DiffsolRtError>;
69    fn set_rtol(&mut self, rtol: f64);
70    fn rtol(&self) -> f64;
71    fn set_atol(&mut self, atol: f64);
72    fn atol(&self) -> f64;
73    fn set_t0(&mut self, t0: f64);
74    fn t0(&self) -> f64;
75    fn set_h0(&mut self, h0: f64);
76    fn h0(&self) -> f64;
77    fn set_integrate_out(&mut self, integrate_out: bool);
78    fn integrate_out(&self) -> bool;
79    fn set_sens_rtol(&mut self, sens_rtol: Option<f64>);
80    fn sens_rtol(&self) -> Option<f64>;
81    fn set_sens_atol(&mut self, sens_atol: Option<f64>);
82    fn sens_atol(&self) -> Option<f64>;
83    fn set_out_rtol(&mut self, out_rtol: Option<f64>);
84    fn out_rtol(&self) -> Option<f64>;
85    fn set_out_atol(&mut self, out_atol: Option<f64>);
86    fn out_atol(&self) -> Option<f64>;
87    fn set_param_rtol(&mut self, param_rtol: Option<f64>);
88    fn param_rtol(&self) -> Option<f64>;
89    fn set_param_atol(&mut self, param_atol: Option<f64>);
90    fn param_atol(&self) -> Option<f64>;
91    fn serialized_diffsl(&self) -> Result<Vec<u8>, DiffsolRtError>;
92
93    // New API: solution object support
94    fn solve(
95        &mut self,
96        method: OdeSolverType,
97        linear_solver: LinearSolverType,
98        params: TensorView<'_>,
99        final_time: f64,
100    ) -> SolveResult;
101
102    fn solve_dense(
103        &mut self,
104        method: OdeSolverType,
105        linear_solver: LinearSolverType,
106        params: TensorView<'_>,
107        t_eval: TensorView<'_>,
108    ) -> SolveResult;
109
110    fn solve_fwd_sens(
111        &mut self,
112        method: OdeSolverType,
113        linear_solver: LinearSolverType,
114        params: TensorView<'_>,
115        t_eval: TensorView<'_>,
116    ) -> SolveResult;
117
118    fn solve_continuous_adjoint(
119        &mut self,
120        method: OdeSolverType,
121        linear_solver: LinearSolverType,
122        params: TensorView<'_>,
123        final_time: f64,
124    ) -> Result<(Tensor, Tensor), DiffsolRtError>;
125
126    fn solve_adjoint_fwd(
127        &mut self,
128        method: OdeSolverType,
129        linear_solver: LinearSolverType,
130        params: TensorView<'_>,
131        t_eval: TensorView<'_>,
132    ) -> Result<(Arc<dyn Solution>, AdjointCheckpointWrapper), DiffsolRtError>;
133
134    fn solve_adjoint_bkwd(
135        &mut self,
136        method: OdeSolverType,
137        linear_solver: LinearSolverType,
138        checkpoint: &AdjointCheckpointWrapper,
139        t_eval: TensorView<'_>,
140        dgdu_eval: TensorView<'_>,
141    ) -> Result<Tensor, DiffsolRtError>;
142
143    fn solve_vjp_fwd(
144        &mut self,
145        method: OdeSolverType,
146        linear_solver: LinearSolverType,
147        params: TensorView<'_>,
148        t_eval: TensorView<'_>,
149    ) -> Result<(Arc<dyn Solution>, VjpCheckpointWrapper), DiffsolRtError>;
150
151    fn solve_vjp_bkwd(
152        &mut self,
153        checkpoint: &VjpCheckpointWrapper,
154        dgdu_eval: TensorView<'_>,
155    ) -> Result<Tensor, DiffsolRtError>;
156
157    generate_trait_ic_option_accessors! {
158        use_linesearch: bool,
159        max_linesearch_iterations: usize,
160        max_newton_iterations: usize,
161        max_linear_solver_setups: usize,
162        step_reduction_factor: f64,
163        armijo_constant: f64
164    }
165    generate_trait_ode_option_accessors! {
166        max_nonlinear_solver_iterations: usize,
167        max_error_test_failures: usize,
168        max_nonlinear_solver_failures: usize,
169        min_timestep: f64,
170        update_jacobian_after_steps: usize,
171        update_rhs_jacobian_after_steps: usize,
172        threshold_to_update_jacobian: f64,
173        threshold_to_update_rhs_jacobian: f64,
174        derivative_solve_tolerance: f64,
175        derivative_solve_max_iter: usize
176    }
177}
178pub(crate) struct GenericSolve<M, Eqn>
179where
180    M: Matrix<T: Scalar>,
181    M::V: Vector,
182    Eqn: OdeEquations<M = M, T = M::T, V = M::V, C = M::C>,
183{
184    pub(crate) problem: OdeSolverProblem<Eqn>,
185}
186
187/// Read the times `t_eval` as `T` values. All the lanes use the same times, so `t_eval` is a 1D
188/// `[nt]` array, or a 2D `[nt, 1]` array with one lane.
189fn read_times<T: Scalar>(times: TensorView<'_>) -> Result<Vec<T>, DiffsolRtError> {
190    match times.shape() {
191        [_] | [_, 1] => column_major_values(times).map_err(|e| e.named("t_eval")),
192        shape => {
193            let message = format!("Expected an [nt] or an [nt, 1] array, got the shape {shape:?}");
194            Err(DiffsolRtError::from(DiffsolError::Other(message)).named("t_eval"))
195        }
196    }
197}
198
199/// The first value of lane 0 of a tolerance vector, or `None` if it has no values.
200///
201/// This API sets all values of a tolerance vector to one value, in all lanes, so the first
202/// value is the value that was set.
203fn first_lane_value<V: Vector<T: Scalar>>(v: &V) -> Option<f64> {
204    (v.len() > 0).then(|| v.get_batch(0).get_index(0).to_f64().unwrap())
205}
206
207impl<M, Eqn> GenericSolve<M, Eqn>
208where
209    M: Matrix<T: Scalar>,
210    M::V: Vector + DefaultDenseMatrix,
211    Eqn: OdeEquations<M = M, T = M::T, V = M::V, C = M::C> + SolveEqn,
212{
213    pub(crate) fn from_eqn(eqn: Eqn) -> Result<Self, DiffsolRtError> {
214        // the builder needs one block of parameters for each lane
215        let ctx = eqn.context().clone();
216        let default_p = vec![0.0; eqn.nparams() * ctx.nbatch()];
217        let problem = OdeBuilder::<M>::new()
218            .context(ctx)
219            .p(default_p)
220            .build_from_eqn(eqn)?;
221        Ok(GenericSolve { problem })
222    }
223
224    /// Finish a solve with `result`.
225    ///
226    /// If the solve succeeded, discard the errors that the equations recorded. The solver
227    /// recovered from them, for example by rejecting a step.
228    ///
229    /// If the solve failed, return the error of the solver. If the equations recorded an error
230    /// during the solve, add it to the message. It is often the cause of the failure, but it can
231    /// also be an earlier failure that the solver recovered from.
232    fn finish<R, E: Into<DiffsolRtError>>(
233        &self,
234        result: Result<R, E>,
235    ) -> Result<R, DiffsolRtError> {
236        let recorded = self.problem.eqn.take_error();
237        result.map_err(|err| {
238            let err = err.into();
239            match recorded {
240                Some(recorded) => DiffsolError::Other(format!(
241                    "{err} (a callback failed during the solve: {recorded})"
242                ))
243                .into(),
244                None => err,
245            }
246        })
247    }
248
249    /// Convert `tensor` (the argument `name`) to a diffsol value with the lanes of the problem.
250    /// Each lane must have the shape `lane_shape`. The errors start with `name`.
251    fn read_input<T: FromTensor<Context = M::C>>(
252        &self,
253        tensor: TensorView<'_>,
254        lane_shape: &[usize],
255        name: &str,
256    ) -> Result<T, DiffsolRtError> {
257        self.read_input_with_context(tensor, self.problem.context(), lane_shape, name)
258    }
259
260    /// As [Self::read_input], but with the lanes of `ctx`.
261    fn read_input_with_context<T: FromTensor<Context = M::C>>(
262        &self,
263        tensor: TensorView<'_>,
264        ctx: &M::C,
265        lane_shape: &[usize],
266        name: &str,
267    ) -> Result<T, DiffsolRtError> {
268        // the conversion checks the number of axes and the number of lanes
269        let value = T::from_tensor(tensor, ctx).map_err(|e| e.named(name))?;
270        // the last axis is the lane axis
271        let got = &tensor.shape()[..tensor.shape().len() - 1];
272        if got != lane_shape {
273            let message = match (lane_shape, got) {
274                ([n], [m]) => format!("Expected {n} values in each lane, got {m}"),
275                ([r, c], [gr, gc]) => format!(
276                    "Expected {r} rows and {c} columns in each lane, got {gr} rows and {gc} columns"
277                ),
278                _ => format!("Expected the shape {lane_shape:?} in each lane, got {got:?}"),
279            };
280            return Err(DiffsolRtError::from(DiffsolError::Other(message)).named(name));
281        }
282        Ok(value)
283    }
284
285    /// Convert an adjoint gradient to a Tensor with the shape `[nparams, nchannels, nbatch]`.
286    ///
287    /// The gradient has one lane for each (batch lane, output channel) pair, with lane
288    /// `b * nchannels + j` for batch lane `b` and channel `j`.
289    fn gradient_to_tensor(&self, gradient: M::V) -> Tensor
290    where
291        M::V: ToTensor,
292    {
293        // `[nparams, nbatch * nchannels]`, with the lane axis split into `[nchannels, nbatch]`
294        let tensor = gradient.to_tensor();
295        let nbatch = self.problem.context().nbatch();
296        let nchannels = tensor.shape()[1] / nbatch;
297        let shape = [tensor.shape()[0], nchannels, nbatch];
298        let strides = [
299            tensor.strides()[0],
300            tensor.strides()[1],
301            nchannels * tensor.strides()[1],
302        ];
303        tensor.with_layout(&shape, &strides)
304    }
305
306    /// Return the error recorded by the equations during the last evaluation, if any.
307    fn take_eqn_error(&self) -> Result<(), DiffsolRtError> {
308        match self.problem.eqn.take_error() {
309            Some(message) => Err(DiffsolError::Other(message).into()),
310            None => Ok(()),
311        }
312    }
313
314    pub(crate) fn setup_problem(&mut self, params: TensorView<'_>) -> Result<(), DiffsolRtError>
315    where
316        M::V: FromTensor<Context = M::C>,
317    {
318        self.reset_eval_state();
319        let nparams = self.problem.eqn.nparams();
320        let params = self.read_input(params, &[nparams], "params")?;
321        // a single lane of parameters broadcasts to all the lanes of the problem
322        self.problem.eqn.set_params(&params);
323        Ok(())
324    }
325
326    /// Discard any error left over from an earlier evaluation, and reset the model index to 0
327    /// for hybrid solves.
328    fn reset_eval_state(&mut self) {
329        let _ = self.problem.eqn.take_error();
330        self.problem.eqn.set_model_index(0);
331    }
332}
333
334impl<M, Eqn> Solve for GenericSolve<M, Eqn>
335where
336    M: Matrix<T: Scalar + ToPrimitive>
337        + DefaultSolver
338        + LuValidator<M>
339        + KluValidator<M>
340        + MatrixKind,
341    Eqn: OdeEquationsImplicitSens<M = M, T = M::T, V = M::V, C = M::C>
342        + OdeEquationsImplicitAdjoint<M = M, T = M::T, V = M::V, C = M::C>
343        + SolveEqn
344        + Send
345        + 'static,
346    <M::V as DefaultDenseMatrix>::M: TensorLayout + FromTensor<Context = M::C>,
347    M::V: TensorLayout + FromTensor<Context = M::C>,
348    M::V: DefaultDenseMatrix + Send + Sync + 'static,
349    <M::V as DefaultDenseMatrix>::M: Send + Sync,
350    for<'b> &'b M::V: VectorRef<M::V>,
351    for<'b> &'b M: MatrixRef<M>,
352{
353    fn matrix_type(&self) -> MatrixType {
354        MatrixType::from_diffsol::<M>()
355    }
356
357    fn nstates(&self) -> usize {
358        self.problem.eqn.nstates()
359    }
360
361    fn nparams(&self) -> usize {
362        self.problem.eqn.nparams()
363    }
364
365    fn nout(&self) -> usize {
366        self.problem.eqn.nout()
367    }
368
369    fn has_stop(&self) -> bool {
370        self.problem.eqn.root().is_some()
371    }
372
373    fn nbatch(&self) -> usize {
374        self.problem.context().nbatch()
375    }
376
377    fn check(&self, linear_solver: LinearSolverType) -> Result<(), DiffsolRtError> {
378        validate_linear_solver::<M>(linear_solver)
379    }
380
381    fn set_atol(&mut self, atol: f64) {
382        self.problem.atol.fill(M::T::from_f64(atol).unwrap());
383    }
384
385    fn atol(&self) -> f64 {
386        first_lane_value(&self.problem.atol).expect("atol has no values")
387    }
388
389    fn serialized_diffsl(&self) -> Result<Vec<u8>, DiffsolRtError> {
390        self.problem.eqn.serialized()
391    }
392
393    fn set_rtol(&mut self, rtol: f64) {
394        self.problem.rtol = M::T::from_f64(rtol).unwrap();
395    }
396
397    fn rtol(&self) -> f64 {
398        self.problem.rtol.to_f64().unwrap()
399    }
400
401    fn set_t0(&mut self, t0: f64) {
402        self.problem.t0 = M::T::from_f64(t0).unwrap();
403    }
404
405    fn t0(&self) -> f64 {
406        self.problem.t0.to_f64().unwrap()
407    }
408
409    fn set_h0(&mut self, h0: f64) {
410        self.problem.h0 = M::T::from_f64(h0).unwrap();
411    }
412
413    fn h0(&self) -> f64 {
414        self.problem.h0.to_f64().unwrap()
415    }
416
417    fn set_integrate_out(&mut self, integrate_out: bool) {
418        self.problem.integrate_out = integrate_out;
419    }
420
421    fn integrate_out(&self) -> bool {
422        self.problem.integrate_out
423    }
424
425    fn set_sens_rtol(&mut self, sens_rtol: Option<f64>) {
426        self.problem.sens_rtol = sens_rtol.map(|value| M::T::from_f64(value).unwrap());
427    }
428
429    fn sens_rtol(&self) -> Option<f64> {
430        self.problem.sens_rtol.map(|value| value.to_f64().unwrap())
431    }
432
433    fn set_sens_atol(&mut self, sens_atol: Option<f64>) {
434        // one batch lane per (problem batch, parameter), as the augmented state is laid out
435        self.problem.sens_atol = sens_atol.map(|value| {
436            let nparams = self.problem.eqn.nparams();
437            let ctx = self.problem.context();
438            let aug_ctx = ctx
439                .clone_with_nbatch(ctx.nbatch() * nparams)
440                .expect("failed to create the augmented context");
441            M::V::from_element(
442                self.problem.eqn.nstates(),
443                M::T::from_f64(value).unwrap(),
444                aug_ctx,
445            )
446        });
447    }
448
449    fn sens_atol(&self) -> Option<f64> {
450        self.problem.sens_atol.as_ref().and_then(first_lane_value)
451    }
452
453    fn set_out_rtol(&mut self, out_rtol: Option<f64>) {
454        self.problem.out_rtol = out_rtol.map(|value| M::T::from_f64(value).unwrap());
455    }
456
457    fn out_rtol(&self) -> Option<f64> {
458        self.problem.out_rtol.map(|value| value.to_f64().unwrap())
459    }
460
461    fn set_out_atol(&mut self, out_atol: Option<f64>) {
462        self.problem.out_atol = out_atol.map(|value| {
463            let len = if self.problem.eqn.out().is_some() {
464                self.problem.eqn.nout()
465            } else {
466                self.problem.eqn.nstates()
467            };
468            M::V::from_element(
469                len,
470                M::T::from_f64(value).unwrap(),
471                self.problem.context().clone(),
472            )
473        });
474    }
475
476    fn out_atol(&self) -> Option<f64> {
477        self.problem.out_atol.as_ref().and_then(first_lane_value)
478    }
479
480    fn set_param_rtol(&mut self, param_rtol: Option<f64>) {
481        self.problem.param_rtol = param_rtol.map(|value| M::T::from_f64(value).unwrap());
482    }
483
484    fn param_rtol(&self) -> Option<f64> {
485        self.problem.param_rtol.map(|value| value.to_f64().unwrap())
486    }
487
488    fn set_param_atol(&mut self, param_atol: Option<f64>) {
489        self.problem.param_atol = param_atol.map(|value| {
490            M::V::from_element(
491                self.problem.eqn.nparams(),
492                M::T::from_f64(value).unwrap(),
493                self.problem.context().clone(),
494            )
495        });
496    }
497
498    fn param_atol(&self) -> Option<f64> {
499        self.problem.param_atol.as_ref().and_then(first_lane_value)
500    }
501
502    generate_ic_option_accessors! {
503        use_linesearch: bool,
504        max_linesearch_iterations: usize,
505        max_newton_iterations: usize,
506        max_linear_solver_setups: usize,
507        step_reduction_factor: f64,
508        armijo_constant: f64
509    }
510
511    generate_ode_option_accessors! {
512        max_nonlinear_solver_iterations: usize,
513        max_error_test_failures: usize,
514        max_nonlinear_solver_failures: usize,
515        min_timestep: f64,
516        update_jacobian_after_steps: usize,
517        update_rhs_jacobian_after_steps: usize,
518        threshold_to_update_jacobian: f64,
519        threshold_to_update_rhs_jacobian: f64,
520        derivative_solve_tolerance: f64,
521        derivative_solve_max_iter: usize
522    }
523
524    fn y0(&mut self, params: TensorView<'_>) -> Result<Tensor, DiffsolRtError> {
525        self.setup_problem(params)?;
526        let n = self.problem.eqn.nstates();
527        let mut y0 = M::V::zeros(n, self.problem.context().clone());
528        let t0 = self.problem.t0;
529        self.problem.eqn.init().call_inplace(t0, &mut y0);
530        self.take_eqn_error()?;
531        Ok(y0.to_tensor())
532    }
533
534    fn rhs(
535        &mut self,
536        params: TensorView<'_>,
537        t: f64,
538        y: TensorView<'_>,
539    ) -> Result<Tensor, DiffsolRtError> {
540        let n = self.problem.eqn.nstates();
541        let y = self.read_input(y, &[n], "y")?;
542        self.setup_problem(params)?;
543        let mut dydt = M::V::zeros(n, self.problem.context().clone());
544        self.problem
545            .eqn
546            .rhs()
547            .call_inplace(&y, M::T::from_f64(t).unwrap(), &mut dydt);
548        self.take_eqn_error()?;
549        Ok(dydt.to_tensor())
550    }
551
552    fn rhs_jac_mul(
553        &mut self,
554        params: TensorView<'_>,
555        t: f64,
556        y: TensorView<'_>,
557        v: TensorView<'_>,
558    ) -> Result<Tensor, DiffsolRtError> {
559        let n = self.problem.eqn.nstates();
560        let y = self.read_input(y, &[n], "y")?;
561        let v = self.read_input(v, &[n], "v")?;
562        self.setup_problem(params)?;
563        let mut dydt = M::V::zeros(n, self.problem.context().clone());
564        self.problem
565            .eqn
566            .rhs()
567            .jac_mul_inplace(&y, M::T::from_f64(t).unwrap(), &v, &mut dydt);
568        self.take_eqn_error()?;
569        Ok(dydt.to_tensor())
570    }
571
572    fn solve(
573        &mut self,
574        method: OdeSolverType,
575        linear_solver: LinearSolverType,
576        params: TensorView<'_>,
577        final_time: f64,
578    ) -> SolveResult {
579        self.problem
580            .eqn
581            .check_requirement(EqnRequirement::Solve(method))?;
582        self.setup_problem(params)?;
583        let final_time = M::T::from_f64(final_time).unwrap();
584        let soln = match linear_solver {
585            LinearSolverType::Default => {
586                method.solve::<M, Eqn, <M as DefaultSolver>::LS>(&mut self.problem, final_time)
587            }
588            LinearSolverType::Lu => {
589                method.solve::<M, Eqn, <M as LuValidator<M>>::LS>(&mut self.problem, final_time)
590            }
591            LinearSolverType::Klu => {
592                method.solve::<M, Eqn, <M as KluValidator<M>>::LS>(&mut self.problem, final_time)
593            }
594        };
595        Ok(Arc::new(self.finish(soln)?))
596    }
597
598    fn solve_dense(
599        &mut self,
600        method: OdeSolverType,
601        linear_solver: LinearSolverType,
602        params: TensorView<'_>,
603        t_eval: TensorView<'_>,
604    ) -> SolveResult {
605        self.problem
606            .eqn
607            .check_requirement(EqnRequirement::Solve(method))?;
608        self.setup_problem(params)?;
609
610        let t_eval: Vec<M::T> = read_times(t_eval)?;
611        let soln =
612            match linear_solver {
613                LinearSolverType::Default => method
614                    .solve_dense::<M, Eqn, <M as DefaultSolver>::LS>(&mut self.problem, &t_eval),
615                LinearSolverType::Lu => method
616                    .solve_dense::<M, Eqn, <M as LuValidator<M>>::LS>(&mut self.problem, &t_eval),
617                LinearSolverType::Klu => method
618                    .solve_dense::<M, Eqn, <M as KluValidator<M>>::LS>(&mut self.problem, &t_eval),
619            };
620        Ok(Arc::new(self.finish(soln)?))
621    }
622
623    fn solve_fwd_sens(
624        &mut self,
625        method: OdeSolverType,
626        linear_solver: LinearSolverType,
627        params: TensorView<'_>,
628        t_eval: TensorView<'_>,
629    ) -> SolveResult {
630        self.problem
631            .eqn
632            .check_requirement(EqnRequirement::FwdSens)?;
633        self.setup_problem(params)?;
634
635        let t_eval: Vec<M::T> = read_times(t_eval)?;
636        let soln = match linear_solver {
637            LinearSolverType::Default => method
638                .solve_fwd_sens::<M, Eqn, <M as DefaultSolver>::LS>(&mut self.problem, &t_eval),
639            LinearSolverType::Lu => method
640                .solve_fwd_sens::<M, Eqn, <M as LuValidator<M>>::LS>(&mut self.problem, &t_eval),
641            LinearSolverType::Klu => method
642                .solve_fwd_sens::<M, Eqn, <M as KluValidator<M>>::LS>(&mut self.problem, &t_eval),
643        };
644        Ok(Arc::new(self.finish(soln)?))
645    }
646
647    fn solve_continuous_adjoint(
648        &mut self,
649        method: OdeSolverType,
650        linear_solver: LinearSolverType,
651        params: TensorView<'_>,
652        final_time: f64,
653    ) -> Result<(Tensor, Tensor), DiffsolRtError> {
654        self.problem
655            .eqn
656            .check_requirement(EqnRequirement::Adjoint(method))?;
657        self.setup_problem(params)?;
658
659        let final_time = M::T::from_f64(final_time).unwrap();
660        let integrate_out = self.problem.integrate_out;
661        self.problem.integrate_out = true;
662        let result = match linear_solver {
663            LinearSolverType::Default => method
664                .solve_continuous_adjoint::<M, Eqn, <M as DefaultSolver>::LS>(
665                    &mut self.problem,
666                    final_time,
667                ),
668            LinearSolverType::Lu => method
669                .solve_continuous_adjoint::<M, Eqn, <M as LuValidator<M>>::LS>(
670                    &mut self.problem,
671                    final_time,
672                ),
673            LinearSolverType::Klu => method
674                .solve_continuous_adjoint::<M, Eqn, <M as KluValidator<M>>::LS>(
675                    &mut self.problem,
676                    final_time,
677                ),
678        };
679        self.problem.integrate_out = integrate_out;
680        let (integral, gradient) = self.finish(result)?;
681        Ok((integral.to_tensor(), self.gradient_to_tensor(gradient)))
682    }
683
684    fn solve_adjoint_fwd(
685        &mut self,
686        method: OdeSolverType,
687        linear_solver: LinearSolverType,
688        params: TensorView<'_>,
689        t_eval: TensorView<'_>,
690    ) -> Result<(Arc<dyn Solution>, AdjointCheckpointWrapper), DiffsolRtError> {
691        self.problem
692            .eqn
693            .check_requirement(EqnRequirement::Adjoint(method))?;
694        self.setup_problem(params)?;
695
696        let t_eval: Vec<M::T> = read_times(t_eval)?;
697        let result = match linear_solver {
698            LinearSolverType::Default => method
699                .solve_adjoint_fwd::<M, Eqn, <M as DefaultSolver>::LS>(
700                    &mut self.problem,
701                    &t_eval,
702                    linear_solver,
703                ),
704            LinearSolverType::Lu => method.solve_adjoint_fwd::<M, Eqn, <M as LuValidator<M>>::LS>(
705                &mut self.problem,
706                &t_eval,
707                linear_solver,
708            ),
709            LinearSolverType::Klu => method
710                .solve_adjoint_fwd::<M, Eqn, <M as KluValidator<M>>::LS>(
711                    &mut self.problem,
712                    &t_eval,
713                    linear_solver,
714                ),
715        };
716        let (soln, checkpoint) = self.finish(result)?;
717        Ok((Arc::new(soln), AdjointCheckpointWrapper::new(checkpoint)))
718    }
719
720    fn solve_adjoint_bkwd(
721        &mut self,
722        method: OdeSolverType,
723        linear_solver: LinearSolverType,
724        checkpoint: &AdjointCheckpointWrapper,
725        t_eval: TensorView<'_>,
726        dgdu_eval: TensorView<'_>,
727    ) -> Result<Tensor, DiffsolRtError> {
728        self.problem
729            .eqn
730            .check_requirement(EqnRequirement::Adjoint(method))?;
731        let checkpoint = checkpoint.guard()?;
732        // the backward pass sets the parameters that the checkpoint holds
733        self.reset_eval_state();
734
735        let t_eval: Vec<M::T> = read_times(t_eval)?;
736        let lane_shape = [self.problem.eqn.nout(), t_eval.len()];
737        let dgdu_eval: <M::V as DefaultDenseMatrix>::M =
738            self.read_input(dgdu_eval, &lane_shape, "dgdu_eval")?;
739
740        let result = match linear_solver {
741            LinearSolverType::Default => method
742                .solve_adjoint_bkwd::<M, Eqn, <M as DefaultSolver>::LS>(
743                    &mut self.problem,
744                    checkpoint.as_ref(),
745                    &dgdu_eval,
746                    &t_eval,
747                ),
748            LinearSolverType::Lu => method.solve_adjoint_bkwd::<M, Eqn, <M as LuValidator<M>>::LS>(
749                &mut self.problem,
750                checkpoint.as_ref(),
751                &dgdu_eval,
752                &t_eval,
753            ),
754            LinearSolverType::Klu => method
755                .solve_adjoint_bkwd::<M, Eqn, <M as KluValidator<M>>::LS>(
756                    &mut self.problem,
757                    checkpoint.as_ref(),
758                    &dgdu_eval,
759                    &t_eval,
760                ),
761        };
762        let gradient = self.finish(result)?;
763        Ok(self.gradient_to_tensor(gradient))
764    }
765
766    fn solve_vjp_fwd(
767        &mut self,
768        method: OdeSolverType,
769        linear_solver: LinearSolverType,
770        params: TensorView<'_>,
771        t_eval: TensorView<'_>,
772    ) -> Result<(Arc<dyn Solution>, VjpCheckpointWrapper), DiffsolRtError> {
773        self.problem
774            .eqn
775            .check_requirement(EqnRequirement::Vjp(method))?;
776        if self.problem.eqn.root().is_some() {
777            return Err(DiffsolError::Other(
778                "A VJP solve does not support models with a stop or reset condition".to_string(),
779            )
780            .into());
781        }
782        if self.problem.integrate_out {
783            return Err(DiffsolError::Other(
784                "A VJP solve does not support integrating the output".to_string(),
785            )
786            .into());
787        }
788        self.setup_problem(params)?;
789
790        let t_eval: Vec<M::T> = read_times(t_eval)?;
791        let result = match linear_solver {
792            LinearSolverType::Default => method.solve_vjp_fwd::<M, Eqn, <M as DefaultSolver>::LS>(
793                &mut self.problem,
794                &t_eval,
795                linear_solver,
796            ),
797            LinearSolverType::Lu => method.solve_vjp_fwd::<M, Eqn, <M as LuValidator<M>>::LS>(
798                &mut self.problem,
799                &t_eval,
800                linear_solver,
801            ),
802            LinearSolverType::Klu => method.solve_vjp_fwd::<M, Eqn, <M as KluValidator<M>>::LS>(
803                &mut self.problem,
804                &t_eval,
805                linear_solver,
806            ),
807        };
808        let (soln, checkpoint) = self.finish(result)?;
809        Ok((Arc::new(soln), VjpCheckpointWrapper::new(checkpoint)))
810    }
811
812    fn solve_vjp_bkwd(
813        &mut self,
814        checkpoint: &VjpCheckpointWrapper,
815        dgdu_eval: TensorView<'_>,
816    ) -> Result<Tensor, DiffsolRtError> {
817        let checkpoint = checkpoint.guard()?;
818        self.problem
819            .eqn
820            .check_requirement(EqnRequirement::Vjp(checkpoint.method()))?;
821        // the backward pass sets the parameters that the checkpoint holds
822        self.reset_eval_state();
823
824        // `nbatch * ndirections` lanes, lane `b * ndirections + k` for direction `k` of batch
825        // lane `b`, or 1 lane for the same values in all the batch lanes
826        let ctx = self.problem.context();
827        let nbatch = ctx.nbatch();
828        let nlanes = dgdu_eval.shape().last().copied().unwrap_or(0);
829        let ctx = if nlanes == 1 {
830            ctx.clone()
831        } else if nlanes > 0 && nlanes.is_multiple_of(nbatch) {
832            ctx.clone_with_nbatch(nlanes).map_err(DiffsolError::from)?
833        } else {
834            return Err(DiffsolRtError::from(DiffsolError::Other(format!(
835                "Expected 1 lane or a multiple of {nbatch} lanes (nbatch * ndirections), got {nlanes}"
836            )))
837            .named("dgdu_eval"));
838        };
839        let lane_shape = [self.problem.eqn.nout(), checkpoint.nt()];
840        let dgdu_eval: <M::V as DefaultDenseMatrix>::M =
841            self.read_input_with_context(dgdu_eval, &ctx, &lane_shape, "dgdu_eval")?;
842        let result = OdeSolverType::solve_vjp_bkwd::<M, Eqn>(
843            &mut self.problem,
844            checkpoint.as_ref(),
845            &dgdu_eval,
846        );
847        let gradient = self.finish(result)?;
848        Ok(self.gradient_to_tensor(gradient))
849    }
850}
diffsol_c/solve.rs

diffsol_c/
solve.rs
