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

diffsol_c/
adjoint_checkpoint.rs
