wip: fixing ann benchmarks bench

nishaq503 · Nov 29, 2024 · c459abf · c459abf
1 parent 8621eb8
commit c459abf
Show file tree

Hide file tree

Showing 8 changed files with 441 additions and 49 deletions.
diff --git a/crates/abd-clam/benches/ann_benchmarks.rs b/crates/abd-clam/benches/ann_benchmarks.rs
@@ -7,7 +7,7 @@ use std::{
 
 use abd_clam::{
     cakes::PermutedBall,
-    cluster::{adapter::ParBallAdapter, ParPartition},
+    cluster::{adapter::ParBallAdapter, BalancedBall, ParPartition},
     metric::{Cosine, Euclidean, ParMetric},
     Ball, Cluster, Dataset, FlatVec,
 };
@@ -116,19 +116,27 @@ fn run_search<M: ParMetric<Vec<f32>, f32>>(
     let ball = Ball::par_new_tree(data, metric, &criteria, seed);
     let (perm_ball, perm_data) = PermutedBall::par_from_ball_tree(ball.clone(), data.clone(), metric);
 
+    let criteria = |c: &BalancedBall<_>| c.cardinality() > 1;
+    let balanced_ball = BalancedBall::par_new_tree(data, metric, &criteria, seed).into_ball();
+    let (perm_balanced_ball, perm_balanced_data) =
+        PermutedBall::par_from_ball_tree(balanced_ball.clone(), data.clone(), metric);
+
     let root_radius = ball.radius();
     let radii = radii_fractions.iter().map(|&r| r * root_radius).collect::<Vec<_>>();
 
     utils::compare_permuted(
         c,
         metric,
         (&ball, data),
+        (&balanced_ball, &perm_balanced_data),
         (&perm_ball, &perm_data),
+        (&perm_balanced_ball, &perm_balanced_data),
+        None,
         None,
         queries,
         &radii,
         ks,
-        true,
+        false,
     );
 }
 

diff --git a/crates/abd-clam/benches/genomic_search.rs b/crates/abd-clam/benches/genomic_search.rs
@@ -2,7 +2,7 @@
 
 use abd_clam::{
     cakes::PermutedBall,
-    cluster::{adapter::BallAdapter, Partition},
+    cluster::{adapter::ParBallAdapter, BalancedBall, ParPartition},
     metric::Levenshtein,
     pancakes::SquishyBall,
     Ball, Cluster, FlatVec,
@@ -56,16 +56,26 @@ fn genomic_search(c: &mut Criterion) {
     let metric = Levenshtein;
 
     let criteria = |c: &Ball<u16>| c.cardinality() > 1;
-    let ball = Ball::new_tree(&data, &metric, &criteria, seed);
-    let (perm_ball, perm_data) = PermutedBall::from_ball_tree(ball.clone(), data.clone(), &metric);
-    let (squishy_ball, dec_data) = SquishyBall::from_ball_tree(ball.clone(), data.clone(), &metric);
+    let ball = Ball::par_new_tree(&data, &metric, &criteria, seed);
+    let (perm_ball, perm_data) = PermutedBall::par_from_ball_tree(ball.clone(), data.clone(), &metric);
+    let (squishy_ball, dec_data) = SquishyBall::par_from_ball_tree(ball.clone(), data.clone(), &metric);
+
+    let criteria = |c: &BalancedBall<u16>| c.cardinality() > 1;
+    let balanced_ball = BalancedBall::par_new_tree(&data, &metric, &criteria, seed).into_ball();
+    let (perm_balanced_ball, perm_balanced_data) =
+        PermutedBall::par_from_ball_tree(balanced_ball.clone(), data.clone(), &metric);
+    let (squishy_balanced_ball, dec_balanced_data) =
+        SquishyBall::par_from_ball_tree(balanced_ball.clone(), data.clone(), &metric);
 
     utils::compare_permuted(
         c,
         &metric,
         (&ball, &data),
+        (&balanced_ball, &perm_balanced_data),
         (&perm_ball, &perm_data),
+        (&perm_balanced_ball, &perm_balanced_data),
         Some((&squishy_ball, &dec_data)),
+        Some((&squishy_balanced_ball, &dec_balanced_data)),
         &queries,
         &radii,
         &ks,

diff --git a/crates/abd-clam/benches/utils/compare_permuted.rs b/crates/abd-clam/benches/utils/compare_permuted.rs
@@ -10,15 +10,21 @@ use abd_clam::{
 use criterion::*;
 use distances::Number;
 
-/// Compare the performance of different algorithms on permuted datasets.
+/// Compare the performance of different algorithms on datasets using different
+/// cluster types.
 ///
 /// # Parameters
 ///
 /// - `c`: The criterion context.
 /// - `metric`: The metric used to measure distances.
 /// - `ball_data`: The original dataset and its ball.
+/// - `balanced_ball_data`: The original dataset and its balanced ball.
 /// - `perm_ball_data`: The permuted dataset and its permuted ball.
+/// - `perm_balanced_ball_data`: The permuted dataset and its permuted balanced
+///   ball.
 /// - `dec_ball_data`: The permuted dataset and its squishy ball, if any.
+/// - `dec_balanced_ball_data`: The permuted dataset and its balanced squishy
+///   ball, if any.
 /// - `queries`: The queries to search for.
 /// - `radii`: The radii to use for the RNN algorithm.
 /// - `ks`: The numbers of neighbors to search for.
@@ -35,8 +41,11 @@ pub fn compare_permuted<I, U, Co, M>(
     c: &mut Criterion,
     metric: &M,
     ball_data: (&Ball<U>, &Co),
+    balanced_ball_data: (&Ball<U>, &Co),
     perm_ball_data: (&PermutedBall<U, Ball<U>>, &Co),
+    perm_balanced_ball_data: (&PermutedBall<U, Ball<U>>, &Co),
     dec_ball_data: Option<(&SquishyBall<U, Ball<U>>, &CodecData<I, usize>)>,
+    dec_balanced_ball_data: Option<(&SquishyBall<U, Ball<U>>, &CodecData<I, usize>)>,
     queries: &[I],
     radii: &[U],
     ks: &[usize],
@@ -55,6 +64,8 @@ pub fn compare_permuted<I, U, Co, M>(
 
     let (ball, data) = ball_data;
     let (off_ball, perm_data) = perm_ball_data;
+    let (balanced_ball, balanced_data) = balanced_ball_data;
+    let (perm_off_ball, perm_balanced_data) = perm_balanced_ball_data;
 
     let mut group = c.benchmark_group(format!("{}-{}-RnnClustered", data.name(), metric.name()));
     group
@@ -70,30 +81,58 @@ pub fn compare_permuted<I, U, Co, M>(
                 b.iter_with_large_drop(|| ball.batch_search(data, metric, queries, alg));
             });
 
+            group.bench_with_input(BenchmarkId::new("BalancedBall", i), &radius, |b, _| {
+                b.iter_with_large_drop(|| balanced_ball.batch_search(balanced_data, metric, queries, alg));
+            });
+
             group.bench_with_input(BenchmarkId::new("OffBall", i), &radius, |b, _| {
                 b.iter_with_large_drop(|| off_ball.batch_search(perm_data, metric, queries, alg));
             });
 
+            group.bench_with_input(BenchmarkId::new("OffBalancedBall", i), &radius, |b, _| {
+                b.iter_with_large_drop(|| perm_off_ball.batch_search(perm_balanced_data, metric, queries, alg));
+            });
+
             if let Some((dec_root, dec_data)) = dec_ball_data {
                 group.bench_with_input(BenchmarkId::new("SquishyBall", i), &radius, |b, _| {
                     b.iter_with_large_drop(|| dec_root.batch_search(dec_data, metric, queries, alg));
                 });
             }
+
+            if let Some((dec_root, dec_data)) = dec_balanced_ball_data {
+                group.bench_with_input(BenchmarkId::new("BalancedSquishyBall", i), &radius, |b, _| {
+                    b.iter_with_large_drop(|| dec_root.batch_search(dec_data, metric, queries, alg));
+                });
+            }
         }
 
         group.bench_with_input(BenchmarkId::new("ParBall", i), &radius, |b, _| {
             b.iter_with_large_drop(|| ball.par_batch_search(data, metric, queries, alg));
         });
 
+        group.bench_with_input(BenchmarkId::new("ParBalancedBall", i), &radius, |b, _| {
+            b.iter_with_large_drop(|| balanced_ball.par_batch_search(balanced_data, metric, queries, alg));
+        });
+
         group.bench_with_input(BenchmarkId::new("ParOffBall", i), &radius, |b, _| {
             b.iter_with_large_drop(|| off_ball.par_batch_search(perm_data, metric, queries, alg));
         });
 
+        group.bench_with_input(BenchmarkId::new("ParOffBalancedBall", i), &radius, |b, _| {
+            b.iter_with_large_drop(|| perm_off_ball.par_batch_search(perm_balanced_data, metric, queries, alg));
+        });
+
         if let Some((dec_root, dec_data)) = dec_ball_data {
             group.bench_with_input(BenchmarkId::new("ParSquishyBall", i), &radius, |b, _| {
                 b.iter_with_large_drop(|| dec_root.par_batch_search(dec_data, metric, queries, alg));
             });
         }
+
+        if let Some((dec_root, dec_data)) = dec_balanced_ball_data {
+            group.bench_with_input(BenchmarkId::new("ParBalancedSquishyBall", i), &radius, |b, _| {
+                b.iter_with_large_drop(|| dec_root.par_batch_search(dec_data, metric, queries, alg));
+            });
+        }
     }
     group.finish();
 
@@ -112,30 +151,58 @@ pub fn compare_permuted<I, U, Co, M>(
                     b.iter_with_large_drop(|| ball.batch_search(data, metric, queries, alg));
                 });
 
+                group.bench_with_input(BenchmarkId::new("BalancedBall", k), &k, |b, _| {
+                    b.iter_with_large_drop(|| balanced_ball.batch_search(balanced_data, metric, queries, alg));
+                });
+
                 group.bench_with_input(BenchmarkId::new("OffBall", k), &k, |b, _| {
                     b.iter_with_large_drop(|| off_ball.batch_search(perm_data, metric, queries, alg));
                 });
 
+                group.bench_with_input(BenchmarkId::new("OffBalancedBall", k), &k, |b, _| {
+                    b.iter_with_large_drop(|| perm_off_ball.batch_search(perm_balanced_data, metric, queries, alg));
+                });
+
                 if let Some((dec_root, dec_data)) = dec_ball_data {
                     group.bench_with_input(BenchmarkId::new("SquishyBall", k), &k, |b, _| {
                         b.iter_with_large_drop(|| dec_root.batch_search(dec_data, metric, queries, alg));
                     });
                 }
+
+                if let Some((dec_root, dec_data)) = dec_balanced_ball_data {
+                    group.bench_with_input(BenchmarkId::new("BalancedSquishyBall", k), &k, |b, _| {
+                        b.iter_with_large_drop(|| dec_root.batch_search(dec_data, metric, queries, alg));
+                    });
+                }
             }
 
             group.bench_with_input(BenchmarkId::new("ParBall", k), &k, |b, _| {
                 b.iter_with_large_drop(|| ball.par_batch_search(data, metric, queries, alg));
             });
 
+            group.bench_with_input(BenchmarkId::new("ParBalancedBall", k), &k, |b, _| {
+                b.iter_with_large_drop(|| balanced_ball.par_batch_search(balanced_data, metric, queries, alg));
+            });
+
             group.bench_with_input(BenchmarkId::new("ParOffBall", k), &k, |b, _| {
                 b.iter_with_large_drop(|| off_ball.par_batch_search(perm_data, metric, queries, alg));
             });
 
+            group.bench_with_input(BenchmarkId::new("ParOffBalancedBall", k), &k, |b, _| {
+                b.iter_with_large_drop(|| perm_off_ball.par_batch_search(perm_balanced_data, metric, queries, alg));
+            });
+
             if let Some((dec_root, dec_data)) = dec_ball_data {
                 group.bench_with_input(BenchmarkId::new("ParSquishyBall", k), &k, |b, _| {
                     b.iter_with_large_drop(|| dec_root.par_batch_search(dec_data, metric, queries, alg));
                 });
             }
+
+            if let Some((dec_root, dec_data)) = dec_balanced_ball_data {
+                group.bench_with_input(BenchmarkId::new("ParBalancedSquishyBall", k), &k, |b, _| {
+                    b.iter_with_large_drop(|| dec_root.par_batch_search(dec_data, metric, queries, alg));
+                });
+            }
         }
         group.finish();
     }

diff --git a/crates/abd-clam/benches/vector_search.rs b/crates/abd-clam/benches/vector_search.rs
@@ -2,7 +2,7 @@
 
 use abd_clam::{
     cakes::PermutedBall,
-    cluster::{adapter::ParBallAdapter, ParPartition},
+    cluster::{adapter::ParBallAdapter, BalancedBall, ParPartition},
     metric::{Euclidean, ParMetric},
     Ball, Cluster, Dataset, FlatVec, Metric,
 };
@@ -60,11 +60,19 @@ fn run_search<M: ParMetric<Vec<f32>, f32>>(
     let ball = Ball::par_new_tree(data, metric, &criteria, seed);
     let (perm_ball, perm_data) = PermutedBall::par_from_ball_tree(ball.clone(), data.clone(), metric);
 
+    let criteria = |c: &BalancedBall<_>| c.cardinality() > 1;
+    let balanced_ball = BalancedBall::par_new_tree(data, metric, &criteria, seed).into_ball();
+    let (perm_balanced_ball, perm_balanced_data) =
+        PermutedBall::par_from_ball_tree(balanced_ball.clone(), data.clone(), metric);
+
     utils::compare_permuted(
         c,
         metric,
         (&ball, data),
+        (&balanced_ball, &perm_balanced_data),
         (&perm_ball, &perm_data),
+        (&perm_balanced_ball, &perm_balanced_data),
+        None,
         None,
         queries,
         radii,