BENCH Add ASV benchmark setup for radius_neighborhood

jjerphan · jjerphan · commit 051a6b81691f · 2022-02-23T10:31:40.000+01:00
diff --git a/asv_benchmarks/asv_scalability.sh b/asv_benchmarks/asv_scalability.sh
@@ -0,0 +1,15 @@
+#!/usr/bin/env bash
+#
+# Launch a ASV benchmark on groups of CPU
+# for scalability inspection
+#
+# Results are saved in given folders and files.
+#
+for i in 20 16 8 4 2 1;
+do
+  last_core=$(($i-1))
+    taskset -c 0-$last_core \
+    asv continuous -b PairwiseDistancesRadiusNeighborhood \
+    -e main pairwise-distances-radius-neighborhood | tee pairwise_distances_radius_neighborhood_asv_${i}_cores.txt
+    cp -R results results_${i}_cores
+done
diff --git a/asv_benchmarks/benchmarks/pdr.py b/asv_benchmarks/benchmarks/pdr.py
@@ -0,0 +1,26 @@
+import numpy as np
+
+from .common import Benchmark
+
+from sklearn.neighbors import NearestNeighbors
+
+
+class PairwiseDistancesRadiusNeighborhoodBenchmark(Benchmark):
+
+    param_names = ["n_train", "n_test", "n_features"]
+    params = [
+        [1000, 10_000],
+        [1000, 10_000, 100_000],
+        [100],
+    ]
+
+    def setup(self, n_train, n_test, n_features):
+        rng = np.random.RandomState(0)
+        X_train = rng.rand(n_train, n_features)
+        self.X_test = rng.rand(n_test, n_features)
+        self.est = NearestNeighbors().fit(X_train)
+
+    def time_pairwise_distances_argmin(self, n_train, n_test, n_features):
+        self.est.radius_neighbors(
+            X=self.X_test, return_distance=True, sort_results=True
+        )
