Add Hamerly's accelerated K-means algorithm to linfa-clustering

algorithms/linfa-clustering/benches/k_means.rs

@@ -5,7 +5,7 @@ use criterion::{
 use linfa::benchmarks::config;
 use linfa::prelude::*;
 use linfa::DatasetBase;
-use linfa_clustering::{IncrKMeansError, KMeans, KMeansInit};
+use linfa_clustering::{IncrKMeansError, KMeans, KMeansAlgorithm, KMeansInit};
 use linfa_datasets::generate;
 use ndarray::Array2;
 use ndarray_rand::RandomExt;
@@ -36,33 +36,41 @@ impl Drop for Stats {
 fn k_means_bench(c: &mut Criterion) {
     let mut rng = Xoshiro256Plus::seed_from_u64(40);
     let cluster_sizes = [(100, 4), (400, 10), (3000, 10)];
+    let algorithms = [KMeansAlgorithm::Lloyd, KMeansAlgorithm::Hamerly];
     let n_features = 3;
 
-    let mut benchmark = c.benchmark_group("naive_k_means");
+    let mut benchmark = c.benchmark_group("k_means");
     config::set_default_benchmark_configs(&mut benchmark);
     benchmark.plot_config(PlotConfiguration::default().summary_scale(AxisScale::Logarithmic));
 
-    for &(cluster_size, n_clusters) in &cluster_sizes {
-        let rng = &mut rng;
-        let centroids =
-            Array2::random_using((n_clusters, n_features), Uniform::new(-30., 30.), rng);
-        let dataset = DatasetBase::from(generate::blobs(cluster_size, &centroids, rng));
-        let mut stats = Stats::default();
+    for &algorithm in &algorithms {
+        for &(cluster_size, n_clusters) in &cluster_sizes {
+            let rng = &mut rng;
+            let centroids =
+                Array2::random_using((n_clusters, n_features), Uniform::new(-30., 30.), rng);
+            let dataset = DatasetBase::from(generate::blobs(cluster_size, &centroids, rng));
+            let mut stats = Stats::default();
 
-        benchmark.bench_function(
-            BenchmarkId::new("naive_k_means", format!("{n_clusters}x{cluster_size}")),
-            |bencher| {
-                bencher.iter(|| {
-                    let m = KMeans::params_with_rng(black_box(n_clusters), black_box(rng.clone()))
-                        .init_method(KMeansInit::KMeansPlusPlus)
-                        .max_n_iterations(black_box(1000))
-                        .tolerance(black_box(1e-3))
-                        .fit(&dataset)
-                        .unwrap();
-                    stats.add(m.inertia());
-                });
-            },
-        );
+            benchmark.bench_function(
+                BenchmarkId::new(
+                    "k_means",
+                    format!("{algorithm:?}:{n_clusters}x{cluster_size}"),
+                ),
+                |bencher| {
+                    bencher.iter(|| {
+                        let m =
+                            KMeans::params_with_rng(black_box(n_clusters), black_box(rng.clone()))
+                                .init_method(KMeansInit::KMeansPlusPlus)
+                                .algorithm(algorithm)
+                                .max_n_iterations(black_box(1000))
+                                .tolerance(black_box(1e-3))
+                                .fit(&dataset)
+                                .unwrap();
+                        stats.add(m.inertia());
+                    });
+                },
+            );
+        }
     }
 
     benchmark.finish();