diffsol_c/
solution_wrapper.rs

1use std::sync::Arc;
2
3use serde::{ser::SerializeStruct, Serialize, Serializer};
4
5use crate::{
6    error::DiffsolRtError,
7    scalar_type::{Scalar, ScalarType},
8    solution::Solution,
9    tensor::Tensor,
10};
11
12/// A handle to the solution of a solve. Clones share the same solution.
13///
14/// The solution does not change after the solve, so the tensors that the methods give share its
15/// data, with no copy. A tensor stays valid after the handle is dropped.
16#[derive(Clone)]
17pub struct SolutionWrapper(Arc<dyn Solution>);
18
19impl SolutionWrapper {
20    pub(crate) fn new(solution: Arc<dyn Solution>) -> Self {
21        Self(solution)
22    }
23
24    /// The solution values, as a `[nout, nt, nbatch]` array.
25    pub fn get_ys(&self) -> Result<Tensor, DiffsolRtError> {
26        Ok(self.0.clone().get_ys())
27    }
28
29    /// The solution times, as a `[nt]` array. All the lanes use the same times.
30    pub fn get_ts(&self) -> Result<Tensor, DiffsolRtError> {
31        Ok(self.0.clone().get_ts())
32    }
33
34    /// The sensitivities, as `nparams` arrays, each `[nout, nt, nbatch]`.
35    pub fn get_sens(&self) -> Result<Vec<Tensor>, DiffsolRtError> {
36        Ok(self.0.clone().get_sens())
37    }
38}
39
40/// Read a 3D array, for example `[nout, nt, nbatch]`, as nested values of the scalar type `T` in
41/// the order of its axes: `values[i][j][k]` is the value at the index `(i, j, k)`.
42fn nested<T: Scalar>(tensor: &Tensor) -> Result<Vec<Vec<Vec<T>>>, DiffsolRtError> {
43    let (n0, n1, n2) = (tensor.shape()[0], tensor.shape()[1], tensor.shape()[2]);
44    // `columns[j][k]` is the column `[:, j, k]`
45    let columns = (0..n1)
46        .map(|j| {
47            (0..n2)
48                .map(|k| tensor.column::<T>(&[j, k]))
49                .collect::<Result<Vec<_>, _>>()
50        })
51        .collect::<Result<Vec<_>, _>>()?;
52    Ok((0..n0)
53        .map(|i| {
54            (0..n1)
55                .map(|j| (0..n2).map(|k| columns[j][k][i]).collect())
56                .collect()
57        })
58        .collect())
59}
60
61impl SolutionWrapper {
62    /// Serialize the solution with the values as `T`, the scalar type of the solver.
63    fn serialize_as<T: Scalar + Serialize, S: Serializer>(
64        &self,
65        serializer: S,
66    ) -> Result<S::Ok, S::Error> {
67        let error = serde::ser::Error::custom;
68        let ts_tensor = self.get_ts().map_err(error)?;
69        let ys_tensor = self.get_ys().map_err(error)?;
70        let ts = ts_tensor.column::<T>(&[]).map_err(error)?;
71        let ys = nested::<T>(&ys_tensor).map_err(error)?;
72        let sensitivities = self
73            .get_sens()
74            .map_err(error)?
75            .iter()
76            .map(nested::<T>)
77            .collect::<Result<Vec<_>, _>>()
78            .map_err(error)?;
79
80        let mut state = serializer.serialize_struct("SolutionWrapper", 4)?;
81        state.serialize_field("dtype", &T::scalar_type())?;
82        state.serialize_field("ts", ts)?;
83        state.serialize_field("ys", &ys)?;
84        state.serialize_field("sensitivities", &sensitivities)?;
85        state.end()
86    }
87}
88
89/// The solution serializes as `{dtype, ts, ys, sensitivities}`. The values have the scalar type
90/// of the solver, `dtype`, so an `f32` solution keeps its values exactly as `f32` (in JSON, `0.1`
91/// and not `0.10000000149011612`).
92impl Serialize for SolutionWrapper {
93    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
94    where
95        S: Serializer,
96    {
97        let dtype = self.get_ts().map_err(serde::ser::Error::custom)?.dtype();
98        match dtype {
99            ScalarType::F32 => self.serialize_as::<f32, S>(serializer),
100            ScalarType::F64 => self.serialize_as::<f64, S>(serializer),
101        }
102    }
103}
104
105#[cfg(test)]
106mod tests {
107    use std::sync::Arc;
108
109    use super::{nested, SolutionWrapper};
110    use crate::solution::Solution;
111    use crate::tensor::{Tensor, ToTensor};
112
113    #[test]
114    fn nested_keep_zero_length_axes() {
115        let no_rows = Tensor::from_vec(Vec::<f64>::new(), &[2, 0, 4]);
116        assert_eq!(
117            nested::<f64>(&no_rows).unwrap(),
118            vec![Vec::<Vec<f64>>::new(); 2]
119        );
120        let no_cols = Tensor::from_vec(Vec::<f64>::new(), &[3, 1, 0]);
121        assert_eq!(
122            nested::<f64>(&no_cols).unwrap(),
123            vec![vec![Vec::<f64>::new()]; 3]
124        );
125    }
126
127    #[test]
128    fn nested_read_f32_values() {
129        // a `[2, 2, 1]` array of `f32` values: a 2x2 column-major matrix with one lane
130        let m = <diffsol::NalgebraMat<f32> as diffsol::DenseMatrix>::from_vec(
131            2,
132            2,
133            vec![1.0, 2.0, 3.0, 4.0],
134            Default::default(),
135        );
136        let tensor = m.to_tensor();
137        assert_eq!(
138            nested::<f32>(&tensor).unwrap(),
139            vec![vec![vec![1.0], vec![3.0]], vec![vec![2.0], vec![4.0]]]
140        );
141        assert!(nested::<f64>(&tensor).is_err());
142    }
143
144    /// A solution of `f32` values with 2 times, 1 output and 1 lane, and no sensitivities.
145    struct F32Solution;
146
147    impl Solution for F32Solution {
148        fn get_ys(self: Arc<Self>) -> Tensor {
149            Tensor::from_vec(vec![0.1f32, 0.2], &[1, 2, 1])
150        }
151        fn get_ts(self: Arc<Self>) -> Tensor {
152            Tensor::from_vec(vec![0.5f32, 1.0], &[2])
153        }
154        fn get_sens(self: Arc<Self>) -> Vec<Tensor> {
155            Vec::new()
156        }
157    }
158
159    #[test]
160    fn solutions_serialize_with_the_scalar_type_of_the_solver() {
161        let solution = SolutionWrapper::new(Arc::new(F32Solution));
162        let json = serde_json::to_string(&solution).unwrap();
163        assert_eq!(
164            json,
165            r#"{"dtype":"f32","ts":[0.5,1.0],"ys":[[[0.1],[0.2]]],"sensitivities":[]}"#
166        );
167    }
168}
diffsol_c/solution_wrapper.rs

diffsol_c/
solution_wrapper.rs
