initial commit

Project.toml

@@ -1,10 +1,16 @@
-name = "STRtree"
+name = "GISTRtree"
 uuid = "746ee33f-1797-42c2-866d-db2fce69d14d"
 authors = ["Max Freudenberg <[email protected]> and contributors"]
-version = "1.0.0-DEV"
+version = "0.1.0"
+
+[deps]
+Extents = "411431e0-e8b7-467b-b5e0-f676ba4f2910"
+GeoInterface = "cf35fbd7-0cd7-5166-be24-54bfbe79505f"
 
 [compat]
-julia = "1.8"
+julia = "1.6"
+Extents = "0.1.0"
+GeoInterface = "1"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

README.md

@@ -1,3 +1,19 @@
-# STRtree
+# GISTRtree
 
-[![Build Status](https://github.com/maxfreu/STRtree.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/maxfreu/STRtree.jl/actions/workflows/CI.yml?query=branch%3Amain)
+[![Build Status](https://github.com/maxfreu/GISTRtree.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/maxfreu/GISTRtree.jl/actions/workflows/CI.yml?query=branch%3Amain)
+
+An STR tree implementation for GeoInterface compatible geometries.
+
+Usage:
+
+```julia
+using GISTRtree
+using Extents
+
+tree = STRtree(geometries)
+query_result = query(tree, Extent(X=(0, 100.5), Y=(0, 1.5)))
+# or 
+query_result = query(tree, query_geometry)
+```
+
+Contributions are welcome! :)

src/GISTRtree.jl

@@ -0,0 +1,133 @@
+module GISTRtree
+
+using Extents
+import GeoInterface as GI
+
+
+"""
+    STRtree(geoms; nodecapacity=10)
+
+Construct an STRtree from a collection of geometries with the given node capacity.
+"""
+struct STRtree{T}
+    rootnode::T
+    function STRtree(geoms; nodecapacity=10)
+        rootnode = build_root_node(geoms, nodecapacity=nodecapacity)
+        return new{typeof(rootnode)}(rootnode)
+    end
+end
+
+
+struct STRNode{E,T}
+    extent::E
+    children::T
+end
+
+
+struct STRLeafNode{E}
+    extents::E
+    indices::Vector{Int}
+end
+
+
+GI.extent(n::STRNode) = n.extent
+GI.extent(n::STRLeafNode) = foldl(Extents.union, n.extents)
+
+
+function leafnodes(geoms; nodecapacity=10)
+    extents_indices = [(GI.extent(geoms[i]), i) for i in eachindex(geoms)]
+    perm = sortperm(extents_indices; by=(v -> ((v[1][1][1] + v[1][1][2]) / 2)))  # [extent/index][dim][min/max] sort by x
+    sorted_extents = extents_indices[perm]
+    r = length(sorted_extents)
+    P = ceil(Int, r / nodecapacity)
+    S = ceil(Int, sqrt(P))
+    x_splits = Iterators.partition(sorted_extents, S * nodecapacity)
+
+    nodes = STRLeafNode{Vector{typeof(extents_indices[1][1])}}[]
+    for x_split in x_splits
+        perm = sortperm(x_split; by=(v -> ((v[1][2][1] + v[1][2][2]) / 2)))  # [extent/index][dim][min/max] sort by y
+        sorted_split = x_split[perm]
+        y_splits = Iterators.partition(sorted_split, nodecapacity)
+        for y_split in y_splits
+            push!(nodes, STRLeafNode(getindex.(y_split,1), getindex.(y_split,2)))
+        end
+    end
+    return nodes
+end
+
+
+# a bit of duplication...
+function parentnodes(nodes; nodecapacity=10)
+    extents_indices = [(GI.extent(node), node) for node in nodes]
+    perm = sortperm(extents_indices; by=(v -> ((v[1][1][1] + v[1][1][2]) / 2)))  # [extent/node][dim][min/max] sort by x
+    sorted_extents = extents_indices[perm]
+    r = length(sorted_extents)
+    P = ceil(Int, r / nodecapacity)
+    S = ceil(Int, sqrt(P))
+    x_splits = Iterators.partition(sorted_extents, S * nodecapacity)
+
+    T = typeof(extents_indices[1][1])
+    N = Vector{typeof(extents_indices[1][2])}
+    nodes = STRNode{T, N}[]
+    for x_split in x_splits
+        perm = sortperm(x_split; by=(v -> ((v[1][2][1] + v[1][2][2]) / 2)))  # [extent/index][dim][min/max] sort by y
+        sorted_split = x_split[perm]
+        y_splits = Iterators.partition(sorted_split, nodecapacity)
+        for y_split in y_splits
+            push!(nodes, STRNode(foldl(Extents.union, getindex.(y_split,1)), getindex.(y_split,2)))
+        end
+    end
+    return nodes
+end
+
+
+"""recursively build root node from geometries and node capacity"""
+function build_root_node(geoms; nodecapacity=10)
+    nodes = leafnodes(geoms, nodecapacity=nodecapacity)
+    while length(nodes) > 1
+        nodes = parentnodes(nodes, nodecapacity=nodecapacity)
+    end
+    return nodes[1]
+end
+
+
+"""
+    query(tree::STRtree, extent::Extent)
+    query(tree::STRtree, geom)
+
+Query the tree for geometries whose extent intersects with the given extent or the extent of the given geometry.
+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.
+"""
+function query end
+
+function query(tree::STRtree, extent::Extent)
+    query_result = Int[]
+    query!(query_result, tree.rootnode, extent)
+    return unique(sort!(query_result))
+end
+
+query(tree::STRtree, geom) = query(tree, GI.extent(geom))
+
+"""recursively query the nodes until a leaf node is reached"""
+function query!(query_result::Vector{Int}, node::STRNode, extent::Extent)
+    if Extents.intersects(node.extent, extent)
+        for child in node.children
+            query!(query_result, child, extent)
+        end
+    end
+    return query_result
+end
+
+"""when leaf node is reached, push indices of geometries to query result"""
+function query!(query_result::Vector{Int}, node::STRLeafNode, extent::Extent)
+    for i in eachindex(node.extents)
+        if Extents.intersects(node.extents[i], extent)
+            push!(query_result, node.indices[i])
+        end
+    end
+end
+
+
+export STRtree, query
+
+end

test/Project.toml

@@ -0,0 +1,5 @@
+[deps]
+ArchGDAL = "c9ce4bd3-c3d5-55b8-8973-c0e20141b8c3"
+Extents = "411431e0-e8b7-467b-b5e0-f676ba4f2910"
+GeoInterface = "cf35fbd7-0cd7-5166-be24-54bfbe79505f"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

test/runtests.jl

@@ -1,6 +1,21 @@
-using STRtree
+using GISTRtree
 using Test
+using Extents
+import ArchGDAL as AG
+import GeoInterface as GI
 
-@testset "STRtree.jl" begin
-    # Write your tests here.
+
+@testset "GISTRtree.jl" begin
+    x = 1:100
+    y = 1:100
+    points = AG.createpoint.(x, y')
+    # polygons = AG.buffer.(points, 0.1)
+    tree = STRtree(points)
+    @test tree.rootnode isa GISTRtree.STRNode
+    @test tree.rootnode.children[1] isa GISTRtree.STRNode
+
+    query_result = query(tree, Extent(X=(0, 100.5), Y=(0, 1.5)))
+    @test query_result isa Vector{Int}
+    @test length(query_result) == 100
+    @test points[query_result] == points[:,1]
 end