diff --git a/src/GeometryOps.jl b/src/GeometryOps.jl index ac2349b86f..58d380eb34 100644 --- a/src/GeometryOps.jl +++ b/src/GeometryOps.jl @@ -55,6 +55,8 @@ include("utils/utils.jl") include("utils/NaturalIndexing.jl") using .NaturalIndexing +include("utils/FlexibleRTrees/FlexibleRTrees.jl") +using .FlexibleRTrees # Load utility modules in using .NaturalIndexing, .SpatialTreeInterface, .LoopStateMachine diff --git a/src/utils/FlexibleRTrees/FlexibleRTrees.jl b/src/utils/FlexibleRTrees/FlexibleRTrees.jl new file mode 100644 index 0000000000..4ba94f558b --- /dev/null +++ b/src/utils/FlexibleRTrees/FlexibleRTrees.jl @@ -0,0 +1,44 @@ +# # FlexibleRTrees + +#= +A packed (bulk-loaded, static) R-tree over `Extents.Extent`s of any +dimensionality, with a pluggable bulk-load algorithm — sort-tile-recursive +([`STR`](@ref)), Hilbert-packed ([`HPR`](@ref)), or none +([`Unsorted`](@ref)) — behind one tree type. + +Storage is flat: `RTree{A, E}` holds a vector of per-level extent vectors +plus a leaf permutation, and is a concrete type at any size or depth. A +bulk-load algorithm chooses only the *leaf order*, via [`loadorder`](@ref); +packing always unions consecutive runs of `nodecapacity` extents, bottom-up. +Upper levels therefore group runs of the leaf order rather than re-tiling +each level: Hilbert order is spatially local at every scale so `HPR` packs +tightly, while `STR`'s upper levels are slightly looser than a re-tiled +pointer tree's. + +The tree implements SpatialTreeInterface, so `depth_first_search` / +`dual_depth_first_search` (and everything built on them) work unchanged. +Leaf queries yield indices into the *original* input collection. + +Parts of the construction logic are adapted from +[SortTileRecursiveTree.jl](https://github.com/maxfreu/SortTileRecursiveTree.jl) (MIT). + +```julia +tree = RTree(HPR(), extents) # or STR(), Unsorted() +hits = query(tree, Extents.Extent(X = (0, 1), Y = (0, 1))) +``` +=# + +module FlexibleRTrees + +import GeoInterface as GI +import Extents +using StaticArrays: MVector + +export RTree, BulkLoadAlgorithm, STR, HPR, Unsorted, query + +include("types.jl") # `BulkLoadAlgorithm`s and the `RTree` type +include("bulk_loading.jl") # `loadorder` methods and bottom-up packing +include("hilbert.jl") # Hilbert keys for `HPR`'s sort +include("interface.jl") # SpatialTreeInterface implementation and `query` + +end # module FlexibleRTrees diff --git a/src/utils/FlexibleRTrees/bulk_loading.jl b/src/utils/FlexibleRTrees/bulk_loading.jl new file mode 100644 index 0000000000..262762be07 --- /dev/null +++ b/src/utils/FlexibleRTrees/bulk_loading.jl @@ -0,0 +1,70 @@ +# # Bulk loading + +# ## Leaf ordering + +""" + loadorder(algorithm, extents::Vector{<:Extents.Extent}, nodecapacity)::Vector{Int} + +The permutation in which `algorithm` packs `extents` into leaves. Implement +this for a new `BulkLoadAlgorithm` subtype to plug in another ordering. +""" +loadorder(::Unsorted, extents, nodecapacity) = collect(1:length(extents)) +loadorder(::HPR, extents, nodecapacity) = sortperm(_hilbert_keys(extents)) +function loadorder(::STR, extents::Vector{E}, nodecapacity) where E + centers = [_center(e) for e in extents] + perm = collect(1:length(extents)) + _str_tile!(perm, centers, 1, length(extents), 1, _ndims(E), nodecapacity) + return perm +end + +#= +One recursion level of N-dimensional sort-tile-recursive: sort the range by +the current dimension's center, cut it into `S ≈ P^(1/remaining)` slabs of +whole leaf pages, and tile each slab along the remaining dimensions. After +the last dimension the consecutive `nodecapacity`-runs are the leaf tiles. +=# +function _str_tile!(perm, centers, lo, hi, dim, ndims, nodecapacity) + len = hi - lo + 1 + len <= nodecapacity && return # a single leaf: internal order doesn't matter + sort!(view(perm, lo:hi); by = i -> @inbounds(centers[i][dim])) + dim == ndims && return + P = cld(len, nodecapacity) # leaf pages in this range + S = ceil(Int, P^(1 / (ndims - dim + 1))) # slabs along this dimension + slab = cld(P, S) * nodecapacity # items per slab (whole pages) + i = lo + while i <= hi + _str_tile!(perm, centers, i, min(i + slab - 1, hi), dim + 1, ndims, nodecapacity) + i += slab + end + return +end + +# ## Packing + +# Union consecutive runs of `nodecapacity` extents, bottom-up, until a level +# fits in one (implicit) root node. Returns levels coarsest-first. +function _pack_levels(leaves::Vector{E}, nodecapacity::Int) where E + levels = [leaves] + current = leaves + while length(current) > nodecapacity + nparents = cld(length(current), nodecapacity) + parents = Vector{E}(undef, nparents) + for p in 1:nparents + lo = (p - 1) * nodecapacity + 1 + hi = min(p * nodecapacity, length(current)) + acc = current[lo] + for j in (lo + 1):hi + acc = Extents.union(acc, @inbounds current[j]) + end + parents[p] = acc + end + push!(levels, parents) + current = parents + end + return reverse!(levels) +end + +# ## Extent helpers + +_ndims(::Type{Extents.Extent{K, V}}) where {K, V} = length(K) +_center(ext::Extents.Extent) = map(b -> (b[1] + b[2]) / 2, values(ext)) diff --git a/src/utils/FlexibleRTrees/hilbert.jl b/src/utils/FlexibleRTrees/hilbert.jl new file mode 100644 index 0000000000..ac4af301b4 --- /dev/null +++ b/src/utils/FlexibleRTrees/hilbert.jl @@ -0,0 +1,95 @@ +# # Hilbert curve encoding +# +#= +The HPR ("Hilbert-packed R-tree") bulk loader sorts items by the position of +their center on a Hilbert space-filling curve. The Hilbert curve visits every +cell of a `2^bits × … × 2^bits` grid exactly once, and consecutive cells along +the curve are always spatially adjacent — so items that are consecutive in +sorted order are close in space, in every dimension, at every scale. That +fractal locality is what makes simple "pack consecutive runs" tree +construction produce tight nodes on all levels. + +This is Skilling's transpose-based algorithm ("Programming the Hilbert +curve", AIP Conf. Proc. 707, 2004), which works in any number of dimensions — +unlike the lookup-table encoders (e.g. JTS's 2D-only `HilbertCode`). The +"transpose" is the Hilbert index of the point stored bit-interleaved across +the N input coordinates; we finish by de-interleaving it into a single +integer sort key. +=# + +""" + hilbert_key(coords::NTuple{N, UInt32}, bits::Int) -> UInt64 + +The Hilbert-curve index of a point on the `N`-dimensional `2^bits` grid, as a +sortable integer. Requires `N * bits <= 64`; each coordinate must be +`< 2^bits`. +""" +function hilbert_key(coords::NTuple{N, UInt32}, bits::Int) where N + X = MVector{N, UInt32}(coords) + M = one(UInt32) << (bits - 1) + # Inverse undo + Q = M + while Q > one(UInt32) + P = Q - one(UInt32) + for i in 1:N + if !iszero(X[i] & Q) + X[1] ⊻= P # invert + else # exchange + t = (X[1] ⊻ X[i]) & P + X[1] ⊻= t + X[i] ⊻= t + end + end + Q >>= 1 + end + # Gray encode + for i in 2:N + X[i] ⊻= X[i - 1] + end + t = zero(UInt32) + Q = M + while Q > one(UInt32) + !iszero(X[N] & Q) && (t ⊻= Q - one(UInt32)) + Q >>= 1 + end + for i in 1:N + X[i] ⊻= t + end + # Interleave the transpose into one key, most significant bits first. + key = zero(UInt64) + for b in (bits - 1):-1:0 + for i in 1:N + key = (key << 1) | UInt64((X[i] >> b) & one(UInt32)) + end + end + return key +end + +# Bits per dimension for the quantization grid: 16 where possible (the JTS +# HPRtree uses 12), fewer in high dimensions so the key fits in 64 bits. +hilbert_bits(N::Int) = min(16, 63 ÷ N) + +#= +Quantize the centers of a set of extents onto the Hilbert grid spanned by +their total extent, and return each center's Hilbert key. Degenerate +dimensions (zero span) collapse to grid coordinate 0. +=# +function _hilbert_keys(extents::Vector{E}) where E <: Extents.Extent + N = _ndims(E) + bits = hilbert_bits(N) + total = reduce(Extents.union, extents) + los = map(first, values(total)) + spans = map(b -> Float64(b[2] - b[1]), values(total)) + scale = Float64((1 << bits) - 1) + keys = Vector{UInt64}(undef, length(extents)) + for (j, e) in enumerate(extents) + c = _center(e) + q = ntuple(Val(N)) do i + span = spans[i] + frac = iszero(span) ? 0.0 : (Float64(c[i]) - Float64(los[i])) / span + UInt32(clamp(round(Int, frac * scale), 0, Int(scale))) + end + keys[j] = hilbert_key(q, bits) + end + return keys +end diff --git a/src/utils/FlexibleRTrees/interface.jl b/src/utils/FlexibleRTrees/interface.jl new file mode 100644 index 0000000000..3e4c140f88 --- /dev/null +++ b/src/utils/FlexibleRTrees/interface.jl @@ -0,0 +1,87 @@ +# # SpatialTreeInterface + +using ..SpatialTreeInterface +import ..SpatialTreeInterface: isspatialtree, isleaf, nchild, getchild, + child_indices_extents, depth_first_search + +""" + RTreeNode{A, E} + +A cursor into one node of an [`RTree`](@ref): the tree, the node's level +(0-based; the children of a level-`l` node live in `levels[l + 1]`), its +position within that level, and its extent. All SpatialTreeInterface +methods traverse the tree through these cursors. The children of one node +occupy one contiguous run of the next level's extent vector, which each +per-child method resolves once per node and then indexes into. At the leaf +level, `child_indices_extents` maps leaf slots through `tree.indices`, so +queries return indices into the original collection despite the packed +reordering. +""" +struct RTreeNode{A <: BulkLoadAlgorithm, E <: Extents.Extent} + tree::RTree{A, E} + level::Int # 0-based; children of a level-l node live in levels[l + 1] + index::Int # position within its level + extent::E +end + +Extents.extent(node::RTreeNode) = node.extent + +isspatialtree(::Type{<:RTree}) = true +isspatialtree(::Type{<:RTreeNode}) = true + +@inline _child_extents(node::RTreeNode) = node.tree.levels[node.level + 1] + +@inline function _child_range(node::RTreeNode, child_extents) + start_idx = (node.index - 1) * node.tree.nodecapacity + 1 + stop_idx = min(start_idx + node.tree.nodecapacity - 1, length(child_extents)) + return start_idx:stop_idx +end + +isleaf(node::RTreeNode) = node.level == length(node.tree.levels) - 1 + +nchild(node::RTreeNode) = length(_child_range(node, _child_extents(node))) + +function getchild(node::RTreeNode, i::Int) + child_index = (node.index - 1) * node.tree.nodecapacity + i + return RTreeNode(node.tree, node.level + 1, child_index, _child_extents(node)[child_index]) +end + +function getchild(node::RTreeNode) + extents = _child_extents(node) + tree, childlevel = node.tree, node.level + 1 + range = _child_range(node, extents) + return (RTreeNode(tree, childlevel, ci, @inbounds extents[ci]) for ci in range) +end + +function child_indices_extents(node::RTreeNode) + extents = _child_extents(node) + indices = node.tree.indices + range = _child_range(node, extents) + return ((@inbounds(indices[i]), @inbounds(extents[i])) for i in range) +end + +# The tree itself acts as the (implicit) root node. +_rootnode(tree::RTree) = RTreeNode(tree, 0, 1, tree.extent) + +isleaf(tree::RTree) = length(tree.levels) == 1 +nchild(tree::RTree) = length(tree.levels[1]) +getchild(tree::RTree) = getchild(_rootnode(tree)) +getchild(tree::RTree, i) = getchild(_rootnode(tree), i) +child_indices_extents(tree::RTree) = child_indices_extents(_rootnode(tree)) + +# ## Queries + +""" + query(tree::RTree, extent_or_geom) + +Indices (into the collection the tree was built from) of every leaf whose +extent intersects the given extent — or the extent of the given geometry — +in ascending order. +""" +query(tree::RTree, ext::Extents.Extent) = + sort!(depth_first_search(Base.Fix1(Extents.intersects, ext), tree)) +function query(tree::RTree, geom) + ext = GI.extent(geom) + isnothing(ext) && throw(ArgumentError("no extent found on $(typeof(geom))")) + return query(tree, ext) +end diff --git a/src/utils/FlexibleRTrees/types.jl b/src/utils/FlexibleRTrees/types.jl new file mode 100644 index 0000000000..dfb2151059 --- /dev/null +++ b/src/utils/FlexibleRTrees/types.jl @@ -0,0 +1,81 @@ +# # Types + +# ## Bulk-load algorithms + +""" + BulkLoadAlgorithm + +Supertype for the algorithms that decide the *leaf order* of an [`RTree`](@ref). +Packing is always "union consecutive runs of `nodecapacity`, bottom-up"; the +algorithm only chooses the order, via a [`loadorder`](@ref) method. +""" +abstract type BulkLoadAlgorithm end + +""" + STR() + +Sort-tile-recursive ordering (Leutenegger et al., 1997), generalized to any +dimensionality: sort by center along the first dimension, cut into slabs, +recurse within each slab on the remaining dimensions. +""" +struct STR <: BulkLoadAlgorithm end + +""" + HPR() + +Hilbert-packed ordering, as in JTS's `HPRtree`: sort by the Hilbert-curve +index of each extent's center. Hilbert order is spatially local at every +scale, which suits this tree's consecutive-run packing particularly well. +""" +struct HPR <: BulkLoadAlgorithm end + +""" + Unsorted() + +Keep the input order (no sort). Equivalent to natural indexing — good when +the input is already spatially coherent (e.g. the edges of a ring), and the +baseline the sorting algorithms have to beat. +""" +struct Unsorted <: BulkLoadAlgorithm end + +# ## The tree + +""" + RTree(algorithm::BulkLoadAlgorithm, data; nodecapacity = 16) + +A packed R-tree over the extents of `data` (anything `GI.extent` accepts — +geometries, or `Extents.Extent`s themselves), of any dimensionality, bulk +loaded in the order chosen by `algorithm`. + +The tree is flat and fully concrete: `levels[1]` is the coarsest level and +`levels[end]` holds the leaf extents in packed order, with `indices` mapping +each leaf slot back to its position in `data`. Queries through +SpatialTreeInterface therefore return indices into the original collection. +""" +struct RTree{A <: BulkLoadAlgorithm, E <: Extents.Extent} + algorithm::A + nodecapacity::Int + extent::E + levels::Vector{Vector{E}} # levels[1] = coarsest, levels[end] = leaf extents (packed order) + indices::Vector{Int} # leaf slot -> index into the original collection +end + +function RTree(algorithm::A, data; nodecapacity::Int = 16) where A <: BulkLoadAlgorithm + nodecapacity >= 2 || throw(ArgumentError("`nodecapacity` must be at least 2, got $nodecapacity")) + isnothing(iterate(data)) && throw(ArgumentError("cannot build an `RTree` from an empty collection")) + E = typeof(GI.extent(first(data))) + extents = E[GI.extent(x) for x in data] + perm = loadorder(algorithm, extents, nodecapacity) + leaves = extents[perm] + levels = _pack_levels(leaves, nodecapacity) + total = reduce(Extents.union, levels[1]) + return RTree{A, E}(algorithm, nodecapacity, total, levels, perm) +end + +Extents.extent(tree::RTree) = tree.extent + +function Base.show(io::IO, tree::RTree{A}) where A + print(io, "RTree{", nameof(A), "}(", length(tree.indices), " leaves, ", + length(tree.levels), " levels, capacity ", tree.nodecapacity, ")") +end +Base.show(io::IO, ::MIME"text/plain", tree::RTree) = Base.show(io, tree) diff --git a/src/utils/NaturalIndexing.jl b/src/utils/NaturalIndexing.jl index 080da277e4..03fb069ed6 100644 --- a/src/utils/NaturalIndexing.jl +++ b/src/utils/NaturalIndexing.jl @@ -155,33 +155,43 @@ Extents.extent(node::NaturalIndexNode) = node.extent SpatialTreeInterface.isspatialtree(::Type{<: NaturalIndex}) = true SpatialTreeInterface.isspatialtree(::Type{<: NaturalIndexNode}) = true -function SpatialTreeInterface.nchild(node::NaturalIndexNode) +# A node's children all live in one extents vector at the next level. +# Resolve that vector once per node and index into it, instead of following +# `parent_index -> levels -> level -> extents` again for every child. +@inline _child_extents(node::NaturalIndexNode) = node.parent_index.levels[node.level + 1].extents + +@inline function _child_range(node::NaturalIndexNode, child_extents) start_idx = (node.index - 1) * node.parent_index.nodecapacity + 1 - stop_idx = min(start_idx + node.parent_index.nodecapacity - 1, length(node.parent_index.levels[node.level+1].extents)) - return stop_idx - start_idx + 1 + stop_idx = min(start_idx + node.parent_index.nodecapacity - 1, length(child_extents)) + return start_idx:stop_idx end +SpatialTreeInterface.nchild(node::NaturalIndexNode) = length(_child_range(node, _child_extents(node))) + function SpatialTreeInterface.getchild(node::NaturalIndexNode, i::Int) child_index = (node.index - 1) * node.parent_index.nodecapacity + i return NaturalIndexNode( - node.parent_index, + node.parent_index, node.level + 1, # increment level by 1 child_index, # index of this particular child - node.parent_index.levels[node.level+1].extents[child_index] # the extent of this child + _child_extents(node)[child_index] # the extent of this child ) end # Get all children of a node function SpatialTreeInterface.getchild(node::NaturalIndexNode) - return (SpatialTreeInterface.getchild(node, i) for i in 1:SpatialTreeInterface.nchild(node)) + extents = _child_extents(node) + parent, childlevel = node.parent_index, node.level + 1 + range = _child_range(node, extents) + return (NaturalIndexNode(parent, childlevel, ci, @inbounds extents[ci]) for ci in range) end SpatialTreeInterface.isleaf(node::NaturalIndexNode) = node.level == length(node.parent_index.levels) - 1 function SpatialTreeInterface.child_indices_extents(node::NaturalIndexNode) - start_idx = (node.index - 1) * node.parent_index.nodecapacity + 1 - stop_idx = min(start_idx + node.parent_index.nodecapacity - 1, length(node.parent_index.levels[node.level+1].extents)) - return ((i, node.parent_index.levels[node.level+1].extents[i]) for i in start_idx:stop_idx) + extents = _child_extents(node) + range = _child_range(node, extents) + return ((i, @inbounds extents[i]) for i in range) end # implementation for "root node" / top level tree diff --git a/test/runtests.jl b/test/runtests.jl index 86c92d3002..f67aceee5f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -19,6 +19,7 @@ end @safetestset "Utils" begin include("utils/utils.jl") end @safetestset "LoopStateMachine" begin include("utils/LoopStateMachine.jl") end @safetestset "SpatialTreeInterface" begin include("utils/SpatialTreeInterface.jl") end +@safetestset "FlexibleRTrees" begin include("utils/FlexibleRTrees.jl") end @safetestset "UnitSpherical" begin include("utils/unitspherical.jl") end @safetestset "RobustCrossProduct" begin include("utils/robustcrossproduct.jl") end # Methods diff --git a/test/utils/FlexibleRTrees.jl b/test/utils/FlexibleRTrees.jl new file mode 100644 index 0000000000..58a57d17c9 --- /dev/null +++ b/test/utils/FlexibleRTrees.jl @@ -0,0 +1,86 @@ +using Test + +import GeoInterface as GI +import GeometryOps as GO +import GeometryOps.FlexibleRTrees as FRT +import GeometryOps.FlexibleRTrees: RTree, STR, HPR, Unsorted, query, hilbert_key +import GeometryOps.SpatialTreeInterface as STI +import Extents +using Random: Xoshiro + +# Random boxes with side lengths ~5% of the unit cube, in N dims. +function random_extents(rng, n, N) + dims = (:X, :Y, :Z, :M)[1:N] + return [begin + lo = ntuple(_ -> rand(rng), N) + hi = lo .+ 0.05 .* ntuple(_ -> rand(rng), N) + Extents.Extent(NamedTuple{dims}(tuple.(lo, hi))) + end for _ in 1:n] +end + +brute_force(ext, extents) = findall(e -> Extents.intersects(ext, e), extents) + +grow(ext, d) = Extents.buffer(ext, NamedTuple{keys(ext)}(ntuple(_ -> d, length(keys(ext))))) + +@testset "query ≡ brute force ($(N)D, $alg, n = $n)" for + N in (2, 3), + alg in (STR(), HPR(), Unsorted()), + n in (1, 5, 16, 17, 100, 256, 1000) + rng = Xoshiro(hash((N, n))) + extents = random_extents(rng, n, N) + tree = RTree(alg, extents; nodecapacity = 8) + queries = vcat( + random_extents(rng, 20, N), # small probes + [reduce(Extents.union, extents)], # everything + [grow(reduce(Extents.union, extents), 10.0)], # superset + [grow(e, 3.0) for e in random_extents(Xoshiro(0), 5, N)], # big probes + ) + for q in queries + @test query(tree, q) == brute_force(q, extents) + end + # A query far outside everything returns nothing. + faraway = Extents.Extent(NamedTuple{((:X, :Y, :Z, :M)[1:N])}(ntuple(_ -> (99.0, 100.0), N))) + @test isempty(query(tree, faraway)) +end + +@testset "construction and type stability" begin + rng = Xoshiro(7) + extents = random_extents(rng, 300, 2) + tree = @inferred RTree(STR(), extents) + @test tree isa RTree{STR, eltype(extents)} + @inferred RTree(HPR(), extents) + @inferred RTree(Unsorted(), extents) + # The query path is inferrable too (depth_first_search returns Vector{Int}). + q = Extents.Extent(X = (0.2, 0.4), Y = (0.2, 0.4)) + @inferred query(tree, q) + # Deep and shallow trees have the SAME concrete type — the point of the flat layout. + tiny = RTree(STR(), extents[1:3]) + @test typeof(tiny) === typeof(tree) + + # Geometries as input work through GI.extent. + lines = GO.to_edgelist(GI.LinearRing([(0.0, 0.0), (1.0, 0.0), (1.0, 1.0), (0.0, 0.0)])) + gtree = RTree(HPR(), lines) + # Edge 2 is the right side; edge 3 (the closing diagonal) has a bbox + # covering the whole square, so extent-intersects finds it too. + @test query(gtree, Extents.Extent(X = (0.9, 1.1), Y = (0.4, 0.6))) == [2, 3] + + @test_throws ArgumentError RTree(STR(), Extents.Extent{(:X, :Y)}[]) + @test_throws ArgumentError RTree(STR(), random_extents(rng, 5, 2); nodecapacity = 1) + @test occursin("RTree{HPR}", sprint(show, gtree)) +end + +@testset "Hilbert curve properties" begin + # Order-1 2D curve: the classic U through the four quadrants. + keys1 = [hilbert_key((UInt32(x), UInt32(y)), 1) for (x, y) in ((0, 0), (0, 1), (1, 1), (1, 0))] + @test keys1 == [0, 1, 2, 3] + # In any dimension: the curve visits every grid cell exactly once + # (bijectivity), and consecutive cells are adjacent (unit Manhattan step). + for (N, bits) in ((2, 4), (3, 2)) + side = 2^bits + cells = vec(collect(Iterators.product(ntuple(_ -> 0:(side - 1), N)...))) + keys = [hilbert_key(UInt32.(c), bits) for c in cells] + @test allunique(keys) + path = cells[sortperm(keys)] + @test all(sum(abs.(path[i + 1] .- path[i])) == 1 for i in 1:(length(path) - 1)) + end +end