Skip to main content

diffsol_c/
solution.rs

1use std::sync::Arc;
2
3use diffsol::{DefaultDenseMatrix, DenseMatrix, Solution as DiffsolSolution, Vector};
4
5use crate::tensor::{shared_tensor, Tensor, TensorLayout};
6
7/// A solution of a solve. It does not change after the solve, so the tensors that the methods
8/// give share its data, with no copy.
9///
10/// Each method takes an `Arc` of the solution: the tensor holds it, so the data stays valid
11/// while the tensor exists, also after the [crate::SolutionWrapper] is dropped.
12pub(crate) trait Solution: Send + Sync {
13    /// The solution values, as a `[nout, nt, nbatch]` array.
14    fn get_ys(self: Arc<Self>) -> Tensor;
15    /// The solution times, as a `[nt]` array. All the lanes use the same times.
16    fn get_ts(self: Arc<Self>) -> Tensor;
17    /// The sensitivities, as `nparams` arrays, each `[nout, nt, nbatch]`.
18    fn get_sens(self: Arc<Self>) -> Vec<Tensor>;
19}
20
21impl<V> Solution for DiffsolSolution<V>
22where
23    V: Vector + DefaultDenseMatrix + Send + Sync + 'static,
24    <V as DefaultDenseMatrix>::M: DenseMatrix + Send + Sync + TensorLayout,
25    Vec<V::T>: TensorLayout,
26{
27    fn get_sens(self: Arc<Self>) -> Vec<Tensor> {
28        (0..self.y_sens.len())
29            .map(|i| {
30                let nt = self.ts.len();
31                // a matrix can have more columns than times: keep a column for each time
32                shared_tensor(self.clone(), |s| &s.y_sens[i]).truncate_axis(1, nt)
33            })
34            .collect()
35    }
36
37    fn get_ts(self: Arc<Self>) -> Tensor {
38        shared_tensor(self, |s| &s.ts)
39    }
40
41    fn get_ys(self: Arc<Self>) -> Tensor {
42        let nt = self.ts.len();
43        // a solve to a final time grows `ys` with spare columns: keep a column for each time
44        shared_tensor(self, |s| &s.ys).truncate_axis(1, nt)
45    }
46}