diffsol_c/
adjoint_checkpoint.rs

1use std::{
2    any::Any,
3    sync::{Arc, Mutex, MutexGuard},
4};
5
6use diffsol::{CheckpointingPath, DefaultSolver, 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 AdjointCheckpoint: Any + Send {
14    fn as_any(&self) -> &dyn Any;
15    fn method(&self) -> OdeSolverType;
16    fn linear_solver(&self) -> LinearSolverType;
17    fn params(&self) -> &[f64];
18}
19
20impl dyn AdjointCheckpoint + '_ {
21    pub(crate) fn data<M, Eqn, Tag>(
22        &self,
23    ) -> Result<&AdjointCheckpointData<M, Eqn, Tag>, DiffsolError>
24    where
25        M: diffsol::Matrix<T: Scalar> + DefaultSolver + 'static,
26        Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C> + 'static,
27        Tag: OdeSolverMethodTag<M, Eqn> + 'static,
28    {
29        self.as_any()
30            .downcast_ref::<AdjointCheckpointData<M, Eqn, Tag>>()
31            .ok_or_else(|| {
32                DiffsolError::Other(
33                    "Adjoint checkpoint is incompatible with this ODE solver".to_string(),
34                )
35            })
36    }
37}
38
39pub(crate) struct AdjointCheckpointData<M, Eqn, Tag>
40where
41    M: diffsol::Matrix<T: Scalar> + DefaultSolver,
42    Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C>,
43    Tag: OdeSolverMethodTag<M, Eqn>,
44{
45    pub(crate) checkpointing: CheckpointingPath<Eqn, Tag::State>,
46    params: Vec<f64>,
47    method: OdeSolverType,
48    linear_solver: LinearSolverType,
49}
50
51impl<M, Eqn, Tag> AdjointCheckpointData<M, Eqn, Tag>
52where
53    M: diffsol::Matrix<T: Scalar> + DefaultSolver,
54    Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C>,
55    Tag: OdeSolverMethodTag<M, Eqn>,
56{
57    pub(crate) fn new(
58        checkpointing: CheckpointingPath<Eqn, Tag::State>,
59        params: Vec<f64>,
60        method: OdeSolverType,
61        linear_solver: LinearSolverType,
62    ) -> Self {
63        Self {
64            checkpointing,
65            params,
66            method,
67            linear_solver,
68        }
69    }
70}
71
72impl<M, Eqn, Tag> AdjointCheckpoint for AdjointCheckpointData<M, Eqn, Tag>
73where
74    M: diffsol::Matrix<T: Scalar> + DefaultSolver + 'static,
75    Eqn: OdeEquationsImplicit<M = M, T = M::T, V = M::V, C = M::C> + 'static,
76    Tag: OdeSolverMethodTag<M, Eqn> + 'static,
77{
78    fn as_any(&self) -> &dyn Any {
79        self
80    }
81
82    fn method(&self) -> OdeSolverType {
83        self.method
84    }
85
86    fn linear_solver(&self) -> LinearSolverType {
87        self.linear_solver
88    }
89
90    fn params(&self) -> &[f64] {
91        &self.params
92    }
93}
94
95/// Wrapper for an adjoint checkpoint that allows it to be shared across threads and mutated safely.    
96///
97/// Mutex is required because Checkpointing uses interior mutability when interpolating checkpoints
98#[derive(Clone)]
99pub struct AdjointCheckpointWrapper(Arc<Mutex<Box<dyn AdjointCheckpoint>>>);
100
101impl AdjointCheckpointWrapper {
102    pub(crate) fn new(checkpoint: Box<dyn AdjointCheckpoint>) -> Self {
103        Self(Arc::new(Mutex::new(checkpoint)))
104    }
105
106    pub(crate) fn guard(
107        &self,
108    ) -> Result<MutexGuard<'_, Box<dyn AdjointCheckpoint>>, DiffsolRtError> {
109        self.0.lock().map_err(|_| {
110            DiffsolError::Other("Adjoint checkpoint mutex poisoned".to_string()).into()
111        })
112    }
113}
diffsol_c/adjoint_checkpoint.rs

diffsol_c/
adjoint_checkpoint.rs
