add SpatialTreeInterface implementation

(but no tests yet, need to add those!)

Author: asinghvi17

Project.toml

@@ -4,6 +4,7 @@ authors = ["Anshul Singhvi <[email protected]>", "Rafael Schouten <rafaels
 version = "0.1.18"
 
 [deps]
+AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 AdaptivePredicates = "35492f91-a3bd-45ad-95db-fcad7dcfedb7"
 CoordinateTransformations = "150eb455-5306-5404-9cee-2592286d6298"
 DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
@@ -32,6 +33,7 @@ GeometryOpsProjExt = "Proj"
 GeometryOpsTGGeometryExt = "TGGeometry"
 
 [compat]
+AbstractTrees = "0.4"
 AdaptivePredicates = "1.2"
 CoordinateTransformations = "0.5, 0.6"
 DataAPI = "1"

src/utils/SpatialTreeInterface/README.md

@@ -0,0 +1,47 @@
+# SpatialTreeInterface.jl
+
+A simple interface for spatial tree types.
+
+## What is a spatial tree?
+
+- 2 dimensional extents
+- Parent nodes encompass all leaf nodes
+- Leaf nodes contain references to the geometries they represent as indices (or so we assume here)
+
+## Why is this useful?
+
+- It allows us to write algorithms that can work with any spatial tree type, without having to know the details of the tree type.
+    - for example, dual tree traversal / queries
+- It allows us to flexibly and easily swap out and use different tree types, depending on the problem at hand.
+
+This is also a zero cost interface if implemented correctly!  Verified implementations exist for "flat" trees like the "Natural Index" from `tg`, and "hierarchical" trees like the `STRtree` from `SortTileRecursiveTree.jl`.
+
+## Interface
+
+- `isspatialtree(tree)::Bool`
+- `isleaf(node)::Bool` - is the node a leaf node?  In this context, a leaf node is a node that does not have other nodes as its children, but stores a list of indices and extents (even if implicit).
+- `getchild(node)` - get the children of a node.  This may be materialized if necessary or available, but can also be lazy (like a generator).
+- `getchild(node, i)` - get the `i`-th child of a node.
+- `nchild(node)::Int` - the number of children of a node.
+- `child_indices_extents(node)` - an iterator over the indices and extents of the children of a **leaf** node.
+
+These are the only methods that are required to be implemented.  
+
+Optionally, one may define:
+- `node_extent(node)` - get the extent of a node.  This falls back to `GI.extent` but can potentially be overridden if you want to return a different but extent-like object.
+
+They enable the generic query functions described below:
+
+## Query functions
+
+- `do_query(f, predicate, node)` - call `f(i)` for each index `i` in `node` that satisfies `predicate(extent(i))`.
+- `do_dual_query(f, predicate, tree1, tree2)` - call `f(i1, i2)` for each index `i1` in `tree1` and `i2` in `tree2` that satisfies `predicate(extent(i1), extent(i2))`.
+
+These are both completely non-allocating, and will only call `f` for indices that satisfy the predicate.
+You can of course build a standard query interface on top of `do_query` if you want - that's simply:
+```julia
+a = Int[]
+do_query(Base.Fix1(push!, a), predicate, node)
+```
+where `predicate` might be `Base.Fix1(Extents.intersects, extent_to_query)`.
+

src/utils/SpatialTreeInterface/SpatialTreeInterface.jl

@@ -0,0 +1,244 @@
+module SpatialTreeInterface
+
+import ..LoopStateMachine: @controlflow
+
+import Extents
+import GeoInterface as GI
+import AbstractTrees
+
+# public isspatialtree, getchild, nchild, child_indices_extents, node_extent
+export query, do_query
+
+# ## Interface
+# Interface definition for spatial tree types.
+# There is no abstract supertype here since it's impossible to enforce,
+# but we do have a few methods that are common to all spatial tree types.
+
+"""
+    isspatialtree(tree)::Bool
+
+Return true if the object is a spatial tree, false otherwise.
+
+## Implementation notes
+
+For type stability, if your spatial tree type is `MyTree`, you should define
+`isspatialtree(::Type{MyTree}) = true`, and `isspatialtree(::MyTree)` will forward
+to that method automatically.
+"""
+isspatialtree(::T) where T = isspatialtree(T)
+isspatialtree(::Type{<: Any}) = false
+
+
+"""
+    getchild(node)
+
+Return an iterator over all the children of a node.
+This may be materialized if necessary or available,
+but can also be lazy (like a generator).
+"""
+getchild(node) = AbstractTrees.children(node)
+
+"""
+    getchild(node, i)
+
+Return the `i`-th child of a node.
+"""
+getchild(node, i) = getchild(node)[i]
+
+"""
+    nchild(node)
+
+Return the number of children of a node.
+"""
+nchild(node) = length(getchild(node))
+
+"""
+    isleaf(node)
+
+Return true if the node is a leaf node, i.e., there are no "children" below it.
+[`getchild`](@ref) should still work on leaf nodes, though, returning an iterator over the extents stored in the node - and similarly for `getnodes.`
+"""
+isleaf(node) = error("isleaf is not implemented for node type $(typeof(node))")
+
+"""
+    child_indices_extents(node)
+
+Return an iterator over the indices and extents of the children of a node.
+
+Each value of the iterator should take the form `(i, extent)`.
+
+This can only be invoked on leaf nodes!
+"""
+function child_indices_extents(node)
+    return zip(1:nchild(node), getchild(node))
+end
+
+"""
+    node_extent(node)
+
+Return the extent like object of the node.  
+Falls back to `GI.extent` by default, which falls back
+to `Extents.extent`.  
+
+Generally, defining `Extents.extent(node)` is sufficient here, and you
+won't need to define this
+
+The reason we don't use that directly is to give users of this interface
+a way to define bounding boxes that are not extents, like spherical caps 
+and other such things.
+"""
+node_extent(node) = GI.extent(node)
+
+# ## Query functions
+# These are generic functions that work with any spatial tree type that implements the interface.
+
+
+"""
+    do_query(f, predicate, tree)
+
+Call `f(i)` for each index `i` in the tree that satisfies `predicate(extent(i))`.
+
+This is generic to anything that implements the SpatialTreeInterface, particularly the methods
+[`isleaf`](@ref), [`getchild`](@ref), and [`child_extents`](@ref).
+"""
+function do_query(f::F, predicate::P, node::N) where {F, P, N}
+    if isleaf(node)
+        for (i, leaf_geometry_extent) in child_indices_extents(node)
+            if predicate(leaf_geometry_extent)
+                @controlflow f(i)
+            end
+        end
+    else
+        for child in getchild(node)
+            if predicate(node_extent(child))
+                @controlflow do_query(f, predicate, child)
+            end
+        end
+    end
+end
+function do_query(predicate, node)
+    a = Int[]
+    do_query(Base.Fix1(push!, a), predicate, node)
+    return a
+end
+
+
+"""
+    query(tree, predicate)
+
+Return a sorted list of indices of the tree that satisfy the predicate.
+"""
+function query(tree, predicate)
+    a = Int[]
+    do_query(Base.Fix1(push!, a), sanitize_predicate(predicate), tree)
+    return sort!(a)
+end
+
+
+"""
+    sanitize_predicate(pred)
+
+Convert a predicate to a function that returns a Boolean.
+
+If `pred` is an Extent, convert it to a function that returns a Boolean by intersecting with the extent.
+If `pred` is a geometry, convert it to an extent first, then wrap in Extents.intersects.
+
+Otherwise, return the predicate unchanged.
+
+
+Users and developers may overload this function to provide custom behaviour when something is passed in.
+"""
+sanitize_predicate(pred) = sanitize_predicate(GI.trait(pred), pred)
+sanitize_predicate(::Nothing, pred) = pred
+sanitize_predicate(::GI.AbstractTrait, pred) = sanitize_predicate(GI.extent(pred))
+sanitize_predicate(pred::Extents.Extent) = Base.Fix1(Extents.intersects, pred)
+
+
+"""
+    do_dual_query(f, predicate, node1, node2)
+
+Call `f(i1, i2)` for each index `i1` in `node1` and `i2` in `node2` that satisfies `predicate(extent(i1), extent(i2))`.
+
+This is generic to anything that implements the SpatialTreeInterface, particularly the methods
+[`isleaf`](@ref), [`getchild`](@ref), and [`child_extents`](@ref).
+"""
+function do_dual_query(f::F, predicate::P, node1::N1, node2::N2) where {F, P, N1, N2}
+    if isleaf(node1) && isleaf(node2)
+        # both nodes are leaves, so we can just iterate over the indices and extents
+        for (i1, extent1) in child_indices_extents(node1)
+            for (i2, extent2) in child_indices_extents(node2)
+                if predicate(extent1, extent2)
+                    @controlflow f(i1, i2)
+                end
+            end
+        end
+    elseif isleaf(node1) # node2 is not a leaf, node1 is - recurse further into node2
+        for child in getchild(node2)
+            if predicate(node_extent(node1), node_extent(child))
+                @controlflow do_dual_query(f, predicate, node1, child)
+            end
+        end
+    elseif isleaf(node2) # node1 is not a leaf, node2 is - recurse further into node1
+        for child in getchild(node1)
+            if predicate(node_extent(child), node_extent(node2))
+                @controlflow do_dual_query(f, predicate, child, node2)
+            end
+        end
+    else # neither node is a leaf, recurse into both children
+        for child1 in getchild(node1)
+            for child2 in getchild(node2)
+                if predicate(node_extent(child1), node_extent(child2))
+                    @controlflow do_dual_query(f, predicate, child1, child2)
+                end
+            end
+        end
+    end
+end
+
+# Finally, here's a sample implementation of the interface for STRtrees
+
+using SortTileRecursiveTree: STRtree, STRNode, STRLeafNode
+
+nchild(tree::STRtree) = nchild(tree.rootnode)
+getchild(tree::STRtree) = getchild(tree.rootnode)
+getchild(tree::STRtree, i) = getchild(tree.rootnode, i)
+isleaf(tree::STRtree) = isleaf(tree.rootnode)
+child_indices_extents(tree::STRtree) = child_indices_extents(tree.rootnode)
+
+
+nchild(node::STRNode) = length(node.children)
+getchild(node::STRNode) = node.children
+getchild(node::STRNode, i) = node.children[i]
+isleaf(node::STRNode) = false # STRNodes are not leaves by definition
+
+isleaf(node::STRLeafNode) = true
+child_indices_extents(node::STRLeafNode) = zip(node.indices, node.extents)
+
+
+"""
+    FlatNoTree(iterable_of_geoms_or_extents)
+
+Represents a flat collection with no tree structure, i.e., a brute force search.
+This is cost free, so particularly useful when you don't want to build a tree!
+"""
+struct FlatNoTree{T}
+    geometries::T
+end
+
+isleaf(tree::FlatNoTree) = true
+
+# NOTE: use pairs instead of enumerate here, so that we can support 
+# iterators or collections that define custom `pairs` methods.
+# This includes things like filtered extent lists, for example,
+# so we can perform extent thinning with no allocations.
+function child_indices_extents(tree::FlatNoTree{T}) where T
+    # This test only applies at compile time and should be optimized away in any case.
+    # And we can use multiple dispatch to override anyway, but it should be cost free I think.
+    if applicable(Base.keys, T) 
+        return ((i, GI.extent(obj)) for (i, obj) in pairs(tree.geometries))
+    else
+        return ((i, GI.extent(obj)) for (i, obj) in enumerate(tree.geometries))
+    end
+end
+
+end # module SpatialTreeInterface