automatic array conversion and check for invalid shape

Author: Felix Igelbrink

numba_kdtree/kd_tree.py

@@ -382,6 +382,9 @@ def _query_radius_impl(self, X, r, p=2.0, eps=0.0, return_sorted=False, return_l
             else:
                 r_ = _convert_to_valid_input/(r, 1, dtype_npy).squeeze()
 
+            if r_.shape != (n_queries,):
+                raise ValueError("Invalid shape for r. Must be broadcastable to the number of queries.")
+
             if p < 1:
                 raise ValueError("Only p-norms with 1<=p<=infinity permitted")
 

tests/test_kdtree.py

@@ -179,6 +179,10 @@ def test_kdtree_query_radius_array(data, kdtree, scipy_kdtree):
         assert len(ii[i]) == len(ii_scipy[i])
         assert np.all(ii[i] == ii_scipy[i]), "Not equal for i={}".format(i)
 
+    # invalid array shapes should not work
+    with pytest.raises(ValueError):
+        ii_invalid = kdtree.query_radius(data[:100], r=radii[:50], return_sorted=True)
+
     num_executions = 5
     r_benchmark = np.linspace(0.01, 0.05, 500)
     runtime_kdtree_query = timeit(lambda: kdtree.query_radius(data[:500], r=r_benchmark, return_sorted=True),