1use std::cell::RefCell;
2
3use crate::{
4 jacobian::{
5 find_adjoint_non_zeros, find_jacobian_non_zeros, find_sens_adjoint_non_zeros,
6 JacobianColoring,
7 },
8 Matrix, MatrixSparsity, NonLinearOp, NonLinearOpAdjoint, NonLinearOpJacobian,
9 NonLinearOpSensAdjoint, Op, Vector,
10};
11
12use super::{BuilderOp, OpStatistics, ParameterisedOp};
13
14#[derive(Clone)]
15pub struct ClosureWithAdjoint<M, F, G, H, I>
16where
17 M: Matrix,
18 F: Fn(&M::V, &M::V, M::T, &mut M::V),
19 G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
20 H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
21 I: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
22{
23 func: F,
24 jacobian_action: G,
25 jacobian_adjoint_action: H,
26 sens_adjoint_action: I,
27 nstates: usize,
28 nout: usize,
29 nparams: usize,
30 coloring: Option<JacobianColoring<M>>,
31 sparsity: Option<M::Sparsity>,
32 sparsity_adjoint: Option<M::Sparsity>,
33 coloring_adjoint: Option<JacobianColoring<M>>,
34 sens_sparsity: Option<M::Sparsity>,
35 coloring_sens_adjoint: Option<JacobianColoring<M>>,
36 statistics: RefCell<OpStatistics>,
37 ctx: M::C,
38}
39
40impl<M, F, G, H, I> ClosureWithAdjoint<M, F, G, H, I>
41where
42 M: Matrix,
43 F: Fn(&M::V, &M::V, M::T, &mut M::V),
44 G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
45 H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
46 I: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
47{
48 #[allow(clippy::too_many_arguments)]
49 pub fn new(
50 func: F,
51 jacobian_action: G,
52 jacobian_adjoint_action: H,
53 sens_adjoint_action: I,
54 nstates: usize,
55 nout: usize,
56 nparams: usize,
57 ctx: M::C,
58 ) -> Self {
59 Self {
60 func,
61 jacobian_action,
62 jacobian_adjoint_action,
63 sens_adjoint_action,
64 nstates,
65 nout,
66 nparams,
67 statistics: RefCell::new(OpStatistics::default()),
68 coloring: None,
69 sparsity: None,
70 sparsity_adjoint: None,
71 coloring_adjoint: None,
72 sens_sparsity: None,
73 coloring_sens_adjoint: None,
74 ctx,
75 }
76 }
77
78 pub fn calculate_jacobian_sparsity(&mut self, y0: &M::V, t0: M::T, p: &M::V) {
79 let op = ParameterisedOp { op: self, p };
80 let non_zeros = find_jacobian_non_zeros(&op, y0, t0);
81 self.sparsity = Some(
82 MatrixSparsity::try_from_indices(self.nout(), self.nstates(), non_zeros.clone())
83 .expect("invalid sparsity pattern"),
84 );
85 self.coloring = Some(JacobianColoring::new(
86 self.sparsity.as_ref().unwrap(),
87 &non_zeros,
88 self.ctx.clone(),
89 ));
90 }
91
92 pub fn calculate_adjoint_sparsity(&mut self, y0: &M::V, t0: M::T, p: &M::V) {
93 let op = ParameterisedOp { op: self, p };
94 let non_zeros = find_adjoint_non_zeros(&op, y0, t0);
95 self.sparsity_adjoint = Some(
96 MatrixSparsity::try_from_indices(self.nstates, self.nout, non_zeros.clone())
97 .expect("invalid sparsity pattern"),
98 );
99 self.coloring_adjoint = Some(JacobianColoring::new(
100 self.sparsity_adjoint.as_ref().unwrap(),
101 &non_zeros,
102 self.ctx.clone(),
103 ));
104 }
105
106 pub fn calculate_sens_adjoint_sparsity(&mut self, y0: &M::V, t0: M::T, p: &M::V) {
107 let op = ParameterisedOp { op: self, p };
108 let non_zeros = find_sens_adjoint_non_zeros(&op, y0, t0);
109 let nparams = p.len();
110 self.sens_sparsity = Some(
111 MatrixSparsity::try_from_indices(nparams, self.nstates, non_zeros.clone())
112 .expect("invalid sparsity pattern"),
113 );
114 self.coloring_sens_adjoint = Some(JacobianColoring::new(
115 self.sens_sparsity.as_ref().unwrap(),
116 &non_zeros,
117 self.ctx.clone(),
118 ));
119 }
120}
121
122impl<M, F, G, H, I> Op for ClosureWithAdjoint<M, F, G, H, I>
123where
124 M: Matrix,
125 F: Fn(&M::V, &M::V, M::T, &mut M::V),
126 G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
127 H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
128 I: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
129{
130 type V = M::V;
131 type T = M::T;
132 type M = M;
133 type C = M::C;
134 fn nstates(&self) -> usize {
135 self.nstates
136 }
137 fn nout(&self) -> usize {
138 self.nout
139 }
140 fn nparams(&self) -> usize {
141 self.nparams
142 }
143 fn statistics(&self) -> OpStatistics {
144 self.statistics.borrow().clone()
145 }
146 fn context(&self) -> &Self::C {
147 &self.ctx
148 }
149}
150
151impl<M, F, G, H, I> BuilderOp for ClosureWithAdjoint<M, F, G, H, I>
152where
153 M: Matrix,
154 F: Fn(&M::V, &M::V, M::T, &mut M::V),
155 G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
156 H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
157 I: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
158{
159 fn calculate_sparsity(&mut self, y0: &Self::V, t0: Self::T, p: &Self::V) {
160 self.calculate_jacobian_sparsity(y0, t0, p);
161 self.calculate_adjoint_sparsity(y0, t0, p);
162 self.calculate_sens_adjoint_sparsity(y0, t0, p);
163 }
164 fn set_nstates(&mut self, nstates: usize) {
165 self.nstates = nstates;
166 }
167 fn set_nout(&mut self, nout: usize) {
168 self.nout = nout;
169 }
170 fn set_nparams(&mut self, nparams: usize) {
171 self.nparams = nparams;
172 }
173}
174
175impl<M, F, G, H, I> NonLinearOp for ParameterisedOp<'_, ClosureWithAdjoint<M, F, G, H, I>>
176where
177 M: Matrix,
178 F: Fn(&M::V, &M::V, M::T, &mut M::V),
179 G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
180 H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
181 I: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
182{
183 fn call_inplace(&self, x: &M::V, t: M::T, y: &mut M::V) {
184 self.op.statistics.borrow_mut().increment_call();
185 (self.op.func)(x, self.p, t, y)
186 }
187}
188
189impl<M, F, G, H, I> NonLinearOpJacobian for ParameterisedOp<'_, ClosureWithAdjoint<M, F, G, H, I>>
190where
191 M: Matrix,
192 F: Fn(&M::V, &M::V, M::T, &mut M::V),
193 G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
194 H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
195 I: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
196{
197 fn jac_mul_inplace(&self, x: &M::V, t: M::T, v: &M::V, y: &mut M::V) {
198 self.op.statistics.borrow_mut().increment_jac_mul();
199 (self.op.jacobian_action)(x, self.p, t, v, y)
200 }
201 fn jacobian_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
202 self.op.statistics.borrow_mut().increment_matrix();
203 if let Some(coloring) = self.op.coloring.as_ref() {
204 coloring.jacobian_inplace(self, x, t, y);
205 } else {
206 self._default_jacobian_inplace(x, t, y);
207 }
208 }
209 fn jacobian_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
210 self.op.sparsity.clone()
211 }
212}
213
214impl<M, F, G, H, I> NonLinearOpAdjoint for ParameterisedOp<'_, ClosureWithAdjoint<M, F, G, H, I>>
215where
216 M: Matrix,
217 F: Fn(&M::V, &M::V, M::T, &mut M::V),
218 G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
219 H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
220 I: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
221{
222 fn jac_transpose_mul_inplace(&self, x: &Self::V, t: Self::T, v: &Self::V, y: &mut Self::V) {
223 self.op.statistics.borrow_mut().increment_jac_adj_mul();
224 (self.op.jacobian_adjoint_action)(x, self.p, t, v, y);
225 }
226
227 fn adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
228 if let Some(coloring) = self.op.coloring_adjoint.as_ref() {
229 coloring.adjoint_inplace(self, x, t, y);
230 } else {
231 self._default_adjoint_inplace(x, t, y);
232 }
233 }
234 fn adjoint_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
235 self.op.sparsity_adjoint.clone()
236 }
237}
238
239impl<M, F, G, H, I> NonLinearOpSensAdjoint
240 for ParameterisedOp<'_, ClosureWithAdjoint<M, F, G, H, I>>
241where
242 M: Matrix,
243 F: Fn(&M::V, &M::V, M::T, &mut M::V),
244 G: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
245 H: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
246 I: Fn(&M::V, &M::V, M::T, &M::V, &mut M::V),
247{
248 fn sens_transpose_mul_inplace(&self, _x: &Self::V, _t: Self::T, _v: &Self::V, y: &mut Self::V) {
249 (self.op.sens_adjoint_action)(_x, self.p, _t, _v, y);
250 }
251 fn sens_adjoint_inplace(&self, x: &Self::V, t: Self::T, y: &mut Self::M) {
252 if let Some(coloring) = self.op.coloring_sens_adjoint.as_ref() {
253 coloring.sens_adjoint_inplace(self, x, t, y);
254 } else {
255 self._default_sens_adjoint_inplace(x, t, y);
256 }
257 }
258 fn sens_adjoint_sparsity(&self) -> Option<<Self::M as Matrix>::Sparsity> {
259 self.op.sens_sparsity.clone()
260 }
261}