Merge pull request #8 from martinjrobins/sparse-jacobian

Sparse jacobian

Cargo.toml

@@ -46,7 +46,8 @@ diffsl14-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffs
 diffsl15-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm15-0"], optional = true }
 diffsl16-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm16-0"], optional = true }
 diffsl17-0 = { package = "diffsl", git = "https://github.com/martinjrobins/diffsl.git", version = ">=0.1.1", features = ["llvm17-0"], optional = true }
+petgraph = "0.6.4"
 
 
 [dev-dependencies]
-insta = { version = "1.34.0", features = ["yaml"] }
+insta = { version = "1.34.0", features = ["yaml"] }

README.md

@@ -72,28 +72,27 @@ respectively. We set the problem up with the following code:
 ```rust
 type T = f64;
 type V = nalgebra::DVector<T>;
-let p = V::from_vec(vec![0.04, 1.0e4, 3.0e7]);
-let mut problem = OdeSolverProblem::new_ode(
-    // The rhs function `f` 
-    | x: &V, p: &V, _t: T, y: &mut V | {
-        y[0] = -p[0] * x[0] + p[1] * x[1] * x[2];
-        y[1] = p[0] * x[0] - p[1] * x[1] * x[2] - p[2] * x[1] * x[1];
-        y[2] = p[2] * x[1] * x[1];
-    },
-    // The jacobian function `Jv`
-    | x: &V, p: &V, _t: T, v: &V, y: &mut V | {
-        y[0] = -p[0] * v[0] + p[1] * v[1] * x[2] + p[1] * x[1] * v[2];
-        y[1] = p[0] * v[0] - p[1] * v[1] * x[2] - p[1] * x[1] * v[2]  - 2.0 * p[2] * x[1] * v[1];
-        y[2] = 2.0 * p[2] * x[1] * v[1];
-    },
-    // The initial condition
-    | _p: &V, _t: T | {
-        V::from_vec(vec![1.0, 0.0, 0.0])
-    },
-    p,
-);
-problem.rtol = 1.0e-4;
-problem.atol = Rc::new(V::from_vec(vec![1.0e-8, 1.0e-6, 1.0e-6]));
+let problem = OdeBuilder::new()
+    .p([0.04, 1.0e4, 3.0e7])
+    .rtol(1e-4)
+    .atol([1.0e-8, 1.0e-6, 1.0e-6])
+    .build_ode(
+        |x: &V, p: &V, _t: T, y: &mut V| {
+            y[0] = -p[0] * x[0] + p[1] * x[1] * x[2];
+            y[1] = p[0] * x[0] - p[1] * x[1] * x[2] - p[2] * x[1] * x[1];
+            y[2] = p[2] * x[1] * x[1];
+        },
+        |x: &V, p: &V, _t: T, v: &V, y: &mut V| {
+            y[0] = -p[0] * v[0] + p[1] * v[1] * x[2] + p[1] * x[1] * v[2];
+            y[1] = p[0] * v[0]
+                - p[1] * v[1] * x[2]
+                - p[1] * x[1] * v[2]
+                - 2.0 * p[2] * x[1] * v[1];
+            y[2] = 2.0 * p[2] * x[1] * v[1];
+        },
+        |_p: &V, _t: T| V::from_vec(vec![1.0, 0.0, 0.0]),
+    )
+    .unwrap();
 
 let mut solver = Bdf::default();
 ```

src/jacobian/coloring.rs

@@ -0,0 +1,47 @@
+// Translated from https://github.com/JuliaDiff/SparseDiffTools.jl under an MIT license
+
+use petgraph::graph::NodeIndex;
+
+use super::graph::Graph;
+
+/// Returns a vector of rows where each row contains
+/// a vector of its column indices.
+fn cols_by_rows(non_zeros: &[(usize, usize)]) -> Vec<Vec<usize>> {
+    let nrows = if non_zeros.is_empty() {
+        0
+    } else {
+        *non_zeros.iter().map(|(i, _j)| i).max().unwrap() + 1
+    };
+    let mut cols_by_rows = vec![Vec::new(); nrows];
+    for (i, j) in non_zeros.iter() {
+        cols_by_rows[*i].push(*j);
+    }
+    cols_by_rows
+}
+
+/// A utility function to generate a graph from input sparse matrix, columns are represented
+/// with vertices and 2 vertices are connected with an edge only if the two columns are mutually
+/// orthogonal.
+///
+/// non_zeros: A vector of indices (i, j) where i is the row index and j is the column index
+pub fn nonzeros2graph(non_zeros: &[(usize, usize)], ncols: usize) -> Graph {
+    let cols_by_rows = cols_by_rows(non_zeros);
+    let mut edges = Vec::new();
+    for (cur_row, cur_col) in non_zeros.iter() {
+        if !cols_by_rows[*cur_row].is_empty() {
+            for next_col in cols_by_rows[*cur_row].iter() {
+                if next_col < cur_col {
+                    edges.push((*cur_col, *next_col));
+                }
+            }
+        }
+    }
+    let mut graph = Graph::with_capacity(ncols, edges.len());
+    for _ in 0..ncols {
+        graph.add_node(());
+    }
+    for (i, j) in edges.iter() {
+        graph.add_edge(NodeIndex::new(*i), NodeIndex::new(*j), ());
+    }
+    graph
+}

src/jacobian/graph.rs

@@ -0,0 +1 @@
+pub type Graph = petgraph::graph::Graph<(), (), petgraph::Directed>;

src/jacobian/greedy_coloring.rs

@@ -0,0 +1,32 @@
+// Translated from https://github.com/JuliaDiff/SparseDiffTools.jl under an MIT license
+
+use petgraph::graph::NodeIndex;
+
+use super::graph::Graph;
+
+/// Find a coloring of a given input graph such that
+/// no two vertices connected by an edge have the same
+/// color using greedy approach. The number of colors
+/// used may be equal or greater than the chromatic
+/// number `χ(G)` of the graph.
+pub fn color_graph_greedy(graph: &Graph) -> Vec<usize> {
+    let mut result = vec![0; graph.node_count()];
+    result[0] = 1;
+    let mut available = vec![false; graph.node_count()];
+
+    for ii in 1..graph.node_count() {
+        for j in graph.neighbors(NodeIndex::new(ii)) {
+            if result[j.index()] != 0 {
+                available[result[j.index()] - 1] = true;
+            }
+        }
+        for (i, a) in available.iter().enumerate() {
+            if !a {
+                result[ii] = i + 1;
+                break;
+            }
+        }
+        available.iter_mut().for_each(|x| *x = false);
+    }
+    result
+}

src/jacobian/mod.rs

@@ -0,0 +1,189 @@
+use crate::op::NonLinearOp;
+use crate::vector::Vector;
+use crate::Scalar;
+use num_traits::{One, Zero};
+
+use self::{coloring::nonzeros2graph, greedy_coloring::color_graph_greedy};
+
+pub mod coloring;
+pub mod graph;
+pub mod greedy_coloring;
+
+/// Find the non-zero entries of the Jacobian matrix of a non-linear operator.
+/// This is used as the default `find_non_zeros` function for the `NonLinearOp` and `LinearOp` traits.
+/// Users can override this function with a more efficient and reliable implementation if desired.
+pub fn find_non_zeros<F: NonLinearOp + ?Sized>(op: &F, x: &F::V, t: F::T) -> Vec<(usize, usize)> {
+    let mut v = F::V::zeros(op.nstates());
+    let mut col = F::V::zeros(op.nout());
+    let mut triplets = Vec::with_capacity(op.nstates());
+    for j in 0..op.nstates() {
+        v[j] = F::T::NAN;
+        op.jac_mul_inplace(x, t, &v, &mut col);
+        for i in 0..op.nout() {
+            if col[i].is_nan() {
+                triplets.push((i, j));
+            }
+            col[i] = F::T::zero();
+        }
+        v[j] = F::T::zero();
+    }
+    triplets
+}
+
+/// Find the non-zero entries of the Jacobian matrix of a non-linear operator.
+/// This is used in the default `jacobian` method of the `NonLinearOp` and `LinearOp` traits.
+pub fn find_non_zero_entries<F: NonLinearOp + ?Sized>(
+    op: &F,
+    x: &F::V,
+    t: F::T,
+) -> Vec<(usize, usize, F::T)> {
+    let mut v = F::V::zeros(op.nstates());
+    let mut col = F::V::zeros(op.nout());
+    let mut triplets = Vec::with_capacity(op.nstates());
+    for j in 0..op.nstates() {
+        v[j] = F::T::one();
+        op.jac_mul_inplace(x, t, &v, &mut col);
+        for i in 0..op.nout() {
+            if col[i] != F::T::zero() {
+                triplets.push((i, j, col[i]));
+            }
+        }
+        v[j] = F::T::zero();
+    }
+    triplets
+}
+
+pub struct JacobianColoring {
+    cols_per_color: Vec<Vec<usize>>,
+    ij_per_color: Vec<Vec<(usize, usize)>>,
+}
+
+impl JacobianColoring {
+    pub fn new<F: NonLinearOp>(op: &F, x: &F::V, t: F::T) -> Self {
+        let non_zeros = op.find_non_zeros(x, t);
+        let ncols = op.nstates();
+        let graph = nonzeros2graph(non_zeros.as_slice(), ncols);
+        let coloring = color_graph_greedy(&graph);
+        let max_color = coloring.iter().max().copied().unwrap_or(0);
+        let mut cols_per_color = vec![Vec::new(); max_color];
+        let mut ij_per_color = vec![Vec::new(); max_color];
+        for c in 1..=max_color {
+            for (i, j) in non_zeros.iter() {
+                if coloring[*j] == c {
+                    cols_per_color[c - 1].push(*j);
+                    ij_per_color[c - 1].push((*i, *j));
+                }
+            }
+        }
+        Self {
+            cols_per_color,
+            ij_per_color,
+        }
+    }
+
+    pub fn find_non_zero_entries<F: NonLinearOp>(
+        &self,
+        op: &F,
+        x: &F::V,
+        t: F::T,
+    ) -> Vec<(usize, usize, F::T)> {
+        let mut triplets = Vec::with_capacity(op.nstates());
+        let mut v = F::V::zeros(op.nstates());
+        let mut col = F::V::zeros(op.nout());
+        for (cols, ijs) in self.cols_per_color.iter().zip(self.ij_per_color.iter()) {
+            for j in cols {
+                v[*j] = F::T::one();
+            }
+            op.jac_mul_inplace(x, t, &v, &mut col);
+            for (i, j) in ijs {
+                if col[*i] != F::T::zero() {
+                    triplets.push((*i, *j, col[*i]));
+                }
+            }
+            for j in cols {
+                v[*j] = F::T::zero();
+            }
+        }
+        triplets
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::rc::Rc;
+
+    use crate::op::Op;
+    use crate::{
+        jacobian::{coloring::nonzeros2graph, greedy_coloring::color_graph_greedy},
+        op::{closure::Closure, NonLinearOp},
+    };
+    use nalgebra::{DMatrix, DVector};
+
+    fn helper_triplets2op(
+        triplets: &[(usize, usize, f64)],
+        nrows: usize,
+        ncols: usize,
+    ) -> impl NonLinearOp<M = DMatrix<f64>, V = DVector<f64>, T = f64> + '_ {
+        let nstates = ncols;
+        let nout = nrows;
+        let f = move |x: &DVector<f64>, y: &mut DVector<f64>| {
+            for (i, j, v) in triplets {
+                y[*i] += x[*j] * v;
+            }
+        };
+        Closure::new(
+            move |x: &DVector<f64>, _p: &DVector<f64>, _t, y| {
+                f(x, y);
+            },
+            move |_x: &DVector<f64>, _p: &DVector<f64>, _t, v, y| {
+                f(v, y);
+            },
+            nstates,
+            nout,
+            Rc::new(DVector::zeros(0)),
+        )
+    }
+    #[test]
+    fn find_non_zeros() {
+        let test_triplets = vec![
+            vec![(0, 0, 1.0), (1, 1, 1.0)],
+            vec![(0, 0, 1.0), (0, 1, 1.0), (1, 1, 1.0)],
+            vec![(1, 1, 1.0)],
+            vec![(0, 0, 1.0), (1, 0, 1.0), (0, 1, 1.0), (1, 1, 1.0)],
+        ];
+        for triplets in test_triplets {
+            let op = helper_triplets2op(triplets.as_slice(), 2, 2);
+            let x = DVector::from_vec(vec![1.0, 1.0]);
+            let t = 0.0;
+            let non_zeros = op.find_non_zeros(&x, t);
+            let expect = triplets
+                .iter()
+                .map(|(i, j, _v)| (*i, *j))
+                .collect::<Vec<_>>();
+            assert_eq!(non_zeros, expect);
+        }
+    }
+
+    #[test]
+    fn build_coloring() {
+        let test_triplets = vec![
+            vec![(0, 0, 1.0), (1, 1, 1.0)],
+            vec![(0, 0, 1.0), (0, 1, 1.0), (1, 1, 1.0)],
+            vec![(1, 1, 1.0)],
+            vec![(0, 0, 1.0), (1, 0, 1.0), (0, 1, 1.0), (1, 1, 1.0)],
+        ];
+        let expect = vec![vec![1, 1], vec![1, 2], vec![1, 1], vec![1, 2]];
+        for (triplets, expect) in test_triplets.iter().zip(expect) {
+            let op = helper_triplets2op(triplets.as_slice(), 2, 2);
+            let x = DVector::from_vec(vec![1.0, 1.0]);
+            let t = 0.0;
+
+            let non_zeros = op.find_non_zeros(&x, t);
+            let ncols = op.nstates();
+            let graph = nonzeros2graph(non_zeros.as_slice(), ncols);
+            let coloring = color_graph_greedy(&graph);
+
+            assert_eq!(coloring, expect);
+        }
+    }
+}