reformat

jackd · May 11, 2020 · 82014e9 · 82014e9
1 parent 8473021
commit 82014e9
Show file tree

Hide file tree

Showing 24 changed files with 852 additions and 786 deletions.
diff --git a/.isort.cfg b/.isort.cfg
@@ -0,0 +1,6 @@
+[settings]
+multi_line_output=3
+include_trailing_comma=True
+force_grid_wrap=0
+use_parentheses=True
+line_length=88
diff --git a/benchmarks/binary_tree/find_node_split_dim.py b/benchmarks/binary_tree/find_node_split_dim.py
@@ -1,11 +1,7 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import numpy as np
-from numba_neighbors.benchmark_utils import run_benchmarks, benchmark
+
 from numba_neighbors import binary_tree as bt
-import sklearn.neighbors
+from numba_neighbors.benchmark_utils import benchmark, run_benchmarks
 
 N = 2048
 NS = 1024
@@ -29,7 +25,7 @@ def get_data():
 #         data, node_indices, n_features, n_points)
 
 
-@benchmark('numba')
+@benchmark("numba")
 def numba_impl():
     data = get_data()
     return bt.find_node_split_dim(*data)
@@ -40,7 +36,7 @@ def numpy_find(data, node_indices):
     return np.argmax(np.ptp(data))
 
 
-@benchmark('numpy')
+@benchmark("numpy")
 def numpy_impl():
     return numpy_find(*get_data())
 

diff --git a/benchmarks/binary_tree/simultaneous_sort.py b/benchmarks/binary_tree/simultaneous_sort.py
@@ -1,12 +1,9 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import numpy as np
-from numba_neighbors.benchmark_utils import run_benchmarks
-from numba_neighbors import binary_tree as bt
 import sklearn.neighbors
 
+from numba_neighbors import binary_tree as bt
+from numba_neighbors.benchmark_utils import run_benchmarks
+
 N = 2048
 k = 32
 
@@ -44,9 +41,5 @@ def numpy_impl():
 
 
 run_benchmarks(
-    10,
-    100,
-    ('sklearn', sklearn_impl),
-    ('numba', numba_impl),
-    ('numpy', numpy_impl),
+    10, 100, ("sklearn", sklearn_impl), ("numba", numba_impl), ("numpy", numpy_impl),
 )
diff --git a/benchmarks/index_heap.py b/benchmarks/index_heap.py
@@ -1,22 +1,19 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import heapq
+
+import numpy as np
+from numba import njit
+
+from numba_neighbors.benchmark_utils import BenchmarkManager
 from numba_neighbors.index_heap import IndexHeap
 from numba_neighbors.index_heap2 import IndexHeap as IndexHeap2
-from numba_neighbors.index_heap import padded_index_heap
-from numba_neighbors.benchmark_utils import BenchmarkManager
-from numba import njit
-import numpy as np
 
 length = 1024
 max_length = 32 * length
 
 heapify_bm = BenchmarkManager()
 
 
-@heapify_bm.benchmark('heapq')
+@heapify_bm.benchmark("heapq")
 def heapq_impl():
     np.random.seed(123)
     priorities = np.random.random(size=(length,)).astype(np.float32)
@@ -25,7 +22,7 @@ def heapq_impl():
     heapq.heapify(heap)
 
 
-@heapify_bm.benchmark('index_heap')
+@heapify_bm.benchmark("index_heap")
 @njit()
 def index_heap_impl():
     np.random.seed(123)
@@ -37,7 +34,7 @@ def index_heap_impl():
     iheap.heapify()
 
 
-@heapify_bm.benchmark('index_heap2')
+@heapify_bm.benchmark("index_heap2")
 @njit()
 def index_heap2_impl():
     np.random.seed(123)

diff --git a/benchmarks/kd_tree/build.py b/benchmarks/kd_tree/build.py
@@ -1,13 +1,9 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import numpy as np
-from numba_neighbors.benchmark_utils import run_benchmarks, benchmark
+import sklearn.neighbors
+
 from numba_neighbors import binary_tree as bt
 from numba_neighbors import kd_tree as kd
-# from numba_neighbors import kd_tree2 as kd2
-import sklearn.neighbors
+from numba_neighbors.benchmark_utils import benchmark, run_benchmarks
 
 N = 1024
 D = 3
@@ -17,30 +13,24 @@
 leaf_size = 16
 
 
-@benchmark('sklearn')
+@benchmark("sklearn")
 def sklearn_impl():
     return sklearn.neighbors.kd_tree.KDTree(data, leaf_size=leaf_size)
 
 
-@benchmark('base')
+@benchmark("base")
 def numba_impl():
     return bt.get_tree_data(data, leaf_size=leaf_size)
 
 
-@benchmark('BinaryTree')
+@benchmark("BinaryTree")
 def binary_tree():
-    # return kd.KDTree(data, leaf_size=leaf_size)
     return bt.BinaryTree(data, leaf_size=leaf_size)
 
 
-@benchmark('KDTree')
+@benchmark("KDTree")
 def kd_tree():
-    # return kd.KDTree(data, leaf_size=leaf_size)
     return kd.KDTree(data, leaf_size=leaf_size)
 
 
-# @benchmark('cls2')
-# def numba2():
-#     return kd2.KDTree(data, leaf_size=leaf_size)
-
 run_benchmarks(10, 100)
diff --git a/benchmarks/kd_tree/ifp_sample.py b/benchmarks/kd_tree/ifp_sample.py
@@ -1,20 +1,9 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
-
 # os.environ['NUMBA_DISABLE_JIT'] = '1'
 import numpy as np
-from numba import njit
-from numba_neighbors.benchmark_utils import run_benchmarks, benchmark
-from numba_neighbors import kd_tree as kd
+
 from dcbs.np_utils.ifp.sample_query import ifp_sample_and_query
-import functools
+from numba_neighbors import kd_tree as kd
+from numba_neighbors.benchmark_utils import benchmark, run_benchmarks
 
 N = 1024
 sample_size = 512
@@ -32,37 +21,48 @@
 @benchmark()
 def ifp():
     tree = kd.KDTree(data, leaf_size)
-    return tree.ifp_sample_query(query_r**2, tree.get_node_indices(),
-                                 sample_size, max_neighbors)
+    return tree.ifp_sample_query(
+        query_r ** 2, tree.get_node_indices(), sample_size, max_neighbors
+    )
 
 
 @benchmark()
 def rejection_ifp():
     tree = kd.KDTree(data, leaf_size)
-    return tree.rejection_ifp_sample_query(rejection_r**2, query_r**2,
-                                           tree.get_node_indices(), sample_size,
-                                           max_neighbors)
+    return tree.rejection_ifp_sample_query(
+        rejection_r ** 2,
+        query_r ** 2,
+        tree.get_node_indices(),
+        sample_size,
+        max_neighbors,
+    )
 
 
 @benchmark()
 def ifp3():
     tree = kd.KDTree3(data, leaf_size)
-    return tree.ifp_sample_query(query_r**2, tree.get_node_indices(),
-                                 sample_size, max_neighbors)
+    return tree.ifp_sample_query(
+        query_r ** 2, tree.get_node_indices(), sample_size, max_neighbors
+    )
 
 
 @benchmark()
 def rejection_ifp3():
     tree = kd.KDTree3(data, leaf_size)
-    return tree.rejection_ifp_sample_query(rejection_r**2, query_r**2,
-                                           tree.get_node_indices(), sample_size,
-                                           max_neighbors)
+    return tree.rejection_ifp_sample_query(
+        rejection_r ** 2,
+        query_r ** 2,
+        tree.get_node_indices(),
+        sample_size,
+        max_neighbors,
+    )
 
 
 @benchmark()
 def base():
-    return ifp_sample_and_query(data, query_r, sample_size, max_neighbors,
-                                   max_neighbors)
+    return ifp_sample_and_query(
+        data, query_r, sample_size, max_neighbors, max_neighbors
+    )
 
 
 run_benchmarks(20, 100)

diff --git a/benchmarks/kd_tree/query_radius.py b/benchmarks/kd_tree/query_radius.py
@@ -1,13 +1,9 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
 import numpy as np
-from numba_neighbors.benchmark_utils import run_benchmarks, benchmark
-from numba_neighbors import kd_tree as kd
-# from numba_neighbors import kd_tree2 as kd2
 import sklearn.neighbors
 
+from numba_neighbors import kd_tree as kd
+from numba_neighbors.benchmark_utils import benchmark, run_benchmarks
+
 N = 1024
 n = 256
 D = 3
@@ -24,15 +20,15 @@
 sk_tree = sklearn.neighbors.kd_tree.KDTree(data, leaf_size=leaf_size)
 
 
-@benchmark('sklearn')
+@benchmark("sklearn")
 def sklearn_impl():
     return sk_tree.query_radius(X, r, return_distance=True)
 
 
 numba_tree = kd.KDTree(data, leaf_size=leaf_size)
 
 
-@benchmark('numba_pre')
+@benchmark("numba_pre")
 def numba_prealloc():
     dists = np.full((n, max_neighbors), np.inf, dtype=np.float32)
     indices = np.zeros((n, max_neighbors), dtype=np.int64)
@@ -45,20 +41,21 @@ def numba():
     return numba_tree.query_radius(X, r2, max_neighbors)
 
 
-@benchmark('numba_bu')
+@benchmark("numba_bu")
 def numba_bottom_up():
     start_nodes = numba_tree.get_node_indices()[X_indices]
     return numba_tree.query_radius_bottom_up(X, r2, start_nodes, max_neighbors)
 
 
-@benchmark('numba_bu_pre')
+@benchmark("numba_bu_pre")
 def numba_bottom_up_prealloc():
     dists = np.full((n, max_neighbors), np.inf, dtype=np.float32)
     indices = np.zeros((n, max_neighbors), dtype=np.int64)
     counts = np.zeros((n,), dtype=np.int64)
     start_nodes = numba_tree.get_node_indices()[X_indices]
-    numba_tree.query_radius_bottom_up_prealloc(X, r2, start_nodes, dists,
-                                               indices, counts)
+    numba_tree.query_radius_bottom_up_prealloc(
+        X, r2, start_nodes, dists, indices, counts
+    )
 
 
 numba_tree3 = kd.KDTree3(data, leaf_size=leaf_size)
@@ -69,20 +66,21 @@ def numba3():
     return numba_tree3.query_radius(X, r2, max_neighbors)
 
 
-@benchmark('numba3_bu')
+@benchmark("numba3_bu")
 def numba3_bottom_up():
     start_nodes = numba_tree3.get_node_indices()[X_indices]
     return numba_tree3.query_radius_bottom_up(X, r2, start_nodes, max_neighbors)
 
 
-@benchmark('numba3_bu_pre')
+@benchmark("numba3_bu_pre")
 def numba3_bottom_up_prealloc():
     dists = np.full((n, max_neighbors), np.inf, dtype=np.float32)
     indices = np.zeros((n, max_neighbors), dtype=np.int64)
     counts = np.zeros((n,), dtype=np.int64)
     start_nodes = numba_tree3.get_node_indices()[X_indices]
-    numba_tree3.query_radius_bottom_up_prealloc(X, r2, start_nodes, dists,
-                                                indices, counts)
+    numba_tree3.query_radius_bottom_up_prealloc(
+        X, r2, start_nodes, dists, indices, counts
+    )
 
 
 # numba_tree2 = kd2.KDTree(data, leaf_size=leaf_size)

diff --git a/benchmarks/kd_tree/rejection_sample.py b/benchmarks/kd_tree/rejection_sample.py
@@ -1,21 +1,11 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
-
-import os
+import heapq
 
 # os.environ['NUMBA_DISABLE_JIT'] = '1'
 import numpy as np
-from numba import njit
-from numba_neighbors.benchmark_utils import run_benchmarks, benchmark
+
 from numba_neighbors import kd_tree as kd
-from dcbs.sparse.sample import ragged_in_place_and_down_sample_query_np
-import functools
-import heapq
+from numba_neighbors.benchmark_utils import benchmark, run_benchmarks
+
 heapq.heappush
 
 N = 1024
@@ -31,20 +21,28 @@
 data /= np.linalg.norm(data, axis=-1, keepdims=True)
 
 
-@benchmark('numba')
+@benchmark("numba")
 def numba_impl():
     tree = kd.KDTree(data, leaf_size)
-    return tree.rejection_sample_query(rejection_r**2, query_r**2,
-                                       tree.get_node_indices(), max_sample_size,
-                                       max_neighbors)
+    return tree.rejection_sample_query(
+        rejection_r ** 2,
+        query_r ** 2,
+        tree.get_node_indices(),
+        max_sample_size,
+        max_neighbors,
+    )
 
 
-@benchmark('numba3')
+@benchmark("numba3")
 def numba3_impl():
     tree = kd.KDTree3(data, leaf_size)
-    return tree.rejection_sample_query(rejection_r**2, query_r**2,
-                                       tree.get_node_indices(), max_sample_size,
-                                       max_neighbors)
+    return tree.rejection_sample_query(
+        rejection_r ** 2,
+        query_r ** 2,
+        tree.get_node_indices(),
+        max_sample_size,
+        max_neighbors,
+    )
 
 
 # @benchmark()