diffsol/jacobian/
mod.rs

1use std::collections::HashSet;
2
3use crate::{
4    LinearOp, LinearOpTranspose, Matrix, MatrixSparsity, NonLinearOp, NonLinearOpAdjoint,
5    NonLinearOpJacobian, NonLinearOpSens, NonLinearOpSensAdjoint, Scalar, Vector, VectorIndex,
6};
7use num_traits::{One, Zero};
8
9use self::{coloring::nonzeros2graph, greedy_coloring::color_graph_greedy};
10
11pub mod coloring;
12pub mod graph;
13pub mod greedy_coloring;
14
15macro_rules! gen_find_non_zeros_nonlinear {
16    ($name:ident, $op_fn:ident, $op_trait:ident) => {
17        /// Find the non-zero entries of the $name matrix of a non-linear operator.
18        /// TODO: This function is not efficient for non-host vectors and could be part of the Vector trait
19        ///       to allow for more efficient implementations. It's ok for now since this is only used once
20        ///       during the setup phase.
21        pub fn $name<F: NonLinearOp + $op_trait + ?Sized>(
22            op: &F,
23            x: &F::V,
24            t: F::T,
25        ) -> Vec<(usize, usize)> {
26            let mut v = F::V::zeros(op.nstates(), op.context().clone());
27            let mut col = F::V::zeros(op.nout(), op.context().clone());
28            let mut triplets = Vec::with_capacity(op.nstates());
29            for j in 0..op.nstates() {
30                v.set_index(j, F::T::NAN);
31                op.$op_fn(x, t, &v, &mut col);
32                for i in 0..op.nout() {
33                    if col.get_index(i).is_nan() {
34                        triplets.push((i, j));
35                    }
36                    col.set_index(i, F::T::zero());
37                }
38                // OR:
39                //col.clone_as_vec().into_iter().for_each(|v| {
40                //    if v.is_nan() {
41                //        triplets.push((0, 0));
42                //    }
43                //});
44                col.fill(F::T::zero());
45                v.set_index(j, F::T::zero());
46            }
47            triplets
48        }
49    };
50}
51
52gen_find_non_zeros_nonlinear!(
53    find_jacobian_non_zeros,
54    jac_mul_inplace,
55    NonLinearOpJacobian
56);
57gen_find_non_zeros_nonlinear!(
58    find_adjoint_non_zeros,
59    jac_transpose_mul_inplace,
60    NonLinearOpAdjoint
61);
62gen_find_non_zeros_nonlinear!(find_sens_non_zeros, sens_mul_inplace, NonLinearOpSens);
63gen_find_non_zeros_nonlinear!(
64    find_sens_adjoint_non_zeros,
65    sens_transpose_mul_inplace,
66    NonLinearOpSensAdjoint
67);
68
69macro_rules! gen_find_non_zeros_linear {
70    ($name:ident, $op_fn:ident $(, $op_trait:tt )?) => {
71        /// Find the non-zero entries of the $name matrix of a non-linear operator.
72        /// TODO: This function is not efficient for non-host vectors and could be part of the Vector trait
73        ///       to allow for more efficient implementations. It's ok for now since this is only used once
74        ///       during the setup phase.
75        pub fn $name<F: LinearOp + ?Sized $(+ $op_trait)?>(op: &F, t: F::T) -> Vec<(usize, usize)> {
76            let mut v = F::V::zeros(op.nstates(), op.context().clone());
77            let mut col = F::V::zeros(op.nout(), op.context().clone());
78            let mut triplets = Vec::with_capacity(op.nstates());
79            for j in 0..op.nstates() {
80                v.set_index(j, F::T::NAN);
81                op.$op_fn(&v, t, &mut col);
82                for i in 0..op.nout() {
83                    if col.get_index(i).is_nan() {
84                        triplets.push((i, j));
85                    }
86                    col.set_index(i, F::T::zero());
87                }
88                // OR:
89                //col.clone_as_vec().into_iter().for_each(|v| {
90                //    if v.is_nan() {
91                //        triplets.push((0, 0));
92                //    }
93                //});
94                v.set_index(j, F::T::zero());
95            }
96            triplets
97        }
98    };
99}
100
101gen_find_non_zeros_linear!(find_matrix_non_zeros, call_inplace);
102gen_find_non_zeros_linear!(
103    find_transpose_non_zeros,
104    call_transpose_inplace,
105    LinearOpTranspose
106);
107
108#[derive(Clone)]
109pub struct JacobianColoring<M: Matrix> {
110    dst_indices_per_color: Vec<<M::V as Vector>::Index>,
111    src_indices_per_color: Vec<<M::V as Vector>::Index>,
112    input_indices_per_color: Vec<<M::V as Vector>::Index>,
113}
114
115impl<M: Matrix> JacobianColoring<M> {
116    pub fn new(sparsity: &impl MatrixSparsity<M>, non_zeros: &[(usize, usize)], ctx: M::C) -> Self {
117        let ncols = sparsity.ncols();
118        let graph = nonzeros2graph(non_zeros, ncols);
119        let coloring = color_graph_greedy(&graph);
120        let max_color = coloring.iter().max().copied().unwrap_or(0);
121        let mut dst_indices_per_color = Vec::new();
122        let mut src_indices_per_color = Vec::new();
123        let mut input_indices_per_color = Vec::new();
124        for c in 1..=max_color {
125            let mut rows = Vec::new();
126            let mut cols = Vec::new();
127            let mut indices = Vec::new();
128            for (i, j) in non_zeros {
129                if coloring[*j] == c {
130                    rows.push(*i);
131                    cols.push(*j);
132                    indices.push((*i, *j));
133                }
134            }
135            let dst_indices = sparsity.get_index(indices.as_slice(), ctx.clone());
136            let src_indices = <M::V as Vector>::Index::from_vec(rows, ctx.clone());
137            let unique_cols: HashSet<_> = HashSet::from_iter(cols.iter().cloned());
138            let unique_cols = unique_cols.into_iter().collect::<Vec<_>>();
139            let input_indices = <M::V as Vector>::Index::from_vec(unique_cols, ctx.clone());
140            dst_indices_per_color.push(dst_indices);
141            src_indices_per_color.push(src_indices);
142            input_indices_per_color.push(input_indices);
143        }
144        Self {
145            dst_indices_per_color,
146            src_indices_per_color,
147            input_indices_per_color,
148        }
149    }
150
151    //pub fn new_non_linear<F: NonLinearOp<M = M>>(op: &F) -> Self {
152    //    let non_zeros = find_non_zeros_nonlinear(op, &F::V::zeros(op.nstates()), F::T::zero());
153    //    Self::new_from_non_zeros(op, non_zeros)
154    //}
155
156    //pub fn new_linear<F: LinearOp<M = M>>(op: &F) -> Self {
157    //    let non_zeros = find_non_zeros_linear(op, F::T::zero());
158    //    Self::new_from_non_zeros(op, non_zeros)
159    //}
160
161    pub fn jacobian_inplace<F: NonLinearOpJacobian<M = M, V = M::V, T = M::T, C = M::C>>(
162        &self,
163        op: &F,
164        x: &F::V,
165        t: F::T,
166        y: &mut F::M,
167    ) {
168        let mut v = F::V::zeros(op.nstates(), op.context().clone());
169        let mut col = F::V::zeros(op.nout(), op.context().clone());
170        for c in 0..self.dst_indices_per_color.len() {
171            let input = &self.input_indices_per_color[c];
172            let dst_indices = &self.dst_indices_per_color[c];
173            let src_indices = &self.src_indices_per_color[c];
174            v.assign_at_indices(input, F::T::one());
175            op.jac_mul_inplace(x, t, &v, &mut col);
176            y.set_data_with_indices(dst_indices, src_indices, &col);
177            v.assign_at_indices(input, F::T::zero());
178        }
179    }
180
181    pub fn adjoint_inplace<F: NonLinearOpAdjoint<M = M, V = M::V, T = M::T, C = M::C>>(
182        &self,
183        op: &F,
184        x: &F::V,
185        t: F::T,
186        y: &mut F::M,
187    ) {
188        let mut v = F::V::zeros(op.nstates(), op.context().clone());
189        let mut col = F::V::zeros(op.nout(), op.context().clone());
190        for c in 0..self.dst_indices_per_color.len() {
191            let input = &self.input_indices_per_color[c];
192            let dst_indices = &self.dst_indices_per_color[c];
193            let src_indices = &self.src_indices_per_color[c];
194            v.assign_at_indices(input, F::T::one());
195            op.jac_transpose_mul_inplace(x, t, &v, &mut col);
196            y.set_data_with_indices(dst_indices, src_indices, &col);
197            v.assign_at_indices(input, F::T::zero());
198        }
199    }
200
201    pub fn sens_adjoint_inplace<F: NonLinearOpSensAdjoint<M = M, V = M::V, T = M::T, C = M::C>>(
202        &self,
203        op: &F,
204        x: &F::V,
205        t: F::T,
206        y: &mut F::M,
207    ) {
208        let mut v = F::V::zeros(op.nstates(), op.context().clone());
209        let mut col = F::V::zeros(op.nout(), op.context().clone());
210        for c in 0..self.dst_indices_per_color.len() {
211            let input = &self.input_indices_per_color[c];
212            let dst_indices = &self.dst_indices_per_color[c];
213            let src_indices = &self.src_indices_per_color[c];
214            v.assign_at_indices(input, F::T::one());
215            op.sens_transpose_mul_inplace(x, t, &v, &mut col);
216            y.set_data_with_indices(dst_indices, src_indices, &col);
217            v.assign_at_indices(input, F::T::zero());
218        }
219    }
220
221    pub fn matrix_inplace<F: LinearOp<M = M, V = M::V, T = M::T, C = M::C>>(
222        &self,
223        op: &F,
224        t: F::T,
225        y: &mut F::M,
226    ) {
227        let mut v = F::V::zeros(op.nstates(), op.context().clone());
228        let mut col = F::V::zeros(op.nout(), op.context().clone());
229        for c in 0..self.dst_indices_per_color.len() {
230            let input = &self.input_indices_per_color[c];
231            let dst_indices = &self.dst_indices_per_color[c];
232            let src_indices = &self.src_indices_per_color[c];
233            v.assign_at_indices(input, F::T::one());
234            op.call_inplace(&v, t, &mut col);
235            y.set_data_with_indices(dst_indices, src_indices, &col);
236            v.assign_at_indices(input, F::T::zero());
237        }
238    }
239}
240
241#[cfg(test)]
242mod tests {
243
244    use crate::jacobian::{find_jacobian_non_zeros, JacobianColoring};
245    use crate::matrix::dense_nalgebra_serial::NalgebraMat;
246    use crate::matrix::Matrix;
247    use crate::op::linear_closure::LinearClosure;
248    use crate::op::ParameterisedOp;
249    use crate::vector::Vector;
250    use crate::{
251        jacobian::{coloring::nonzeros2graph, greedy_coloring::color_graph_greedy},
252        op::closure::Closure,
253        LinearOp, Op,
254    };
255    use crate::{scale, FaerSparseMat, NonLinearOpJacobian};
256    use num_traits::{One, Zero};
257    use std::ops::MulAssign;
258
259    #[allow(clippy::type_complexity)]
260    fn helper_triplets2op_nonlinear<'a, M: Matrix + 'a>(
261        triplets: &'a [(usize, usize, M::T)],
262        p: &'a M::V,
263        nrows: usize,
264        ncols: usize,
265    ) -> Closure<
266        M,
267        impl Fn(&M::V, &M::V, M::T, &mut M::V) + use<'a, M>,
268        impl Fn(&M::V, &M::V, M::T, &M::V, &mut M::V) + use<'a, M>,
269    > {
270        let nstates = ncols;
271        let nout = nrows;
272        let f = move |x: &M::V, y: &mut M::V| {
273            for (i, j, v) in triplets {
274                y.set_index(*i, y.get_index(*i) + x.get_index(*j) * *v);
275            }
276        };
277        let mut ret = Closure::new(
278            move |x: &M::V, _p: &M::V, _t, y: &mut M::V| {
279                y.fill(M::T::zero());
280                f(x, y);
281            },
282            move |_x: &M::V, _p: &M::V, _t, v, y: &mut M::V| {
283                y.fill(M::T::zero());
284                f(v, y);
285            },
286            nstates,
287            nout,
288            p.len(),
289            p.context().clone(),
290        );
291        let y0 = M::V::zeros(nstates, p.context().clone());
292        let t0 = M::T::zero();
293        ret.calculate_sparsity(&y0, t0, p);
294        ret
295    }
296
297    #[allow(clippy::type_complexity)]
298    fn helper_triplets2op_linear<'a, M: Matrix + 'a>(
299        triplets: &'a [(usize, usize, M::T)],
300        p: &'a M::V,
301        nrows: usize,
302        ncols: usize,
303    ) -> LinearClosure<M, impl Fn(&M::V, &M::V, M::T, M::T, &mut M::V) + use<'a, M>> {
304        let nstates = ncols;
305        let nout = nrows;
306        let f = move |x: &M::V, y: &mut M::V| {
307            for (i, j, v) in triplets {
308                y.set_index(*i, y.get_index(*i) + x.get_index(*j) * *v);
309            }
310        };
311        let mut ret = LinearClosure::new(
312            move |x: &M::V, _p: &M::V, _t, beta, y: &mut M::V| {
313                y.mul_assign(scale(beta));
314                f(x, y);
315            },
316            nstates,
317            nout,
318            p.len(),
319            p.context().clone(),
320        );
321        let t0 = M::T::zero();
322        ret.calculate_sparsity(t0, p);
323        ret
324    }
325
326    fn find_non_zeros<M: Matrix>() {
327        let test_triplets = vec![
328            vec![(0, 0, M::T::one()), (1, 1, M::T::one())],
329            vec![
330                (0, 0, M::T::one()),
331                (0, 1, M::T::one()),
332                (1, 1, M::T::one()),
333            ],
334            vec![(1, 1, M::T::one())],
335            vec![
336                (0, 0, M::T::one()),
337                (1, 0, M::T::one()),
338                (0, 1, M::T::one()),
339                (1, 1, M::T::one()),
340            ],
341        ];
342        let p = M::V::zeros(0, Default::default());
343        for triplets in test_triplets {
344            let op = helper_triplets2op_nonlinear::<M>(triplets.as_slice(), &p, 2, 2);
345            let op = ParameterisedOp::new(&op, &p);
346            let non_zeros =
347                find_jacobian_non_zeros(&op, &M::V::zeros(2, p.context().clone()), M::T::zero());
348            let expect = triplets
349                .iter()
350                .map(|(i, j, _v)| (*i, *j))
351                .collect::<Vec<_>>();
352            assert_eq!(non_zeros, expect);
353        }
354    }
355
356    #[test]
357    fn find_non_zeros_dmatrix() {
358        find_non_zeros::<NalgebraMat<f64>>();
359    }
360
361    #[test]
362    fn find_non_zeros_faer_sparse() {
363        find_non_zeros::<FaerSparseMat<f64>>();
364    }
365
366    fn build_coloring<M: Matrix>() {
367        let test_triplets = [
368            vec![(0, 0, M::T::one()), (1, 1, M::T::one())],
369            vec![
370                (0, 0, M::T::one()),
371                (0, 1, M::T::one()),
372                (1, 1, M::T::one()),
373            ],
374            vec![(1, 1, M::T::one())],
375            vec![
376                (0, 0, M::T::one()),
377                (1, 0, M::T::one()),
378                (0, 1, M::T::one()),
379                (1, 1, M::T::one()),
380            ],
381        ];
382        let expect = vec![vec![1, 1], vec![1, 2], vec![1, 1], vec![1, 2]];
383        let p = M::V::zeros(0, Default::default());
384        for (triplets, expect) in test_triplets.iter().zip(expect) {
385            let op = helper_triplets2op_nonlinear::<M>(triplets.as_slice(), &p, 2, 2);
386            let op = ParameterisedOp::new(&op, &p);
387            let non_zeros =
388                find_jacobian_non_zeros(&op, &M::V::zeros(2, p.context().clone()), M::T::zero());
389            let ncols = op.nstates();
390            let graph = nonzeros2graph(non_zeros.as_slice(), ncols);
391            let coloring = color_graph_greedy(&graph);
392
393            assert_eq!(coloring, expect);
394        }
395    }
396
397    #[test]
398    fn build_coloring_dmatrix() {
399        build_coloring::<NalgebraMat<f64>>();
400    }
401
402    #[test]
403    fn build_coloring_faer_sparse() {
404        build_coloring::<FaerSparseMat<f64>>();
405    }
406
407    fn matrix_coloring<M: Matrix>() {
408        let test_triplets = vec![
409            vec![
410                (0, 0, M::T::one()),
411                (1, 1, M::T::one()),
412                (2, 2, M::T::one()),
413            ],
414            vec![(0, 0, M::T::one()), (1, 1, M::T::one())],
415            vec![
416                (0, 0, M::T::from(0.9)),
417                (1, 0, M::T::from(2.0)),
418                (1, 1, M::T::from(1.1)),
419                (2, 2, M::T::from(1.4)),
420            ],
421        ];
422        let n = 3;
423
424        // test nonlinear functions
425        let p = M::V::zeros(0, Default::default());
426        for triplets in test_triplets.iter() {
427            let op = helper_triplets2op_nonlinear::<M>(triplets.as_slice(), &p, n, n);
428            let op = ParameterisedOp::new(&op, &p);
429            let y0 = M::V::zeros(n, p.context().clone());
430            let t0 = M::T::zero();
431            let nonzeros = triplets
432                .iter()
433                .map(|(i, j, _v)| (*i, *j))
434                .collect::<Vec<_>>();
435            let coloring = JacobianColoring::new(
436                &op.jacobian_sparsity().unwrap(),
437                &nonzeros,
438                p.context().clone(),
439            );
440            let mut jac = M::new_from_sparsity(3, 3, op.jacobian_sparsity(), p.context().clone());
441            coloring.jacobian_inplace(&op, &y0, t0, &mut jac);
442            let mut gemv1 = M::V::zeros(n, p.context().clone());
443            let v = M::V::from_element(3, M::T::one(), p.context().clone());
444            op.jac_mul_inplace(&y0, t0, &v, &mut gemv1);
445            let mut gemv2 = M::V::zeros(n, p.context().clone());
446            jac.gemv(M::T::one(), &v, M::T::zero(), &mut gemv2);
447            gemv1.assert_eq_st(&gemv2, M::T::from(1e-10));
448        }
449
450        // test linear functions
451        let p = M::V::zeros(0, p.context().clone());
452        for triplets in test_triplets {
453            let op = helper_triplets2op_linear::<M>(triplets.as_slice(), &p, n, n);
454            let op = ParameterisedOp::new(&op, &p);
455            let t0 = M::T::zero();
456            let nonzeros = triplets
457                .iter()
458                .map(|(i, j, _v)| (*i, *j))
459                .collect::<Vec<_>>();
460            let coloring =
461                JacobianColoring::new(&op.sparsity().unwrap(), &nonzeros, p.context().clone());
462            let mut jac = M::new_from_sparsity(3, 3, op.sparsity(), p.context().clone());
463            coloring.matrix_inplace(&op, t0, &mut jac);
464            let mut gemv1 = M::V::zeros(n, p.context().clone());
465            let v = M::V::from_element(3, M::T::one(), p.context().clone());
466            op.gemv_inplace(&v, t0, M::T::zero(), &mut gemv1);
467            let mut gemv2 = M::V::zeros(n, p.context().clone());
468            jac.gemv(M::T::one(), &v, M::T::zero(), &mut gemv2);
469            gemv1.assert_eq_st(&gemv2, M::T::from(1e-10));
470        }
471    }
472
473    #[test]
474    fn matrix_coloring_dmatrix() {
475        matrix_coloring::<NalgebraMat<f64>>();
476    }
477
478    #[test]
479    fn matrix_coloring_faer_sparse() {
480        matrix_coloring::<FaerSparseMat<f64>>();
481    }
482}
diffsol/jacobian/mod.rs

diffsol/jacobian/
mod.rs
