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;
12#[cfg(feature = "autodiff")]
13pub mod closure_autodiff;
14pub mod closure_no_jac;
15pub mod closure_with_adjoint;
16pub mod closure_with_sens;
17pub mod constant_closure;
18#[cfg(feature = "autodiff")]
19pub mod constant_closure_autodiff;
20pub mod constant_closure_with_adjoint;
21pub mod constant_closure_with_sens;
22pub mod constant_op;
23pub mod init;
24pub mod linear_closure;
25#[cfg(feature = "autodiff")]
26pub mod linear_closure_autodiff;
27pub mod linear_closure_with_adjoint;
28pub mod linear_op;
29pub mod linearise;
30pub mod matrix;
31pub mod nonlinear_op;
32pub mod sdirk;
33pub mod stoch;
34pub mod unit;
35
36/// A generic operator trait.
37///
38/// Op is a trait for operators that, given a paramter vector `p`, operates on an input vector `x` to produce an output vector `y`.
39/// 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.
40/// It also defines the type of the scalar, vector, and matrices used in the operator.
41pub trait Op {
42    type T: Scalar;
43    type V: Vector<T = Self::T, C = Self::C>;
44    type M: Matrix<T = Self::T, V = Self::V, C = Self::C>;
45    type C: Context;
46
47    /// return the context of the operator
48    fn context(&self) -> &Self::C;
49
50    /// Return the number of input states of the operator.
51    fn nstates(&self) -> usize;
52
53    /// Return the number of outputs of the operator.
54    fn nout(&self) -> usize;
55
56    /// Return the number of parameters of the operator.
57    fn nparams(&self) -> usize;
58
59    /// Return statistics about the operator (e.g. how many times it was called, how many times the jacobian was computed, etc.)
60    fn statistics(&self) -> OpStatistics {
61        OpStatistics::default()
62    }
63}
64
65/// A wrapper for an operator that parameterises it with a parameter vector.
66pub struct ParameterisedOp<'a, C: Op> {
67    pub op: &'a C,
68    pub p: &'a C::V,
69}
70
71impl<'a, C: Op> ParameterisedOp<'a, C> {
72    pub fn new(op: &'a C, p: &'a C::V) -> Self {
73        Self { op, p }
74    }
75}
76
77/// trait interface for operators used in the [builder pattern](crate::OdeBuilder)
78pub trait BuilderOp: Op {
79    fn set_nstates(&mut self, nstates: usize);
80    fn set_nparams(&mut self, nparams: usize);
81    fn set_nout(&mut self, nout: usize);
82    fn calculate_sparsity(&mut self, y0: &Self::V, t0: Self::T, p: &Self::V);
83}
84
85impl<C: Op> Op for ParameterisedOp<'_, C> {
86    type V = C::V;
87    type T = C::T;
88    type M = C::M;
89    type C = C::C;
90    fn nstates(&self) -> usize {
91        self.op.nstates()
92    }
93    fn nout(&self) -> usize {
94        self.op.nout()
95    }
96    fn nparams(&self) -> usize {
97        self.op.nparams()
98    }
99    fn statistics(&self) -> OpStatistics {
100        self.op.statistics()
101    }
102    fn context(&self) -> &Self::C {
103        self.op.context()
104    }
105}
106
107/// Useful statistics about an operator.
108#[derive(Default, Clone, Serialize, Debug)]
109pub struct OpStatistics {
110    /// number of times the operator was called
111    pub number_of_calls: usize,
112    /// number of times the jacobian-vector product was computed
113    pub number_of_jac_muls: usize,
114    /// number of times the jacobian matrix was evaluated (SUNDIALS/CVODE `nje`)
115    pub number_of_matrix_evals: usize,
116    /// number of times the adjoint jacobian-vector product was computed
117    pub number_of_jac_adj_muls: usize,
118}
119
120impl OpStatistics {
121    pub fn new() -> Self {
122        Self {
123            number_of_jac_muls: 0,
124            number_of_calls: 0,
125            number_of_matrix_evals: 0,
126            number_of_jac_adj_muls: 0,
127        }
128    }
129
130    pub fn increment_call(&mut self) {
131        self.number_of_calls += 1;
132    }
133
134    pub fn increment_jac_mul(&mut self) {
135        self.number_of_jac_muls += 1;
136    }
137
138    pub fn increment_jac_adj_mul(&mut self) {
139        self.number_of_jac_adj_muls += 1;
140    }
141
142    pub fn increment_matrix(&mut self) {
143        self.number_of_matrix_evals += 1;
144    }
145}
146
147impl<C: Op> Op for &C {
148    type T = C::T;
149    type V = C::V;
150    type M = C::M;
151    type C = C::C;
152    fn nstates(&self) -> usize {
153        C::nstates(*self)
154    }
155    fn nout(&self) -> usize {
156        C::nout(*self)
157    }
158    fn nparams(&self) -> usize {
159        C::nparams(*self)
160    }
161    fn statistics(&self) -> OpStatistics {
162        C::statistics(*self)
163    }
164    fn context(&self) -> &Self::C {
165        C::context(*self)
166    }
167}
168
169impl<C: Op> Op for &mut C {
170    type T = C::T;
171    type V = C::V;
172    type M = C::M;
173    type C = C::C;
174    fn nstates(&self) -> usize {
175        C::nstates(*self)
176    }
177    fn nout(&self) -> usize {
178        C::nout(*self)
179    }
180    fn nparams(&self) -> usize {
181        C::nparams(*self)
182    }
183    fn statistics(&self) -> OpStatistics {
184        C::statistics(*self)
185    }
186    fn context(&self) -> &Self::C {
187        C::context(*self)
188    }
189}
190
191impl<C: NonLinearOp> NonLinearOp for &C {
192    fn call_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::V) {
193        C::call_inplace(*self, x, t, y)
194    }
195}
196
197impl<C: NonLinearOpJacobian> NonLinearOpJacobian for &C {
198    fn jac_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) {
199        C::jac_mul_inplace(*self, x, t, v, y)
200    }
201    fn jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
202        C::jacobian_inplace(*self, x, t, y)
203    }
204    fn jacobian_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
205        C::jacobian_sparsity(*self)
206    }
207}
208
209impl<C: NonLinearOpAdjoint> NonLinearOpAdjoint for &C {
210    fn adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
211        C::adjoint_inplace(*self, x, t, y)
212    }
213    fn adjoint_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
214        C::adjoint_sparsity(*self)
215    }
216    fn jac_transpose_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) {
217        C::jac_transpose_mul_inplace(*self, x, t, v, y)
218    }
219}
220
221impl<C: NonLinearOpSens> NonLinearOpSens for &C {
222    fn sens_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) {
223        C::sens_mul_inplace(*self, x, t, v, y)
224    }
225    fn sens_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
226        C::sens_inplace(*self, x, t, y)
227    }
228
229    fn sens_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
230        C::sens_sparsity(*self)
231    }
232}
233
234impl<C: NonLinearOpSensAdjoint> NonLinearOpSensAdjoint for &C {
235    fn sens_transpose_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) {
236        C::sens_transpose_mul_inplace(*self, x, t, v, y)
237    }
238    fn sens_adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
239        C::sens_adjoint_inplace(*self, x, t, y)
240    }
241    fn sens_adjoint_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
242        C::sens_adjoint_sparsity(*self)
243    }
244}
245
246impl<C: LinearOp> LinearOp for &C {
247    fn gemv_inplace(&self, x: &Self::V, t: Self::T, beta: Self::T, y: &mut Self::V) {
248        C::gemv_inplace(*self, x, t, beta, y)
249    }
250    fn sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
251        C::sparsity(*self)
252    }
253    fn matrix_inplace(&self, t: Self::T, y: &mut Self::M) {
254        C::matrix_inplace(*self, t, y)
255    }
256}
257
258impl<C: LinearOpTranspose> LinearOpTranspose for &C {
259    fn gemv_transpose_inplace(&self, x: &Self::V, t: Self::T, beta: Self::T, y: &mut Self::V) {
260        C::gemv_transpose_inplace(*self, x, t, beta, y)
261    }
262    fn transpose_inplace(&self, t: Self::T, y: &mut Self::M) {
263        C::transpose_inplace(*self, t, y)
264    }
265    fn transpose_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
266        C::transpose_sparsity(*self)
267    }
268}
269
270impl<C: ConstantOp> ConstantOp for &C {
271    fn call_inplace(&self, t: Self::T, y: &mut Self::V) {
272        C::call_inplace(*self, t, y)
273    }
274}
275
276impl<C: ConstantOpSens> ConstantOpSens for &C {
277    fn sens_mul_inplace(&self, t: Self::T, v: &Self::V, y: &mut Self::V) {
278        C::sens_mul_inplace(*self, t, v, y)
279    }
280    fn sens_inplace(&self, t: Self::T, y: &mut Self::M) {
281        C::sens_inplace(*self, t, y)
282    }
283    fn sens_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
284        C::sens_sparsity(*self)
285    }
286}
287
288impl<C: ConstantOpSensAdjoint> ConstantOpSensAdjoint for &C {
289    fn sens_transpose_mul_inplace(&self, t: Self::T, v: &Self::V, y: &mut Self::V) {
290        C::sens_transpose_mul_inplace(*self, t, v, y)
291    }
292    fn sens_adjoint_inplace(&self, t: Self::T, y: &mut Self::M) {
293        C::sens_adjoint_inplace(*self, t, y)
294    }
295    fn sens_adjoint_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
296        C::sens_adjoint_sparsity(*self)
297    }
298}
299
300#[cfg(test)]
301mod tests {
302    use std::cell::RefCell;
303
304    use crate::{
305        context::nalgebra::NalgebraContext, matrix::dense_nalgebra_serial::NalgebraMat, ConstantOp,
306        ConstantOpSens, ConstantOpSensAdjoint, LinearOp, LinearOpTranspose, NonLinearOp,
307        NonLinearOpAdjoint, NonLinearOpJacobian, NonLinearOpSens, NonLinearOpSensAdjoint, Vector,
308    };
309
310    use super::{Op, OpStatistics, ParameterisedOp};
311
312    type M = NalgebraMat<f64>;
313
314    struct ForwardingOp {
315        ctx: NalgebraContext,
316        stats: RefCell<OpStatistics>,
317    }
318
319    impl ForwardingOp {
320        fn new() -> Self {
321            Self {
322                ctx: NalgebraContext::default(),
323                stats: RefCell::new(OpStatistics::new()),
324            }
325        }
326    }
327
328    impl Op for ForwardingOp {
329        type T = f64;
330        type V = crate::NalgebraVec<f64>;
331        type M = M;
332        type C = NalgebraContext;
333
334        fn context(&self) -> &Self::C {
335            &self.ctx
336        }
337        fn nstates(&self) -> usize {
338            2
339        }
340        fn nout(&self) -> usize {
341            2
342        }
343        fn nparams(&self) -> usize {
344            2
345        }
346        fn statistics(&self) -> OpStatistics {
347            self.stats.borrow().clone()
348        }
349    }
350
351    impl NonLinearOp for ForwardingOp {
352        fn call_inplace(&self, x: &Self::V, _t: Self::T, y: &mut Self::V) {
353            self.stats.borrow_mut().increment_call();
354            y.copy_from(x);
355        }
356    }
357
358    impl NonLinearOpJacobian for ForwardingOp {
359        fn jac_mul_inplace(&self, _x: &Self::V, _t: Self::T, v: &Self::V, y: &mut Self::V) {
360            self.stats.borrow_mut().increment_jac_mul();
361            y.copy_from(v);
362        }
363    }
364
365    impl NonLinearOpAdjoint for ForwardingOp {
366        fn jac_transpose_mul_inplace(
367            &self,
368            _x: &Self::V,
369            _t: Self::T,
370            v: &Self::V,
371            y: &mut Self::V,
372        ) {
373            self.stats.borrow_mut().increment_jac_adj_mul();
374            y.copy_from(v);
375        }
376    }
377
378    impl NonLinearOpSens for ForwardingOp {
379        fn sens_mul_inplace(&self, _x: &Self::V, _t: Self::T, _v: &Self::V, y: &mut Self::V) {
380            y.fill(0.0);
381        }
382    }
383
384    impl NonLinearOpSensAdjoint for ForwardingOp {
385        fn sens_transpose_mul_inplace(
386            &self,
387            _x: &Self::V,
388            _t: Self::T,
389            _v: &Self::V,
390            y: &mut Self::V,
391        ) {
392            y.fill(0.0);
393        }
394    }
395
396    impl LinearOp for ForwardingOp {
397        fn gemv_inplace(&self, x: &Self::V, _t: Self::T, beta: Self::T, y: &mut Self::V) {
398            self.stats.borrow_mut().increment_call();
399            y.axpy(1.0, x, beta);
400        }
401    }
402
403    impl LinearOpTranspose for ForwardingOp {
404        fn gemv_transpose_inplace(&self, x: &Self::V, _t: Self::T, beta: Self::T, y: &mut Self::V) {
405            self.stats.borrow_mut().increment_jac_adj_mul();
406            y.axpy(1.0, x, beta);
407        }
408    }
409
410    impl ConstantOp for ForwardingOp {
411        fn call_inplace(&self, _t: Self::T, y: &mut Self::V) {
412            self.stats.borrow_mut().increment_call();
413            y.copy_from(&Self::V::from_vec(vec![1.0, 2.0], self.ctx));
414        }
415    }
416
417    impl ConstantOpSens for ForwardingOp {
418        fn sens_mul_inplace(&self, _t: Self::T, _v: &Self::V, y: &mut Self::V) {
419            y.fill(0.0);
420        }
421    }
422
423    impl ConstantOpSensAdjoint for ForwardingOp {
424        fn sens_transpose_mul_inplace(&self, _t: Self::T, _v: &Self::V, y: &mut Self::V) {
425            y.fill(0.0);
426        }
427    }
428
429    #[test]
430    fn op_statistics_increment_methods_update_counters() {
431        let mut stats = OpStatistics::new();
432        stats.increment_call();
433        stats.increment_jac_mul();
434        stats.increment_jac_adj_mul();
435        stats.increment_matrix();
436        assert_eq!(stats.number_of_calls, 1);
437        assert_eq!(stats.number_of_jac_muls, 1);
438        assert_eq!(stats.number_of_jac_adj_muls, 1);
439        assert_eq!(stats.number_of_matrix_evals, 1);
440    }
441
442    #[test]
443    fn parameterised_op_and_reference_forwarding_delegate_to_inner_operator() {
444        let op = ForwardingOp::new();
445        let p = crate::NalgebraVec::from_vec(vec![1.0, 2.0], NalgebraContext::default());
446        let pop = ParameterisedOp::new(&op, &p);
447        assert_eq!(pop.nstates(), 2);
448        assert_eq!(pop.nout(), 2);
449        assert_eq!(pop.nparams(), 2);
450
451        let x = crate::NalgebraVec::from_vec(vec![3.0, 4.0], NalgebraContext::default());
452        let mut y = crate::NalgebraVec::zeros(2, NalgebraContext::default());
453        NonLinearOp::call_inplace(&&op, &x, 0.0, &mut y);
454        y.assert_eq_st(&x, 1e-12);
455
456        op.jac_mul_inplace(&x, 0.0, &x, &mut y);
457        y.assert_eq_st(&x, 1e-12);
458
459        op.jac_transpose_mul_inplace(&x, 0.0, &x, &mut y);
460        y.assert_eq_st(&x, 1e-12);
461
462        NonLinearOpSens::sens_mul_inplace(&&op, &x, 0.0, &x, &mut y);
463        y.assert_eq_st(
464            &crate::NalgebraVec::zeros(2, NalgebraContext::default()),
465            1e-12,
466        );
467
468        NonLinearOpSensAdjoint::sens_transpose_mul_inplace(&&op, &x, 0.0, &x, &mut y);
469        y.assert_eq_st(
470            &crate::NalgebraVec::zeros(2, NalgebraContext::default()),
471            1e-12,
472        );
473
474        op.gemv_inplace(&x, 0.0, 0.0, &mut y);
475        y.assert_eq_st(&x, 1e-12);
476
477        op.gemv_transpose_inplace(&x, 0.0, 0.0, &mut y);
478        y.assert_eq_st(&x, 1e-12);
479
480        let mut y_const = crate::NalgebraVec::zeros(2, NalgebraContext::default());
481        <&ForwardingOp as ConstantOp>::call_inplace(&&op, 0.0, &mut y_const);
482        y_const.assert_eq_st(
483            &crate::NalgebraVec::from_vec(vec![1.0, 2.0], NalgebraContext::default()),
484            1e-12,
485        );
486
487        let op_ref_stats = pop.statistics();
488        assert!(op_ref_stats.number_of_calls >= 1);
489
490        let op_mut = ForwardingOp::new();
491        assert_eq!(op_mut.nstates(), 2);
492        assert_eq!(op_mut.nout(), 2);
493        assert_eq!(op_mut.nparams(), 2);
494    }
495}
diffsol/op/mod.rs

diffsol/op/
mod.rs
