diffsol/op/
mod.rs

1use crate::{
2    ConstantOp, ConstantOpSens, ConstantOpSensAdjoint, Context, LinearOp, LinearOpTranspose,
3    Matrix, NonLinearOp, NonLinearOpAdjoint, NonLinearOpSens, NonLinearOpSensAdjoint, Scalar,
4    Vector,
5};
6
7use nonlinear_op::NonLinearOpJacobian;
8use serde::Serialize;
9
10pub mod bdf;
11pub mod closure;
12pub mod closure_no_jac;
13pub mod closure_with_adjoint;
14pub mod closure_with_sens;
15pub mod constant_closure;
16pub mod constant_closure_with_adjoint;
17pub mod constant_closure_with_sens;
18pub mod constant_op;
19pub mod init;
20pub mod linear_closure;
21pub mod linear_closure_with_adjoint;
22pub mod linear_op;
23pub mod linearise;
24pub mod matrix;
25pub mod nonlinear_op;
26pub mod sdirk;
27pub mod stoch;
28pub mod unit;
29
30/// A generic operator trait.
31///
32/// Op is a trait for operators that, given a paramter vector `p`, operates on an input vector `x` to produce an output vector `y`.
33/// It defines the number of states (i.e. length of `x`), the number of outputs (i.e. length of `y`), and number of parameters (i.e. length of `p`) of the operator.
34/// It also defines the type of the scalar, vector, and matrices used in the operator.
35pub trait Op {
36    type T: Scalar;
37    type V: Vector<T = Self::T, C = Self::C>;
38    type M: Matrix<T = Self::T, V = Self::V, C = Self::C>;
39    type C: Context;
40
41    /// return the context of the operator
42    fn context(&self) -> &Self::C;
43
44    /// Return the number of input states of the operator.
45    fn nstates(&self) -> usize;
46
47    /// Return the number of outputs of the operator.
48    fn nout(&self) -> usize;
49
50    /// Return the number of parameters of the operator.
51    fn nparams(&self) -> usize;
52
53    /// Return statistics about the operator (e.g. how many times it was called, how many times the jacobian was computed, etc.)
54    fn statistics(&self) -> OpStatistics {
55        OpStatistics::default()
56    }
57}
58
59/// A wrapper for an operator that parameterises it with a parameter vector.
60pub struct ParameterisedOp<'a, C: Op> {
61    pub op: &'a C,
62    pub p: &'a C::V,
63}
64
65impl<'a, C: Op> ParameterisedOp<'a, C> {
66    pub fn new(op: &'a C, p: &'a C::V) -> Self {
67        Self { op, p }
68    }
69}
70
71pub trait BuilderOp: Op {
72    fn set_nstates(&mut self, nstates: usize);
73    fn set_nparams(&mut self, nparams: usize);
74    fn set_nout(&mut self, nout: usize);
75    fn calculate_sparsity(&mut self, y0: &Self::V, t0: Self::T, p: &Self::V);
76}
77
78impl<C: Op> Op for ParameterisedOp<'_, C> {
79    type V = C::V;
80    type T = C::T;
81    type M = C::M;
82    type C = C::C;
83    fn nstates(&self) -> usize {
84        self.op.nstates()
85    }
86    fn nout(&self) -> usize {
87        self.op.nout()
88    }
89    fn nparams(&self) -> usize {
90        self.op.nparams()
91    }
92    fn statistics(&self) -> OpStatistics {
93        self.op.statistics()
94    }
95    fn context(&self) -> &Self::C {
96        self.op.context()
97    }
98}
99
100#[derive(Default, Clone, Serialize)]
101pub struct OpStatistics {
102    pub number_of_calls: usize,
103    pub number_of_jac_muls: usize,
104    pub number_of_matrix_evals: usize,
105    pub number_of_jac_adj_muls: usize,
106}
107
108impl OpStatistics {
109    pub fn new() -> Self {
110        Self {
111            number_of_jac_muls: 0,
112            number_of_calls: 0,
113            number_of_matrix_evals: 0,
114            number_of_jac_adj_muls: 0,
115        }
116    }
117
118    pub fn increment_call(&mut self) {
119        self.number_of_calls += 1;
120    }
121
122    pub fn increment_jac_mul(&mut self) {
123        self.number_of_jac_muls += 1;
124    }
125
126    pub fn increment_jac_adj_mul(&mut self) {
127        self.number_of_jac_adj_muls += 1;
128    }
129
130    pub fn increment_matrix(&mut self) {
131        self.number_of_matrix_evals += 1;
132    }
133}
134
135impl<C: Op> Op for &C {
136    type T = C::T;
137    type V = C::V;
138    type M = C::M;
139    type C = C::C;
140    fn nstates(&self) -> usize {
141        C::nstates(*self)
142    }
143    fn nout(&self) -> usize {
144        C::nout(*self)
145    }
146    fn nparams(&self) -> usize {
147        C::nparams(*self)
148    }
149    fn statistics(&self) -> OpStatistics {
150        C::statistics(*self)
151    }
152    fn context(&self) -> &Self::C {
153        C::context(*self)
154    }
155}
156
157impl<C: Op> Op for &mut C {
158    type T = C::T;
159    type V = C::V;
160    type M = C::M;
161    type C = C::C;
162    fn nstates(&self) -> usize {
163        C::nstates(*self)
164    }
165    fn nout(&self) -> usize {
166        C::nout(*self)
167    }
168    fn nparams(&self) -> usize {
169        C::nparams(*self)
170    }
171    fn statistics(&self) -> OpStatistics {
172        C::statistics(*self)
173    }
174    fn context(&self) -> &Self::C {
175        C::context(*self)
176    }
177}
178
179impl<C: NonLinearOp> NonLinearOp for &C {
180    fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) {
181        C::call_inplace(*self, x, t, y)
182    }
183}
184
185impl<C: NonLinearOpJacobian> NonLinearOpJacobian for &C {
186    fn jac_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) {
187        C::jac_mul_inplace(*self, x, t, v, y)
188    }
189    fn jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
190        C::jacobian_inplace(*self, x, t, y)
191    }
192    fn jacobian_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
193        C::jacobian_sparsity(*self)
194    }
195}
196
197impl<C: NonLinearOpAdjoint> NonLinearOpAdjoint for &C {
198    fn adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
199        C::adjoint_inplace(*self, x, t, y)
200    }
201    fn adjoint_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
202        C::adjoint_sparsity(*self)
203    }
204    fn jac_transpose_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) {
205        C::jac_transpose_mul_inplace(*self, x, t, v, y)
206    }
207}
208
209impl<C: NonLinearOpSens> NonLinearOpSens for &C {
210    fn sens_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) {
211        C::sens_mul_inplace(*self, x, t, v, y)
212    }
213    fn sens_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
214        C::sens_inplace(*self, x, t, y)
215    }
216
217    fn sens_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
218        C::sens_sparsity(*self)
219    }
220}
221
222impl<C: NonLinearOpSensAdjoint> NonLinearOpSensAdjoint for &C {
223    fn sens_transpose_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) {
224        C::sens_transpose_mul_inplace(*self, x, t, v, y)
225    }
226    fn sens_adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
227        C::sens_adjoint_inplace(*self, x, t, y)
228    }
229    fn sens_adjoint_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
230        C::sens_adjoint_sparsity(*self)
231    }
232}
233
234impl<C: LinearOp> LinearOp for &C {
235    fn gemv_inplace(&self, x: &Self::V, t: Self::T, beta: Self::T, y: &mut Self::V) {
236        C::gemv_inplace(*self, x, t, beta, y)
237    }
238    fn sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
239        C::sparsity(*self)
240    }
241    fn matrix_inplace(&self, t: Self::T, y: &mut Self::M) {
242        C::matrix_inplace(*self, t, y)
243    }
244}
245
246impl<C: LinearOpTranspose> LinearOpTranspose for &C {
247    fn gemv_transpose_inplace(&self, x: &Self::V, t: Self::T, beta: Self::T, y: &mut Self::V) {
248        C::gemv_transpose_inplace(*self, x, t, beta, y)
249    }
250    fn transpose_inplace(&self, t: Self::T, y: &mut Self::M) {
251        C::transpose_inplace(*self, t, y)
252    }
253    fn transpose_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
254        C::transpose_sparsity(*self)
255    }
256}
257
258impl<C: ConstantOp> ConstantOp for &C {
259    fn call_inplace(&self, t: Self::T, y: &mut Self::V) {
260        C::call_inplace(*self, t, y)
261    }
262}
263
264impl<C: ConstantOpSens> ConstantOpSens for &C {
265    fn sens_mul_inplace(&self, t: Self::T, v: &Self::V, y: &mut Self::V) {
266        C::sens_mul_inplace(*self, t, v, y)
267    }
268    fn sens_inplace(&self, t: Self::T, y: &mut Self::M) {
269        C::sens_inplace(*self, t, y)
270    }
271    fn sens_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
272        C::sens_sparsity(*self)
273    }
274}
275
276impl<C: ConstantOpSensAdjoint> ConstantOpSensAdjoint for &C {
277    fn sens_transpose_mul_inplace(&self, t: Self::T, v: &Self::V, y: &mut Self::V) {
278        C::sens_transpose_mul_inplace(*self, t, v, y)
279    }
280    fn sens_adjoint_inplace(&self, t: Self::T, y: &mut Self::M) {
281        C::sens_adjoint_inplace(*self, t, y)
282    }
283    fn sens_adjoint_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
284        C::sens_adjoint_sparsity(*self)
285    }
286}
diffsol/op/mod.rs

diffsol/op/
mod.rs
