Parallelize distance calculations and eliminate allocations in hot paths (#328)

* Initial plan

* Optimize distance calculations and reduce allocations

Co-authored-by: relf <[email protected]>

* Remove unused import

Co-authored-by: relf <[email protected]>

* Cleanup

* Simplify pdist and warn about the order

* Cleanup

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: relf <[email protected]>
Co-authored-by: relf <[email protected]>

Author: Copilot

Cargo.lock

@@ -28,6 +28,7 @@ ndarray-rand.workspace = true
 ndarray-stats.workspace = true
 num-traits.workspace = true
 rand_xoshiro.workspace = true
+rayon.workspace = true
 serde = { version = "1", optional = true }
 
 [dev-dependencies]

crates/doe/Cargo.toml

@@ -184,7 +184,7 @@ impl<F: Float, R: Rng> Lhs<F, R> {
         lhs_best
     }
 
-    fn _phip(&self, lhs: &ArrayBase<impl Data<Elem = F>, Ix2>, p: F) -> F {
+    fn _phip(&self, lhs: &ArrayBase<impl Data<Elem = F> + Sync, Ix2>, p: F) -> F {
         F::powf(pdist(lhs).mapv(|v| F::powf(v, -p)).sum(), F::one() / p)
     }
 
@@ -208,8 +208,14 @@ impl<F: Float, R: Rng> Lhs<F, R> {
         let mut dist1 = cdist(&x.slice(s![i1..i1 + 1, ..]), &x_rest);
         let mut dist2 = cdist(&x.slice(s![i2..i2 + 1, ..]), &x_rest);
 
-        let m1 = (x_rest.column(k).to_owned() - x[[i1, k]]).map(|v| *v * *v);
-        let m2 = (x_rest.column(k).to_owned() - x[[i2, k]]).map(|v| *v * *v);
+        let m1 = x_rest.column(k).mapv(|v| {
+            let diff = v - x[[i1, k]];
+            diff * diff
+        });
+        let m2 = x_rest.column(k).mapv(|v| {
+            let diff = v - x[[i2, k]];
+            diff * diff
+        });
 
         let two = F::cast(2.);
         let mut d1 = dist1.mapv(|v| v * v) - &m1 + &m2;

crates/doe/src/lhs.rs

@@ -1,26 +1,36 @@
 use linfa::Float;
-use ndarray::{Array1, Array2, ArrayBase, Data, Ix2};
+use ndarray::{Array, Array1, Array2, ArrayBase, Data, Ix2, Zip};
 use ndarray_stats::DeviationExt;
+use rayon::prelude::*;
 
-pub fn pdist<F: Float>(x: &ArrayBase<impl Data<Elem = F>, Ix2>) -> Array1<F> {
+/// Computes the pairwise distances between rows of a 2D-array using parallel processing
+/// Warning : The result is expected to b=e used in a context where order does not matter
+/// (e.g., get min distance) as the order of distances depends on the order of parallel execution
+pub fn pdist<F: Float>(x: &ArrayBase<impl Data<Elem = F> + Sync, Ix2>) -> Array1<F> {
     let nrows = x.nrows();
-    let size: usize = (nrows - 1) * nrows / 2;
-    let mut res: Array1<F> = Array1::zeros(size);
-    let mut k = 0;
-    for i in 0..nrows {
-        for j in (i + 1)..nrows {
+
+    // Parallelize the outer loop for better performance
+    let pairs: Vec<_> = (0..nrows)
+        .flat_map(|i| ((i + 1)..nrows).map(move |j| (i, j)))
+        .collect();
+
+    let distances: Vec<_> = pairs
+        .par_iter()
+        .map(|&(i, j)| {
             let a = x.row(i);
             let b = x.row(j);
-            res[k] = F::cast(a.l2_dist(&b).unwrap());
-            k += 1;
-        }
-    }
-    res
+            F::cast(a.l2_dist(&b).unwrap())
+        })
+        .collect();
+
+    Array::from_vec(distances)
 }
 
+/// Computes the pairwise distances between rows of two 2D arrays using parallel processing
+/// The resulting array has shape (ma, mb) where ma is the number of rows in xa and mb is the number of rows in xb
 pub fn cdist<F: Float>(
-    xa: &ArrayBase<impl Data<Elem = F>, Ix2>,
-    xb: &ArrayBase<impl Data<Elem = F>, Ix2>,
+    xa: &ArrayBase<impl Data<Elem = F> + Sync, Ix2>,
+    xb: &ArrayBase<impl Data<Elem = F> + Sync, Ix2>,
 ) -> Array2<F> {
     let ma = xa.nrows();
     let mb = xb.nrows();
@@ -29,14 +39,15 @@ pub fn cdist<F: Float>(
     if na != nb {
         panic!("cdist: operands should have same nb of columns. Found {na} and {nb}");
     }
+
     let mut res = Array2::zeros((ma, mb));
-    for i in 0..ma {
-        for j in 0..mb {
-            let a = xa.row(i);
-            let b = xb.row(j);
-            res[[i, j]] = F::cast(a.l2_dist(&b).unwrap());
-        }
-    }
+    Zip::from(res.rows_mut())
+        .and(xa.rows())
+        .par_for_each(|mut row_res, row_a| {
+            for (j, row_b) in xb.rows().into_iter().enumerate() {
+                row_res[j] = F::cast(row_a.l2_dist(&row_b).unwrap());
+            }
+        });
 
     res
 }

crates/doe/src/utils.rs

@@ -1,5 +1,5 @@
 use linfa::Float;
-use ndarray::{Array, Array1, Array2, ArrayBase, Axis, Data, Ix1, Ix2, s};
+use ndarray::{Array1, Array2, ArrayBase, Axis, Data, Ix1, Ix2, s};
 #[cfg(feature = "serializable")]
 use serde::{Deserialize, Serialize};
 
@@ -111,13 +111,22 @@ pub fn pairwise_differences<F: Float>(
     y: &ArrayBase<impl Data<Elem = F>, Ix2>,
 ) -> Array2<F> {
     assert!(x.ncols() == y.ncols());
-    let x3 = x.to_owned().insert_axis(Axis(1));
-    let y3 = y.to_owned().insert_axis(Axis(0));
-    let d = x3 - y3;
-    let n = d.len();
-    let res = Array::from_iter(d.iter().cloned());
-    res.into_shape_with_order((n / x.ncols(), x.ncols()))
-        .unwrap()
+
+    let nx = x.nrows();
+    let ny = y.nrows();
+    let ncols = x.ncols();
+    let mut result = Array2::zeros((nx * ny, ncols));
+
+    for (i, x_row) in x.rows().into_iter().enumerate() {
+        for (j, y_row) in y.rows().into_iter().enumerate() {
+            let idx = i * ny + j;
+            for k in 0..ncols {
+                result[[idx, k]] = x_row[k] - y_row[k];
+            }
+        }
+    }
+
+    result
 }
 
 /// Computes differences between x and each element of y