Be more generic

stefan-k · Mar 1, 2024 · e72abb7 · e72abb7
1 parent e15284d
commit e72abb7
Show file tree

Hide file tree

Showing 8 changed files with 40 additions and 53 deletions.
diff --git a/crates/finitediff/Cargo.toml b/crates/finitediff/Cargo.toml
@@ -14,3 +14,4 @@ categories = ["science"]
 
 [dependencies]
 ndarray = { version = "0.15.0", optional = true }
+num = "0.4.1"
diff --git a/crates/finitediff/src/ndarray_m/diff.rs b/crates/finitediff/src/ndarray_m/diff.rs
@@ -16,7 +16,7 @@ pub fn forward_diff_ndarray_f64(
     let mut xt = x.clone();
     (0..x.len())
         .map(|i| {
-            let fx1 = mod_and_calc_ndarray_f64(&mut xt, f, i, EPS_F64.sqrt());
+            let fx1 = mod_and_calc(&mut xt, f, i, EPS_F64.sqrt());
             (fx1 - fx) / (EPS_F64.sqrt())
         })
         .collect()
@@ -29,8 +29,8 @@ pub fn central_diff_ndarray_f64(
     let mut xt = x.clone();
     (0..x.len())
         .map(|i| {
-            let fx1 = mod_and_calc_ndarray_f64(&mut xt, f, i, EPS_F64.sqrt());
-            let fx2 = mod_and_calc_ndarray_f64(&mut xt, f, i, -EPS_F64.sqrt());
+            let fx1 = mod_and_calc(&mut xt, f, i, EPS_F64.sqrt());
+            let fx2 = mod_and_calc(&mut xt, f, i, -EPS_F64.sqrt());
             (fx1 - fx2) / (2.0 * EPS_F64.sqrt())
         })
         .collect()

diff --git a/crates/finitediff/src/ndarray_m/hessian.rs b/crates/finitediff/src/ndarray_m/hessian.rs
@@ -23,7 +23,7 @@ pub fn forward_hessian_ndarray_f64(
     let n = x.len();
     let mut out = Array2::zeros((n, rn));
     for i in 0..n {
-        let fx1 = mod_and_calc_ndarray_f64(&mut xt, grad, i, EPS_F64.sqrt());
+        let fx1 = mod_and_calc(&mut xt, grad, i, EPS_F64.sqrt());
         for j in 0..rn {
             out[(i, j)] = (fx1[j] - fx[j]) / EPS_F64.sqrt();
         }
@@ -43,8 +43,8 @@ pub fn central_hessian_ndarray_f64(
     let n = x.len();
     let mut out = ndarray::Array2::zeros((n, rn));
     for i in 0..n {
-        let fx1 = mod_and_calc_ndarray_f64(&mut xt, grad, i, EPS_F64.sqrt());
-        let fx2 = mod_and_calc_ndarray_f64(&mut xt, grad, i, -EPS_F64.sqrt());
+        let fx1 = mod_and_calc(&mut xt, grad, i, EPS_F64.sqrt());
+        let fx2 = mod_and_calc(&mut xt, grad, i, -EPS_F64.sqrt());
         for j in 0..rn {
             out[(i, j)] = (fx1[j] - fx2[j]) / (2.0 * EPS_F64.sqrt());
         }
@@ -86,7 +86,7 @@ pub fn forward_hessian_nograd_ndarray_f64(
 
     // Precompute f(x + sqrt(EPS) * e_i) for all i
     let fxei: Vec<f64> = (0..n)
-        .map(|i| mod_and_calc_ndarray_f64(&mut xt, f, i, EPS_F64_NOGRAD.sqrt()))
+        .map(|i| mod_and_calc(&mut xt, f, i, EPS_F64_NOGRAD.sqrt()))
         .collect();
 
     let mut out = ndarray::Array2::zeros((n, n));
@@ -129,10 +129,7 @@ pub fn forward_hessian_nograd_sparse_ndarray_f64(
     let mut fxei = KV::new(idxs.len());
 
     for idx in idxs.iter() {
-        fxei.set(
-            *idx,
-            mod_and_calc_ndarray_f64(&mut xt, f, *idx, EPS_F64_NOGRAD.sqrt()),
-        );
+        fxei.set(*idx, mod_and_calc(&mut xt, f, *idx, EPS_F64_NOGRAD.sqrt()));
     }
 
     let mut out = ndarray::Array2::zeros((n, n));

diff --git a/crates/finitediff/src/ndarray_m/jacobian.rs b/crates/finitediff/src/ndarray_m/jacobian.rs
@@ -21,7 +21,7 @@ pub fn forward_jacobian_ndarray_f64(
     let n = x.len();
     let mut out = Array2::zeros((n, rn));
     for i in 0..n {
-        let fx1 = mod_and_calc_ndarray_f64(&mut xt, fs, i, EPS_F64.sqrt());
+        let fx1 = mod_and_calc(&mut xt, fs, i, EPS_F64.sqrt());
         for j in 0..rn {
             out[(i, j)] = (fx1[j] - fx[j]) / EPS_F64.sqrt();
         }
@@ -44,8 +44,8 @@ pub fn central_jacobian_ndarray_f64(
 
     let mut out = Array2::zeros((n, rn));
     for i in 0..n {
-        let fx1 = mod_and_calc_ndarray_f64(&mut xt, fs, i, EPS_F64.sqrt());
-        let fx2 = mod_and_calc_ndarray_f64(&mut xt, fs, i, -EPS_F64.sqrt());
+        let fx1 = mod_and_calc(&mut xt, fs, i, EPS_F64.sqrt());
+        let fx2 = mod_and_calc(&mut xt, fs, i, -EPS_F64.sqrt());
         for j in 0..rn {
             out[(i, j)] = (fx1[j] - fx2[j]) / (2.0 * EPS_F64.sqrt());
         }

diff --git a/crates/finitediff/src/utils.rs b/crates/finitediff/src/utils.rs
@@ -5,28 +5,17 @@
 // http://opensource.org/licenses/MIT>, at your option. This file may not be
 // copied, modified, or distributed except according to those terms.
 
-#[inline(always)]
-pub fn mod_and_calc_vec_f64<T>(
-    x: &mut Vec<f64>,
-    f: &dyn Fn(&Vec<f64>) -> T,
-    idx: usize,
-    y: f64,
-) -> T {
-    let xtmp = x[idx];
-    x[idx] = xtmp + y;
-    let fx1 = (f)(x);
-    x[idx] = xtmp;
-    fx1
-}
+use std::ops::{Add, IndexMut};
 
-#[cfg(feature = "ndarray")]
+use num::{Float, FromPrimitive};
+
+/// Panics when idx > x.len()
 #[inline(always)]
-pub fn mod_and_calc_ndarray_f64<T>(
-    x: &mut ndarray::Array1<f64>,
-    f: &dyn Fn(&ndarray::Array1<f64>) -> T,
-    idx: usize,
-    y: f64,
-) -> T {
+pub fn mod_and_calc<F, C, T>(x: &mut C, f: &dyn Fn(&C) -> T, idx: usize, y: F) -> T
+where
+    F: Add<Output = F> + Copy,
+    C: IndexMut<usize, Output = F>,
+{
     let xtmp = x[idx];
     x[idx] = xtmp + y;
     let fx1 = (f)(x);
@@ -35,10 +24,13 @@ pub fn mod_and_calc_ndarray_f64<T>(
 }
 
 #[inline(always)]
-pub fn restore_symmetry_vec_f64(mut mat: Vec<Vec<f64>>) -> Vec<Vec<f64>> {
+pub fn restore_symmetry_vec<F>(mut mat: Vec<Vec<F>>) -> Vec<Vec<F>>
+where
+    F: Float + FromPrimitive,
+{
     for i in 0..mat.len() {
         for j in (i + 1)..mat[i].len() {
-            let t = (mat[i][j] + mat[j][i]) / 2.0;
+            let t = (mat[i][j] + mat[j][i]) / F::from_f64(2.0).unwrap();
             mat[i][j] = t;
             mat[j][i] = t;
         }
@@ -47,10 +39,10 @@ pub fn restore_symmetry_vec_f64(mut mat: Vec<Vec<f64>>) -> Vec<Vec<f64>> {
 }
 
 #[cfg(feature = "ndarray")]
-#[inline(always)]
 /// Restore symmetry for an array of type `ndarray::Array2<f64>`
 ///
 /// Unfortunately, this is *really* slow!
+#[inline(always)]
 pub fn restore_symmetry_ndarray_f64(mut mat: ndarray::Array2<f64>) -> ndarray::Array2<f64> {
     let (nx, ny) = mat.dim();
     for i in 0..nx {

diff --git a/crates/finitediff/src/vec/diff.rs b/crates/finitediff/src/vec/diff.rs
@@ -13,7 +13,7 @@ pub fn forward_diff_vec_f64(x: &Vec<f64>, f: &dyn Fn(&Vec<f64>) -> f64) -> Vec<f
     let mut xt = x.clone();
     (0..x.len())
         .map(|i| {
-            let fx1 = mod_and_calc_vec_f64(&mut xt, f, i, EPS_F64.sqrt());
+            let fx1 = mod_and_calc(&mut xt, f, i, EPS_F64.sqrt());
             (fx1 - fx) / (EPS_F64.sqrt())
         })
         .collect()
@@ -23,8 +23,8 @@ pub fn central_diff_vec_f64(x: &[f64], f: &dyn Fn(&Vec<f64>) -> f64) -> Vec<f64>
     let mut xt = x.to_owned();
     (0..x.len())
         .map(|i| {
-            let fx1 = mod_and_calc_vec_f64(&mut xt, f, i, EPS_F64.sqrt());
-            let fx2 = mod_and_calc_vec_f64(&mut xt, f, i, -EPS_F64.sqrt());
+            let fx1 = mod_and_calc(&mut xt, f, i, EPS_F64.sqrt());
+            let fx2 = mod_and_calc(&mut xt, f, i, -EPS_F64.sqrt());
             (fx1 - fx2) / (2.0 * EPS_F64.sqrt())
         })
         .collect()

diff --git a/crates/finitediff/src/vec/hessian.rs b/crates/finitediff/src/vec/hessian.rs
@@ -19,7 +19,7 @@ pub fn forward_hessian_vec_f64(
     let mut xt = x.clone();
     let out: Vec<Vec<f64>> = (0..x.len())
         .map(|i| {
-            let fx1 = mod_and_calc_vec_f64(&mut xt, grad, i, EPS_F64.sqrt());
+            let fx1 = mod_and_calc(&mut xt, grad, i, EPS_F64.sqrt());
             fx1.iter()
                 .zip(fx.iter())
                 .map(|(a, b)| (a - b) / (EPS_F64.sqrt()))
@@ -28,15 +28,15 @@ pub fn forward_hessian_vec_f64(
         .collect();
 
     // restore symmetry
-    restore_symmetry_vec_f64(out)
+    restore_symmetry_vec(out)
 }
 
 pub fn central_hessian_vec_f64(x: &[f64], grad: &dyn Fn(&Vec<f64>) -> Vec<f64>) -> Vec<Vec<f64>> {
     let mut xt = x.to_owned();
     let out: Vec<Vec<f64>> = (0..x.len())
         .map(|i| {
-            let fx1 = mod_and_calc_vec_f64(&mut xt, grad, i, EPS_F64.sqrt());
-            let fx2 = mod_and_calc_vec_f64(&mut xt, grad, i, -EPS_F64.sqrt());
+            let fx1 = mod_and_calc(&mut xt, grad, i, EPS_F64.sqrt());
+            let fx2 = mod_and_calc(&mut xt, grad, i, -EPS_F64.sqrt());
             fx1.iter()
                 .zip(fx2.iter())
                 .map(|(a, b)| (a - b) / (2.0 * EPS_F64.sqrt()))
@@ -45,7 +45,7 @@ pub fn central_hessian_vec_f64(x: &[f64], grad: &dyn Fn(&Vec<f64>) -> Vec<f64>)
         .collect();
 
     // restore symmetry
-    restore_symmetry_vec_f64(out)
+    restore_symmetry_vec(out)
 }
 
 pub fn forward_hessian_vec_prod_vec_f64(
@@ -102,7 +102,7 @@ pub fn forward_hessian_nograd_vec_f64(x: &Vec<f64>, f: &dyn Fn(&Vec<f64>) -> f64
 
     // Precompute f(x + sqrt(EPS) * e_i) for all i
     let fxei: Vec<f64> = (0..n)
-        .map(|i| mod_and_calc_vec_f64(&mut xt, f, i, EPS_F64_NOGRAD.sqrt()))
+        .map(|i| mod_and_calc(&mut xt, f, i, EPS_F64_NOGRAD.sqrt()))
         .collect();
 
     let mut out: Vec<Vec<f64>> = vec![vec![0.0; n]; n];
@@ -146,10 +146,7 @@ pub fn forward_hessian_nograd_sparse_vec_f64(
     let mut fxei = KV::new(idxs.len());
 
     for idx in idxs.iter() {
-        fxei.set(
-            *idx,
-            mod_and_calc_vec_f64(&mut xt, f, *idx, EPS_F64_NOGRAD.sqrt()),
-        );
+        fxei.set(*idx, mod_and_calc(&mut xt, f, *idx, EPS_F64_NOGRAD.sqrt()));
     }
 
     let mut out: Vec<Vec<f64>> = vec![vec![0.0; n]; n];

diff --git a/crates/finitediff/src/vec/jacobian.rs b/crates/finitediff/src/vec/jacobian.rs
@@ -14,7 +14,7 @@ pub fn forward_jacobian_vec_f64(x: &Vec<f64>, fs: &dyn Fn(&Vec<f64>) -> Vec<f64>
     let mut xt = x.clone();
     (0..x.len())
         .map(|i| {
-            let fx1 = mod_and_calc_vec_f64(&mut xt, fs, i, EPS_F64.sqrt());
+            let fx1 = mod_and_calc(&mut xt, fs, i, EPS_F64.sqrt());
             fx1.iter()
                 .zip(fx.iter())
                 .map(|(a, b)| (a - b) / EPS_F64.sqrt())
@@ -27,8 +27,8 @@ pub fn central_jacobian_vec_f64(x: &[f64], fs: &dyn Fn(&Vec<f64>) -> Vec<f64>) -
     let mut xt = x.to_owned();
     (0..x.len())
         .map(|i| {
-            let fx1 = mod_and_calc_vec_f64(&mut xt, fs, i, EPS_F64.sqrt());
-            let fx2 = mod_and_calc_vec_f64(&mut xt, fs, i, -EPS_F64.sqrt());
+            let fx1 = mod_and_calc(&mut xt, fs, i, EPS_F64.sqrt());
+            let fx2 = mod_and_calc(&mut xt, fs, i, -EPS_F64.sqrt());
             fx1.iter()
                 .zip(fx2.iter())
                 .map(|(a, b)| (a - b) / (2.0 * EPS_F64.sqrt()))