ENH: query tree with multiple input geometries

[brendan-ward]

This adds a first pass at bulk query against STRtree instance, as the core of a spatial join function for #52.

This extends the functionality of the existing STRtree::query() method to handle multiple input geometries. It returns a tuple of source (left) indexes and target (right) indexes that either overlap the windows of the sources, or also meet the predicate function provided.

Example:

tree = pygeos.STRtree(pygeos.points(np.arange(10), np.arange(10)))
tree.query_bulk([pygeos.points(1, 1)])
# returns: (array([0]), array([1]))

Multiple source geometries, first and third boxes overlap windows:

tree.query_bulk([
    pygeos.box(0,0, 1,1), 
    pygeos.box(100,100,110,110), 
    pygeos.box(2,2,4,4)
])
# returns: (array([0, 0, 2, 2, 2]), array([0, 1, 2, 3, 4]))

This preserves the original query() function as a singular operation.

/cc: @jorisvandenbossche

[brendan-ward]

Test failure appears unrelated to changes here; likely related to #100

[jorisvandenbossche]

Thanks a lot for looking into this! (you're always ahead of me :-))

And indeed, this can be as a method on the tree to do a bulk query, instead of a separate 'sjoin' function receiving a tree object.
Will try to do a closer review later today / this week!

[jorisvandenbossche]

Went through a bit in more detail, in general this is looking nice to me!
Also tried it out with the cities/countries dataset, and it is giving the same result as geopandas.sjoin (as expected it should be)

[brendan-ward]

@jorisvandenbossche Thanks for the thorough review!

What do you think about keeping the singular vs bulk operations split out? I've been pondering that for a while, and I'm a bit torn. On the one hand, I want to preserve the existing singular query operation which only returns a single array of the hits in the tree, but on the other hand it seems like most cases (e.g., spatial join) would use the query_bulk operation. Having 2 different functions feels inconsistent compared to other interfaces in pygeos.

What about coalescing these into one user-facing function query() which returns different shapes of results depending on what you pass in, and documenting that behavior thoroughly?

• If input geometry is a singular geometry, we return an array of indexes of geometries in the tree
• If input geometry is a 1D array of geometries, we return an array of indexes in the source geometries and corresponding indexes of geometries in the tree

[jorisvandenbossche]

Since the return type is really different (single array only pointing into tree geometries vs two arrays of indices), I personally think that is a good reason to keep them as separate functions (in the user facing API).
@caspervdw thoughts?

on the other hand it seems like most cases (e.g., spatial join) would use the query_bulk operation

I think there will also be quite some use cases for the scalar query (eg in geopandas, for clipping, or other operations where not combining two dataframes, but where the right object is a single geometry)

[brendan-ward]

@jorisvandenbossche I think I've addressed all your comments other than changes pending #100. I tried to clarify the inputs in various docstrings, but please let me know if anything should be more clear.

This should probably go in after #100. I'll watch for that and update here as needed.

[caspervdw]

This is a great addition to the strtree, thanks @brendan-ward for looking into this.

In my opinion, the return value of the proposed query_bulk is so different that query and query_bulk should stay different methods.

Another approach would be mimicking scipy's KDTree (https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.KDTree.query_ball_point.html#scipy.spatial.KDTree.query_ball_point). They return a list of integers (for single) or an array of lists. There is something to be said of following their API anyway. It guess it depends strongly on your use case which return value is best.

What do you think of the suggestion of an ndarray of lists? The query could in that case be turned into a O->O ufunc.

[brendan-ward]

@caspervdw I'm afraid I don't quite follow the suggestion of turning query into a ufunc.

How would we pass in the tree in that case?

I also thought that ufuncs needed to be the same length as the input? In this case, the length is variable based on the number of pairs of input geometry, tree geometry with overlapping windows or that meet the predicate.

I'm certainly open to the idea, I just don't have enough familiarity with ufuncs to see how we'd apply them in this case.

[caspervdw]

It’s just an idea to implement the “Geometry or ndarray of Geometry -> list or ndarray of lists” pattern. You could do it in a different way, but the ufunc way adds a bonus that it would accept n-dimensional arrays.

I didn’t think of the tree itself. You could add the tree as an extra parameter to the ufunc and ensure it is a scalar. I did that on the buffer implementation.

[brendan-ward]

Current test failures on AppVeyor due to differences in docstring test results between platforms:

Expected:
    (array([0, 0, 0, 1, 1]), array([2, 3, 4, 5, 6]))
Got:
    (array([0, 0, 0, 1, 1], dtype=int32), array([2, 3, 4, 5, 6], dtype=int32))

Fix coming up.

Investigation into converting to ufunc or modeling on scipy.KDTree query underway, so please don't merge yet.

[jorisvandenbossche]

I am personally not in favor of an interface like scipy.KDTree that returns object arrays of lists. Such arrays are not very efficient to further work with.
Or at least, I would prefer that there is a low-level function available that gives you the integer arrays as done here (eg to implement sjoin), a higher level function could still give the similar API as KDTree.

Do we know of other examples of such trees / APIs ?

[brendan-ward]

@jorisvandenbossche I think I had interpreted the comparsion to scipy.KDTree as a suggestion that we could return an ndarray of ndarrays, so basically what we have now but an ndarray at the top level rather than a tuple. If that is the primary choice, do you know if there is a benefit here to us doing one vs the other?

[jorisvandenbossche]

think I had interpreted the comparsion to scipy.KDTree as a suggestion that we could return an ndarray of ndarrays,

I understood it as a list of indices per geometry you are searching for. Which seems what it is doing from looking at the example of the reference page:

In [3]: >>> from scipy import spatial 
   ...: >>> x, y = np.mgrid[0:5, 0:5] 
   ...: >>> points = np.c_[x.ravel(), y.ravel()] 
   ...: >>> tree = spatial.KDTree(points)            

In [4]: tree.query_ball_point(([2, 0], [3, 3]), 1)     
Out[4]: array([list([5, 10, 11, 15]), list([13, 17, 18, 19, 23])], dtype=object)

So in this PR, you would get something like ([0, 0, 0, 0, 1, 1, 1, 1, 1], [5, 10, 11, 15, 13, 17, 18, 19, 23]) instead.

[brendan-ward]

@jorisvandenbossche more specifically, you'd get something like:

array([array([0, 0, 0, 0, 1, 1, 1, 1, 1]), array([5, 10, 11, 15, 13, 17, 18, 19, 23])])

(or as now, tuple at the top level).

Now that you show the example of KDTree that way (I had misread their example), I definitely prefer our existing approach here, since the primary intended usage is to consume in sjoin operations.

Testing out the ufunc approach now, which could look something like:

tree = pygeos.STRtree(pygeos.points(np.arange(10), np.arange(10)))
query_tree(tree, [pygeos.box(2, 2, 4, 4), pygeos.box(5, 5, 6, 6)])
# => (array([0, 0, 0, 1, 1]), array([2, 3, 4, 5, 6]))

[caspervdw]

I am personally not in favor of an interface like scipy.KDTree that returns object arrays of lists. Such arrays are not very efficient to further work with.

Not necessarily, it depends on what you are doing. If you want to get lists of geometries that correspond to each query geometry, you are faster with the lists because they are already grouped:

> arr = np.arange(1000)  
> inds = np.sort(np.random.randint(0, 1000, (10, 100)), axis=1)   
> inds_lst = np.array([list(i) for i in inds])
> %timeit arr[inds_lst[0]]                                                                                                                                                                                                   
430 ns ± 2.93 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)
> a = np.repeat(np.arange(10), 100)
> b = np.concatenate(inds)
> %timeit arr[b[a == 0]]                                                                                                                                                                                                     
2.06 µs ± 6.02 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

However this could probably be done by using the normal query method.
Looking at the spatial join of geopandas convinced why you want to do it this way. So I am 👍 on the current approach.

Or at least, I would prefer that there is a low-level function available that gives you the integer arrays as done here (eg to implement sjoin), a higher level function could still give the similar API as KDTree.

Agreed.

One more thing: you are returning a tuple of index pairs. Another approach would be to put it in a 2D array with shape (2, N). If you for example would want to list the pairs, you could just transpose it.

[brendan-ward]

@caspervdw I looked into turning this into a ufunc. I don't think that it is a good fit in this specific case:

• from a user standpoint, given that we already have tree.query() it feels like a better API to have this as a method of the tree instead of a standalone function; it feels awkward to pass the tree into another function instead of using the method.
• I couldn't come up with a good use case of passing in a multidimensional array of input geometries to the query operation ufunc, or broadcasting the tree or predicate to its shape.
• It is focused on returning relative indexes within the input geometry, which means that these would be indexes of the innermost array in the input (so we can't flatten the outputs into a 1D arrays).
• while writing a ufunc that returns outputs derived from the shape of the input seems straightforward, I was not readily able to figure out how to return outputs of a different shape.
• it wasn't immediately obvious to me if it would yield performance benefits compared to our current approach here; it felt like it was requiring more complexity than the current approach.

My primary focus for this PR is to make this easy to consume into a spatial join, which would only pass in 1D arrays of input geometries. There may be an excellent future use case for addding a ufunc wrapper to tree query; maybe we consider this when that use case is more clear?

My latest changes return the outputs as a 2D array of shape (2, n), instead of a tuple. Transposition is a nice feature here. Pending final review, this is ready!

[caspervdw]

I read through the C code and it is looking good.
Good test coverage also 👍

[brendan-ward]

Found some unhandled edge cases while testing nearest - working on fixes...

[brendan-ward]

Should be ready to merge...

[jorisvandenbossche]

Thanks for the updates!
The 2D output is fine for me

Added some minor comments

[jorisvandenbossche]

I get a compiler warning here

src/strtree.c:426:20: warning: ‘predicate_func’ may be used uninitialized in this function

[brendan-ward]

I don't get this - what compiler are you using?

I think I mitigated it by assigning it to NULL up front:

FuncGEOS_YpY_b *predicate_func = NULL;

[jorisvandenbossche]

Just gcc on Ubuntu 18.04. But seems to be fixed now!

[jorisvandenbossche]

Thanks a lot @brendan-ward !