|
| 1 | +module GISTRtree |
| 2 | + |
| 3 | +using Extents |
| 4 | +import GeoInterface as GI |
| 5 | + |
| 6 | + |
| 7 | +""" |
| 8 | + STRtree(geoms; nodecapacity=10) |
| 9 | +
|
| 10 | +Construct an STRtree from a collection of geometries with the given node capacity. |
| 11 | +""" |
| 12 | +struct STRtree{T} |
| 13 | + rootnode::T |
| 14 | + function STRtree(geoms; nodecapacity=10) |
| 15 | + rootnode = build_root_node(geoms, nodecapacity=nodecapacity) |
| 16 | + return new{typeof(rootnode)}(rootnode) |
| 17 | + end |
| 18 | +end |
| 19 | + |
| 20 | + |
| 21 | +struct STRNode{E,T} |
| 22 | + extent::E |
| 23 | + children::T |
| 24 | +end |
| 25 | + |
| 26 | + |
| 27 | +struct STRLeafNode{E} |
| 28 | + extents::E |
| 29 | + indices::Vector{Int} |
| 30 | +end |
| 31 | + |
| 32 | + |
| 33 | +GI.extent(n::STRNode) = n.extent |
| 34 | +GI.extent(n::STRLeafNode) = foldl(Extents.union, n.extents) |
| 35 | + |
| 36 | + |
| 37 | +function leafnodes(geoms; nodecapacity=10) |
| 38 | + extents_indices = [(GI.extent(geoms[i]), i) for i in eachindex(geoms)] |
| 39 | + perm = sortperm(extents_indices; by=(v -> ((v[1][1][1] + v[1][1][2]) / 2))) # [extent/index][dim][min/max] sort by x |
| 40 | + sorted_extents = extents_indices[perm] |
| 41 | + r = length(sorted_extents) |
| 42 | + P = ceil(Int, r / nodecapacity) |
| 43 | + S = ceil(Int, sqrt(P)) |
| 44 | + x_splits = Iterators.partition(sorted_extents, S * nodecapacity) |
| 45 | + |
| 46 | + nodes = STRLeafNode{Vector{typeof(extents_indices[1][1])}}[] |
| 47 | + for x_split in x_splits |
| 48 | + perm = sortperm(x_split; by=(v -> ((v[1][2][1] + v[1][2][2]) / 2))) # [extent/index][dim][min/max] sort by y |
| 49 | + sorted_split = x_split[perm] |
| 50 | + y_splits = Iterators.partition(sorted_split, nodecapacity) |
| 51 | + for y_split in y_splits |
| 52 | + push!(nodes, STRLeafNode(getindex.(y_split,1), getindex.(y_split,2))) |
| 53 | + end |
| 54 | + end |
| 55 | + return nodes |
| 56 | +end |
| 57 | + |
| 58 | + |
| 59 | +# a bit of duplication... |
| 60 | +function parentnodes(nodes; nodecapacity=10) |
| 61 | + extents_indices = [(GI.extent(node), node) for node in nodes] |
| 62 | + perm = sortperm(extents_indices; by=(v -> ((v[1][1][1] + v[1][1][2]) / 2))) # [extent/node][dim][min/max] sort by x |
| 63 | + sorted_extents = extents_indices[perm] |
| 64 | + r = length(sorted_extents) |
| 65 | + P = ceil(Int, r / nodecapacity) |
| 66 | + S = ceil(Int, sqrt(P)) |
| 67 | + x_splits = Iterators.partition(sorted_extents, S * nodecapacity) |
| 68 | + |
| 69 | + T = typeof(extents_indices[1][1]) |
| 70 | + N = Vector{typeof(extents_indices[1][2])} |
| 71 | + nodes = STRNode{T, N}[] |
| 72 | + for x_split in x_splits |
| 73 | + perm = sortperm(x_split; by=(v -> ((v[1][2][1] + v[1][2][2]) / 2))) # [extent/index][dim][min/max] sort by y |
| 74 | + sorted_split = x_split[perm] |
| 75 | + y_splits = Iterators.partition(sorted_split, nodecapacity) |
| 76 | + for y_split in y_splits |
| 77 | + push!(nodes, STRNode(foldl(Extents.union, getindex.(y_split,1)), getindex.(y_split,2))) |
| 78 | + end |
| 79 | + end |
| 80 | + return nodes |
| 81 | +end |
| 82 | + |
| 83 | + |
| 84 | +"""recursively build root node from geometries and node capacity""" |
| 85 | +function build_root_node(geoms; nodecapacity=10) |
| 86 | + nodes = leafnodes(geoms, nodecapacity=nodecapacity) |
| 87 | + while length(nodes) > 1 |
| 88 | + nodes = parentnodes(nodes, nodecapacity=nodecapacity) |
| 89 | + end |
| 90 | + return nodes[1] |
| 91 | +end |
| 92 | + |
| 93 | + |
| 94 | +""" |
| 95 | + query(tree::STRtree, extent::Extent) |
| 96 | + query(tree::STRtree, geom) |
| 97 | +
|
| 98 | +Query the tree for geometries whose extent intersects with the given extent or the extent of the given geometry. |
| 99 | +Returns a vector of indices of the geometries that can be used to index into the original collection of geometries under the assumption that the collection has not been modified since the tree was built. |
| 100 | +""" |
| 101 | +function query end |
| 102 | + |
| 103 | +function query(tree::STRtree, extent::Extent) |
| 104 | + query_result = Int[] |
| 105 | + query!(query_result, tree.rootnode, extent) |
| 106 | + return unique(sort!(query_result)) |
| 107 | +end |
| 108 | + |
| 109 | +query(tree::STRtree, geom) = query(tree, GI.extent(geom)) |
| 110 | + |
| 111 | +"""recursively query the nodes until a leaf node is reached""" |
| 112 | +function query!(query_result::Vector{Int}, node::STRNode, extent::Extent) |
| 113 | + if Extents.intersects(node.extent, extent) |
| 114 | + for child in node.children |
| 115 | + query!(query_result, child, extent) |
| 116 | + end |
| 117 | + end |
| 118 | + return query_result |
| 119 | +end |
| 120 | + |
| 121 | +"""when leaf node is reached, push indices of geometries to query result""" |
| 122 | +function query!(query_result::Vector{Int}, node::STRLeafNode, extent::Extent) |
| 123 | + for i in eachindex(node.extents) |
| 124 | + if Extents.intersects(node.extents[i], extent) |
| 125 | + push!(query_result, node.indices[i]) |
| 126 | + end |
| 127 | + end |
| 128 | +end |
| 129 | + |
| 130 | + |
| 131 | +export STRtree, query |
| 132 | + |
| 133 | +end |
0 commit comments