Repository: scipy/scipy

Issue:
cKDTree.query_ball_point() is slower than python loop for large return lists
When `cKDTree.query_ball_point()` is used on a array of points, `x`, and the number of points returned is large (>100) it is faster to just loop over `x` in python rather than calling `query_ball_point(x)`.



### Reproducing code example:
```python
import numpy as np
import scipy.spatial as ss
x = np.random.randn(10000, 1)
xtree = ss.cKDTree(x)

%timeit [xtree.query_ball_point(point, .02, p=np.inf) for point in x]
# 104 ms ± 1.13 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit xtree.query_ball_point(x, .02, p=np.inf)
# 206 ms ± 1.98 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
```
but for smaller number of returned points (~10)
```python
%timeit [xtree.query_ball_point(point, .002, p=np.inf) for point in x]
# 57.2 ms ± 1.34 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit xtree.query_ball_point(x, .002, p=np.inf)
# 45.2 ms ± 922 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

```
and you can confirm
```python
r1 = [xtree.query_ball_point(point, .002, p=np.inf) for point in x]
r2 = xtree.query_ball_point(x, .002, p=np.inf)
all([np.allclose(sorted(rx1), sorted(rx2)) for rx1, rx2 in zip(r1, r2)])
# True
```

### Scipy/Numpy/Python version information:
```python
import sys, scipy, numpy; print(scipy.__version__, numpy.__version__, sys.version_info
# 1.1.0 1.14.2 sys.version_info(major=3, minor=5, micro=5, releaselevel='final', serial=0)
```




Discussion:
Also saw this slowdown in #8818.
I'm guessing this difference in speed has to do with allocating the lists. I really only need the number of points returned (and not their index). I wrote a ```count_ball_points()``` method which just tracks the number of points and allocates and returns an array of ints. It is significantly faster than both of these.

```python
%timeit xtree.count_ball_point(x, .02, p=np.inf)
# 19.9 ms ± 56 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
```
I am quite sure the culprits are lines 958-960 and 1001-1012. We have to rewrite those to make Cython produce efficient C++.
Another cause of slowdown is line 921. This one is much more expensive in the case where ```x``` has a complicated shape. It is hard to get around though because of the shapes we allow ```x``` to have. It is not known in advance.
The easiest solution would be to provide an optimized path for the common case where ```x``` on input is an ndarray of shape ```(n,m)``` with dtype ```float64``` and the last dimension is contiguous.  
@sturlamolden I think it's even simpler. If the input is a single point, the results are [not sorted](https://github.com/scipy/scipy/blob/d66ee2807e1ec8b42d05233fb1eb6ef35e57c315/scipy/spatial/ckdtree.pyx#L937), whereas they [are sorted if more than one point is queried](https://github.com/scipy/scipy/blob/d66ee2807e1ec8b42d05233fb1eb6ef35e57c315/scipy/spatial/ckdtree.pyx#L1009).

I checked this and also tried special casing the ```x.ndim==2``` case (possibly poorly).

tl;dr Removing the ```sorted()``` speeds up a small example by 2-3x. Special casing the x.ndim==2 case also speeds things up by 20%.

```python
import numpy as np
import scipy.spatial as ss
x = np.random.randn(10000, 1)
xtree = ss.cKDTree(x)
```

All sorted, i.e., I added sorting to the ```x.ndim==1``` case:
```python
%timeit [xtree.query_ball_point(point, .02, p=np.inf) for point in x]
# 183 ms ± 2.01 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit xtree.query_ball_point(x, .02, p=np.inf)
# 151 ms ± 1.49 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

xp = np.ascontiguousarray(x.reshape(100,100, 1))
%timeit xtree.query_ball_point(xp, .02, p=np.inf)
# 163 ms ± 1.7 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

None sorted, i.e., I removed sorting from ```x.ndim > 1``` cases:
```python
%timeit [xtree.query_ball_point(point, .02, p=np.inf) for point in x]
# 78.1 ms ± 1.01 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit xtree.query_ball_point(x, .02, p=np.inf)
# 52.1 ms ± 580 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

xp = np.ascontiguousarray(x.reshape(100,100, 1))
%timeit xtree.query_ball_point(xp, .02, p=np.inf)
# 65.1 ms ± 1.14 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

You can see the changes here
https://github.com/JesseLivezey/scipy/blob/5c70d8ea373adf25c91afe5e529ef3bd16d5d5fa/scipy/spatial/ckdtree.pyx#L919
I think it would make sense to fix this, but I'm not sure the best way to do it that doesn't break backwards compatibility.

Although the documentation and behavior never guarantee a sorted set of indices, it's possible that someone could be relying on the sorted behavior in the multi-point case.
We can optionally turn it off with a keyword argument. Default to the current behavior.


Resolving patch:
diff --git a/benchmarks/benchmarks/spatial.py b/benchmarks/benchmarks/spatial.py
--- a/benchmarks/benchmarks/spatial.py
+++ b/benchmarks/benchmarks/spatial.py
@@ -127,6 +127,10 @@ def setup_query_pairs(self, mnr, p, probe_radius, boxsize, leafsize):
     def time_query_ball_point(self, mnr, p, probe_radius, boxsize, leafsize):
         self.T.query_ball_point(self.queries, probe_radius, p=p)
 
+    def time_query_ball_point_nosort(self, mnr, p, probe_radius, boxsize, leafsize):
+        self.T.query_ball_point(self.queries, probe_radius, p=p,
+                                return_sorted=False)
+
     def time_query_pairs(self, mnr, p, probe_radius, boxsize, leafsize):
         self.T.query_pairs(probe_radius, p=p)
 


Test patch:
diff --git a/scipy/spatial/tests/test_kdtree.py b/scipy/spatial/tests/test_kdtree.py
--- a/scipy/spatial/tests/test_kdtree.py
+++ b/scipy/spatial/tests/test_kdtree.py
@@ -1330,3 +1330,28 @@ def test_short_knn():
             [0., 0.01, np.inf, np.inf],
             [0., 0.01, np.inf, np.inf],
             [0., np.inf, np.inf, np.inf]])
+
+class Test_sorted_query_ball_point(object):
+
+    def setup_method(self):
+        np.random.seed(1234)
+        self.x = np.random.randn(100, 1)
+        self.ckdt = cKDTree(self.x)
+
+    def test_return_sorted_True(self):
+        idxs_list = self.ckdt.query_ball_point(self.x, 1., return_sorted=True)
+        for idxs in idxs_list:
+            assert_array_equal(idxs, sorted(idxs))
+
+    def test_return_sorted_None(self):
+        """Previous behavior was to sort the returned indices if there were
+        multiple points per query but not sort them if there was a single point
+        per query."""
+        idxs_list = self.ckdt.query_ball_point(self.x, 1.)
+        for idxs in idxs_list:
+            assert_array_equal(idxs, sorted(idxs))
+
+        idxs_list_single = [self.ckdt.query_ball_point(xi, 1.) for xi in self.x]
+        idxs_list_False = self.ckdt.query_ball_point(self.x, 1., return_sorted=False)
+        for idxs0, idxs1 in zip(idxs_list_False, idxs_list_single):
+            assert_array_equal(idxs0, idxs1)