ENH: add STRtree nearest neighbor(s) function

[brendan-ward]

Replaces #111, resolves #110.

Note: this builds on top of #265, which should be reviewed and merged first.

Based on the feedback in #111, this provides a nearest neighbor function on the tree that overcomes some of the limitations of the GEOS API. Specifically, the GEOS API returns a single, nondeterministic geometry from a set of intersected or equidistant tree geometries, which causes problems with downstream usage.

This provides an index-oriented results (input indexes, target indexes) similar to query_bulk, with an optional distance array.

This implementation should return all intersected or equidistant tree geometries for each input geometry. To do this, it leverages the userdata parameter of the tree's distance callback function to pass in a data structure that allows us to accumulate pairs of nearby geometries, track the minimum distance observed, and slightly modify (read: lie about) the distance returned to the GEOS function to force investigation of tree geometries in adjacent tree nodes. Without this modification, the algorithm stops before visiting additional equidistant or intersected geometries in adjacent tree nodes, producing an incomplete result set.

This implementation also enables pre-screening of input geometries by a max_distance parameter. This is used to create an envelope around the geometry to query against the tree which is used to make sure that there are tree geometries within that distance. Depending on the configuration of the geometry and the distance used, this either incurs extra overhead or it can dramatically reduce the time needed to find nearest items within that distance (vs finding all nearest then filtering later).

This involved adding helper functions to calculate bounds and create a box; these were derived from similar functions in the ufuncs (and could be consolidated later).

Note: this is not intended to be used for kNN functionality; that will require a separate algorithm implemented within GEOS first. The only time this will return multiple neighbors for an input geometry is when there are multiple equidistant or intersected tree geometries.

The doctests for nearest are marked to skip them completely; I could not find a way to conditionally skip them for GEOS < 3.6.

/cc @adriangb

[adriangb]

This looks awesome @brendan-ward! Very excited to implement a nearest sjoin in geopandas based on this.

I'm guessing this is way faster, but have you benchmarked against #111 (comment) or similar?

[brendan-ward]

I just added some benchmarks against the fastest Python implementation I could figure out. I'm sure the Python version could be improved; they aren't directly comparable because the Python version uses a search radius and may not find all nearest neighbors - but also comes with a steep penalty of comparing many geometries from the tree.

I did not compare against an alternative approach in C; that would be a bit of coding just to demonstrate that it is slower.

time_tree_nearest_points                                                        78.5±1ms
time_tree_nearest_poly                                                          217±10ms
time_tree_nearest_poly_python                                                 1.47±0.05s

[adriangb]

I was more thinking more along the lines of using #111 (comment) within the CPython code instead of using the distance callback, userdata data structure, etc. It's basically the same as your Python approach except that tolerance is determined by the distance to the single result returned by GEOSSTRtree_nearest_generic_r. I'll try to code that up in Python just for a rough comparison.

Edit: opened brendan-ward#1 to clarify.

[brendan-ward]

GEOSDistanceIndexed was much slower according to our benchmarks. There may be some situations where it works well, but leveraging prepared geometries and GEOSPreparedDistance in GEOS >= 3.9 is likely to have a bigger benefit.

[adriangb]

Here's a really rough python implementation:

class STRtree:
    def time_tree_nearest_poly_python(self):
        # simulate getting "a" single nearest geometry (similar to what `GEOSSTRtree_nearest_generic_r` gives you)
        left, right = self.tree.nearest(self.points)
        singles = np.unique(left, return_index=True)[1]  # simulate getting "a" _single_ nearest point for each geometry
        left, right = left[singles], right[singles]
        # get the distance to each nearest geometry
        tolerance = pygeos.distance(self.polygons.take(right), self.points.take(left)) * 1.01  # this is some arbitrary relative tolerance
        # buffer each input geometry and find all other geometries in that buffer
        b = pygeos.buffer(self.points, tolerance)
        left, right = self.tree.query_bulk(b, "intersects")
        # find the distance to each of those possibly equidistant geometries
        dist = pygeos.distance(self.points.take(left), self.polygons.take(right))
        
       # a bunch of indexing stuff that should be trivial (computationally)
        ...

With this, I'm getting the following timings (manually timed, not asv):

0.35 sec  # this python version
0.22  # time_tree_nearest_poly

[brendan-ward]

One issue I'm finding in trying to adapt some existing workflows to use this is the performance impact of not having a tolerance or max distance parameter to short-circuit the results: the tree nearest function operates against all input geometries, even if some are very far away from their nearest neighbors

We might be able to introduce a limit here by either:

• pre-screening source geometries so that they are (approximately) no further than a max distance of tree geometries; this could be via a simple bounding box query against the tree (adding max distance to all sides of envelope around input geometry). Those that pass pre-screening then go into the GEOS nearest function as now.

• passing a max distance and tree leaf size into the distance callback, and tracking the number of leaf nodes visited by that callback. If it visits all leaf nodes within a given tree node and none of the distances are less than tolerance, it is unlikely that subsequent tree nodes will yield closer results - so we can bail early. (only mechanism to do so is to pass an error back to GEOS; not a great alternative).

[brendan-ward]

Updated benchmarks:

time_tree_nearest_points                                       80.7±9ms
time_tree_nearest_points_small_max_distance                    18.8±1ms
time_tree_nearest_points_large_max_distance                     152±8ms
time_tree_nearest_poly                                         218±10ms
time_tree_nearest_poly_small_max_distance                      156±20ms
time_tree_nearest_poly_python                                 1.61±0.1s

time_tree_nearest_points_large_max_distance shows the overhead of using a max distance that covers most of the tree geometries, most likely because all are visited by the tree query callback for each input geometry; I couldn't find a way to fail-fast out of the query callback, which would dramatically reduce this overhead.

In terms of real-world data (115k points vs 13k polygons), I found a 12x speedup (0.27s vs 3.4s) by using this function with the same maximum distance that I was using in a python + ufuncs approach that buffered input points, queried tree for intersection with the buffers, then took the smallest distance for each. Both were much faster than the tree nearest function with no limits (45s).

[brendan-ward]

This now has 2 functions: nearest which returns the nearest singular (nondeterministic) result for each input, and nearest_all which returns all equidistant results.

nearest is a stripped-down, simplified version; it does not currently support the same parameters as nearest_all and does not return distances. The goal was to make it as simple as possible.

Updated benchmarks:

time_tree_nearest_points                                                                            21.7±2ms
time_tree_nearest_all_points                                                                        23.6±2ms
time_tree_nearest_all_points_large_max_distance                                                     62.6±3ms
time_tree_nearest_all_points_small_max_distance                                                   6.96±0.2ms
time_tree_nearest_poly                                                                              82.9±9ms
time_tree_nearest_all_poly                                                                          80.2±4ms
time_tree_nearest_all_poly_small_max_distance                                                       44.3±2ms
time_tree_nearest_all_poly_python                                                                  1.65±0.2s

The nearest and nearest_all don't diverge much in the benchmarks.

[caspervdw]

Thanks for addressing my suggestions @brendan-ward . I have some more suggestions, especially about the API of nearest. The other suggestions are more optimizations on the simpler algorithm.

The benchmarks you showed are remarkably close indeed! Maybe the suggested optimizations make the difference a bit larger...

[caspervdw]

I believe that these changes should be in the geos.h file?

[jorisvandenbossche]

My apologies as well for not giving feedback on this topic earlier (and also on the previous issue/PR). But thanks a lot for working on this!

I think the explanation of the current limitations of the GEOS API makes sense (based on my limited experience on this topic). So I think it is good to address this, and I am also fine with having such workarounds in PyGEOS in the short term. However, I would personally be hesitant to commit to supporting this long term in PyGEOS/Shapely, so I would personally like to see at least some discussion with the GEOS team about whether they are interested in such changes / whether there are reasons for the current implementation/limitations.

[brendan-ward]

Updated benchmarks specifically to probe the case of multiple equidistant neighbors:

time_tree_nearest_points                                         37.1±2ms
time_tree_nearest_points_equidistant                             44.0±2ms
time_tree_nearest_all_points                                     31.6±2ms
time_tree_nearest_all_points_equidistant                       40.9±0.5ms

Will need to investigate a bit further as to why the nearest_all variants seem a touch faster than the nearest counterparts, which are intentionally simpler and should be theoretically faster.

[brendan-ward]

Updated benchmarks to include a variant of searching for additional equidistant neighbors ourselves, given singular results from nearest (per comments in the GEOS IRC).

time_tree_nearest_points_equidistant                           36.6±0.4ms
time_tree_nearest_all_points_equidistant                         42.2±2ms
time_tree_nearest_points_equidistant_manual_all                95.1±0.6ms

[caspervdw]

I approved this PR @brendan-ward feel free to merge

[brendan-ward]

Found one small optimization: we were passing in query GeometryObjects to the callback, which required pulling back out the GEOSGeometry on each test against a geometry in the tree. We're able to pass the query GEOSGeometry in directly as one of the inputs to the callback, which cuts down on a little overhead.

Otherwise, the profiles (via py-spy) between nearest and nearest_all look very similar, and their benchmarks are within a reasonable tolerance of each other, so this is probably as good as its going to get for now.

[adriangb]

this addition by you and @adriangb

To be clear, this was all @brendan-ward , awesome work!

I'm excited to integrate this into GeoPandas' sindex.