pointpca2 0.4.2__tar.gz → 0.4.4__tar.gz

{pointpca2-0.4.2 → pointpca2-0.4.4}/Cargo.lock

@@ -450,12 +450,13 @@ dependencies = [
 
 [[package]]
 name = "pointpca2"
-version = "0.4.2"
+version = "0.4.4"
 dependencies = [
  "nalgebra",
  "numpy",
  "pointpca2-rs",
  "pyo3",
+ "rayon",
 ]
 
 [[package]]

{pointpca2-0.4.2 → pointpca2-0.4.4}/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pointpca2"
-version = "0.4.2"
+version = "0.4.4"
 edition = "2021"
 
 [lib]
@@ -10,6 +10,7 @@ crate-type = ["cdylib"]
 [dependencies]
 nalgebra = "0.32.5"
 numpy = "0.15"
+rayon = "1.10.0"
 
 [dependencies.pointpca2-rs]
 git = "https://github.com/akaTsunemori/pointpca2-rs.git"

{pointpca2-0.4.2 → pointpca2-0.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: pointpca2
-Version: 0.4.2
+Version: 0.4.4
 Requires-Dist: numpy >=1.26
 Requires-Dist: toml ==0.10.0
 License-File: LICENSE

{pointpca2-0.4.2 → pointpca2-0.4.4}/python/pointpca2/__init__.py

@@ -2,30 +2,30 @@ import numpy as np
 from .pointpca2 import compute_pointpca2 as __pointpca2_internal
 
 
-def __preprocess_point_cloud(points, colors):
-    if type(points) is np.ndarray:
-        points = np.array(points, dtype=np.double)
-    if type(colors) is np.ndarray:
-        colors = np.array(colors, dtype=np.double)
+def __preprocess_point_cloud(
+    points: np.ndarray,
+    colors: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
     if points.shape != colors.shape:
         raise Exception("Points and colors must have the same shape.")
+    if colors.max() <= 1 and colors.min() >= 0 and colors.dtype is np.double:
+        colors *= 255
     if points.dtype != np.double:
         points = points.astype(np.double)
-    if colors.max() <= 1 and colors.min() >= 0 and colors.dtype == np.double:
-        colors *= 255
     if colors.dtype != np.uint8:
         colors = colors.astype(np.uint8)
     return points, colors
 
 
 def compute_pointpca2(
-    points_reference,
-    colors_reference,
-    points_test,
-    colors_test,
+    points_reference: np.ndarray,
+    colors_reference: np.ndarray,
+    points_test: np.ndarray,
+    colors_test: np.ndarray,
     search_size=81,
+    max_workers=0,
     verbose=False,
-):
+) -> np.ndarray:
     """
     Compute PointPCA2 from reference and test point clouds.
 
@@ -42,7 +42,13 @@ def compute_pointpca2(
     search_size : int, optional
         The k-nearest-neighbors search size.
         Default is 81.
-    verbose: bool, optional
+    max_workers : int, optional
+        Sets the number of threads to be used in the thread-pool.
+        If you specify a non-zero number of threads using this parameter, then the
+        resulting thread-pools are guaranteed to start at most this number of threads.
+        If max_workers is 0, then the number of logical CPUs will be used.
+        Default is 0.
+    verbose : bool, optional
         Whether to display verbose information or not.
         Default is False.
 
@@ -57,7 +63,7 @@ def compute_pointpca2(
     the array will eventually be converted to np.double.
 
     For the colors arguments, it is expected that the colors are
-    on the [0, 1] range, or [0, 255]. Other ranges are not supported.
+    on the [0, 1] range (np.double), or [0, 255]. Other ranges are not supported.
     Any dtype is accepted, but the array will eventually be converted
     to np.uint8.
 
@@ -69,6 +75,12 @@ def compute_pointpca2(
     points_a, colors_a = __preprocess_point_cloud(points_reference, colors_reference)
     points_b, colors_b = __preprocess_point_cloud(points_test, colors_test)
     predictors = __pointpca2_internal(
-        points_a, colors_a, points_b, colors_b, search_size, verbose
+        points_a,
+        colors_a,
+        points_b,
+        colors_b,
+        search_size,
+        max_workers,
+        verbose,
     )
     return predictors

{pointpca2-0.4.2 → pointpca2-0.4.4}/src/lib.rs

@@ -1,20 +1,16 @@
 extern crate nalgebra as na;
 extern crate numpy;
 extern crate pointpca2_rs;
+extern crate rayon;
+
+mod utils;
 
-use na::DMatrix;
 use numpy::PyArray1;
 use numpy::PyReadonlyArray2;
 use pyo3::prelude::*;
 use pyo3::wrap_pyfunction;
-
-fn as_dmatrix<'a, T>(x: &'a PyReadonlyArray2<T>) -> DMatrix<T>
-where
-    T: numpy::Element + na::Scalar,
-{
-    let data: Vec<T> = x.as_array().iter().cloned().collect();
-    DMatrix::from_row_slice(x.shape()[0], x.shape()[1], &data)
-}
+use utils::as_dmatrix;
+use utils::build_thread_pool;
 
 #[pyfunction]
 fn compute_pointpca2<'py>(
@@ -24,24 +20,24 @@ fn compute_pointpca2<'py>(
     points_b: PyReadonlyArray2<'py, f64>,
     colors_b: PyReadonlyArray2<'py, u8>,
     search_size: usize,
+    max_workers: usize,
     verbose: bool,
 ) -> &'py PyArray1<f64> {
-    let (points_a, colors_a, points_b, colors_b) = {
-        (
-            as_dmatrix(&points_a),
-            as_dmatrix(&colors_a),
-            as_dmatrix(&points_b),
-            as_dmatrix(&colors_b),
+    let points_a = as_dmatrix(points_a);
+    let colors_a = as_dmatrix(colors_a);
+    let points_b = as_dmatrix(points_b);
+    let colors_b = as_dmatrix(colors_b);
+    let pool = build_thread_pool(max_workers);
+    let pooled_predictors = pool.install(|| {
+        pointpca2_rs::compute_pointpca2(
+            points_a,
+            colors_a,
+            points_b,
+            colors_b,
+            search_size,
+            verbose,
         )
-    };
-    let pooled_predictors = pointpca2_rs::compute_pointpca2(
-        points_a,
-        colors_a,
-        points_b,
-        colors_b,
-        search_size,
-        verbose,
-    );
+    });
     let py_array = PyArray1::from_iter(_py, pooled_predictors.iter().cloned());
     py_array
 }

pointpca2-0.4.4/src/utils.rs

@@ -0,0 +1,20 @@
+use na::DMatrix;
+use numpy::PyReadonlyArray2;
+use rayon::ThreadPool;
+use rayon::ThreadPoolBuilder;
+
+pub fn as_dmatrix<T>(x: PyReadonlyArray2<T>) -> DMatrix<T>
+where
+    T: numpy::Element + na::Scalar,
+{
+    let data: Vec<T> = x.as_array().iter().cloned().collect();
+    DMatrix::from_row_slice(x.shape()[0], x.shape()[1], &data)
+}
+
+pub fn build_thread_pool(max_workers: usize) -> ThreadPool {
+    let pool = ThreadPoolBuilder::new()
+        .num_threads(max_workers)
+        .build()
+        .expect("Failed to build thread pool");
+    pool
+}