perf: remove allocation in diff update for bdf

examples/robertson.rs

@@ -3,7 +3,7 @@ use diffsol::{ode_solver::test_models::robertson::robertson, Bdf, OdeSolverMetho
 fn main() {
     let mut s = Bdf::default();
     let (problem, _soln) = robertson::<nalgebra::DMatrix<f64>>(false);
-    for _ in 0..10000 {
+    for _ in 0..100000 {
         let _y = s.solve(&problem, 4.0000e+10);
     }
 }

src/matrix/dense_faer_serial.rs

@@ -128,6 +128,23 @@ impl<T: Scalar> DenseMatrix for Mat<T> {
     fn columns(&self, start: usize, nrows: usize) -> MatRef<'_, T> {
         self.get(0..self.nrows(), start..nrows)
     }
+
+    fn column_axpy(&mut self, alpha: Self::T, j: IndexType, beta: Self::T, i: IndexType) {
+        if i > self.ncols() {
+            panic!("Column index out of bounds");
+        }
+        if j > self.ncols() {
+            panic!("Column index out of bounds");
+        }
+        if i == j {
+            panic!("Column index cannot be the same");
+        }
+        for k in 0..self.nrows() {
+            let value =
+                unsafe { beta * self.read_unchecked(k, i) + alpha * self.read_unchecked(k, j) };
+            unsafe { self.write_unchecked(k, i, value) };
+        }
+    }
 }
 
 impl<T: Scalar> Matrix for Mat<T> {
@@ -202,3 +219,28 @@ impl<T: Scalar> Matrix for Mat<T> {
         Self::zeros(nrows, ncols)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_column_axpy() {
+        // M = [1 2]
+        //     [3 4]
+        let mut a = Mat::zeros(2, 2);
+        a[(0, 0)] = 1.0;
+        a[(0, 1)] = 2.0;
+        a[(1, 0)] = 3.0;
+        a[(1, 1)] = 4.0;
+
+        // op is M(:, 1) = 2 * M(:, 0) + M(:, 1)
+        a.column_axpy(2.0, 0, 1.0, 1);
+        // M = [1 4]
+        //     [3 10]
+        assert_eq!(a[(0, 0)], 1.0);
+        assert_eq!(a[(0, 1)], 4.0);
+        assert_eq!(a[(1, 0)], 3.0);
+        assert_eq!(a[(1, 1)], 10.0);
+    }
+}

src/matrix/dense_nalgebra_serial.rs

@@ -1,7 +1,10 @@
 use std::ops::{AddAssign, Mul, MulAssign};
 
 use anyhow::Result;
-use nalgebra::{DMatrix, DMatrixView, DMatrixViewMut, DVector, DVectorView, DVectorViewMut};
+use nalgebra::{
+    DMatrix, DMatrixView, DMatrixViewMut, DVector, DVectorView, DVectorViewMut, RawStorage,
+    RawStorageMut,
+};
 
 use crate::op::NonLinearOp;
 use crate::vector::Vector;
@@ -190,4 +193,43 @@ impl<T: Scalar> DenseMatrix for DMatrix<T> {
     fn columns(&self, start: IndexType, ncols: IndexType) -> Self::View<'_> {
         self.columns(start, ncols)
     }
+    fn column_axpy(&mut self, alpha: Self::T, j: IndexType, beta: Self::T, i: IndexType) {
+        if i > self.ncols() {
+            panic!("Column index out of bounds");
+        }
+        if j > self.ncols() {
+            panic!("Column index out of bounds");
+        }
+        if i == j {
+            panic!("Column index cannot be the same");
+        }
+        for k in 0..self.nrows() {
+            let value = unsafe {
+                beta * *self.data.get_unchecked(k, i) + alpha * *self.data.get_unchecked(k, j)
+            };
+            unsafe {
+                *self.data.get_unchecked_mut(k, i) = value;
+            };
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_column_axpy() {
+        // M = [1 2]
+        //     [3 4]
+        let mut a = DMatrix::from_row_slice(2, 2, &[1.0, 2.0, 3.0, 4.0]);
+        // op is M(:, 1) = 2 * M(:, 0) + M(:, 1)
+        a.column_axpy(2.0, 0, 1.0, 1);
+        // M = [1 4]
+        //     [3 10]
+        assert_eq!(a[(0, 0)], 1.0);
+        assert_eq!(a[(0, 1)], 4.0);
+        assert_eq!(a[(1, 0)], 3.0);
+        assert_eq!(a[(1, 1)], 10.0);
+    }
 }

src/matrix/mod.rs

@@ -284,6 +284,9 @@ pub trait DenseMatrix:
     /// Perform a matrix-matrix multiplication `self = alpha * a * b + beta * self`, where `alpha` and `beta` are scalars, and `a` and `b` are matrices
     fn gemm(&mut self, alpha: Self::T, a: &Self, b: &Self, beta: Self::T);
 
+    /// Performs an axpy operation on two columns of the matrix `M[:, i] = alpha * M[:, j] + M[:, i]`
+    fn column_axpy(&mut self, alpha: Self::T, j: IndexType, beta: Self::T, i: IndexType);
+
     /// Get a matrix view of the columns starting at `start` and ending at `start + ncols`
     fn columns(&self, start: IndexType, ncols: IndexType) -> Self::View<'_>;
 

src/ode_solver/bdf.rs

@@ -291,8 +291,7 @@ where
             .copy_from(&d_minus_order_plus_one);
         diff.column_mut(order + 1).copy_from(d);
         for i in (0..=order).rev() {
-            let tmp = diff.column(i + 1).into_owned();
-            diff.column_mut(i).add_assign(&tmp);
+            diff.column_axpy(Eqn::T::one(), i + 1, Eqn::T::one(), i);
         }
     }