diffsol_c/
vjp_checkpoint.rs

1use std::{
2    any::Any,
3    sync::{Arc, Mutex, MutexGuard},
4};
5
6use diffsol::{DefaultSolver, DenseVjpTape, DiffsolError, OdeEquationsImplicit};
7
8use crate::{
9    error::DiffsolRtError, linear_solver_type::LinearSolverType,
10    ode_solver_tag::OdeSolverMethodTag, ode_solver_type::OdeSolverType, scalar_type::Scalar,
11};
12
13pub(crate) trait VjpCheckpoint: Any + Send {
14    fn as_any(&self) -> &dyn Any;
15    fn method(&self) -> OdeSolverType;
16    fn linear_solver(&self) -> LinearSolverType;
17    /// The number of times of the forward solve.
18    fn nt(&self) -> usize;
19}
20
21impl dyn VjpCheckpoint + '_ {
22    pub(crate) fn data<M, Eqn, Tag>(&self) -> Result<&VjpCheckpointData<M, Eqn, Tag>, DiffsolError>
23    where
24        M: diffsol::Matrix<T: Scalar> + DefaultSolver + 'static,
25        Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C> + 'static,
26        Tag: OdeSolverMethodTag<M, Eqn> + 'static,
27    {
28        self.as_any()
29            .downcast_ref::<VjpCheckpointData<M, Eqn, Tag>>()
30            .ok_or_else(|| {
31                DiffsolError::Other(
32                    "VJP checkpoint is incompatible with this ODE solver".to_string(),
33                )
34            })
35    }
36}
37
38/// The data of a forward pass of [crate::OdeWrapper::solve_vjp_fwd]: the tape of the solve,
39/// the parameters of all the lanes, and the solver that made the tape, which must also do the
40/// backward pass.
41pub(crate) struct VjpCheckpointData<M, Eqn, Tag>
42where
43    M: diffsol::Matrix<T: Scalar> + DefaultSolver,
44    Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C>,
45    Tag: OdeSolverMethodTag<M, Eqn>,
46{
47    pub(crate) tape: DenseVjpTape<M::V, Tag::State>,
48    pub(crate) params: M::V,
49    method: OdeSolverType,
50    linear_solver: LinearSolverType,
51    _eqn: std::marker::PhantomData<fn() -> Eqn>,
52}
53
54impl<M, Eqn, Tag> VjpCheckpointData<M, Eqn, Tag>
55where
56    M: diffsol::Matrix<T: Scalar> + DefaultSolver,
57    Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C>,
58    Tag: OdeSolverMethodTag<M, Eqn>,
59{
60    pub(crate) fn new(
61        tape: DenseVjpTape<M::V, Tag::State>,
62        params: M::V,
63        method: OdeSolverType,
64        linear_solver: LinearSolverType,
65    ) -> Self {
66        Self {
67            tape,
68            params,
69            method,
70            linear_solver,
71            _eqn: std::marker::PhantomData,
72        }
73    }
74}
75
76impl<M, Eqn, Tag> VjpCheckpoint for VjpCheckpointData<M, Eqn, Tag>
77where
78    M: diffsol::Matrix<T: Scalar> + DefaultSolver + 'static,
79    M::V: Send,
80    Tag::State: Send,
81    Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C> + 'static,
82    Tag: OdeSolverMethodTag<M, Eqn> + 'static,
83{
84    fn as_any(&self) -> &dyn Any {
85        self
86    }
87
88    fn method(&self) -> OdeSolverType {
89        self.method
90    }
91
92    fn linear_solver(&self) -> LinearSolverType {
93        self.linear_solver
94    }
95
96    fn nt(&self) -> usize {
97        self.tape.t_eval().len()
98    }
99}
100
101/// The checkpoint of a forward pass of [crate::OdeWrapper::solve_vjp_fwd], for one or more
102/// backward passes with [crate::OdeWrapper::solve_vjp_bkwd]. It can be shared across
103/// threads.
104#[derive(Clone)]
105pub struct VjpCheckpointWrapper(Arc<Mutex<Box<dyn VjpCheckpoint>>>);
106
107impl VjpCheckpointWrapper {
108    pub(crate) fn new(checkpoint: Box<dyn VjpCheckpoint>) -> Self {
109        Self(Arc::new(Mutex::new(checkpoint)))
110    }
111
112    pub(crate) fn guard(&self) -> Result<MutexGuard<'_, Box<dyn VjpCheckpoint>>, DiffsolRtError> {
113        self.0
114            .lock()
115            .map_err(|_| DiffsolError::Other("VJP checkpoint mutex poisoned".to_string()).into())
116    }
117}
diffsol_c/vjp_checkpoint.rs

diffsol_c/
vjp_checkpoint.rs
