WIP: geometry lookup

Project.toml

@@ -16,6 +16,7 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Flatten = "4c728ea3-d9ee-5c9a-9642-b6f7d7dc04fa"
 GeoFormatTypes = "68eda718-8dee-11e9-39e7-89f7f65f511f"
 GeoInterface = "cf35fbd7-0cd7-5166-be24-54bfbe79505f"
+GeometryOps = "3251bfac-6a57-4b6d-aa61-ac1fef2975ab"
 GeometryOpsCore = "05efe853-fabf-41c8-927e-7063c8b9f013"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Missings = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
@@ -25,6 +26,7 @@ ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
+SortTileRecursiveTree = "746ee33f-1797-42c2-866d-db2fce69d14d"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [weakdeps]
@@ -69,6 +71,7 @@ GeoDataFrames = "0.3"
 GeoFormatTypes = "0.4"
 GeoInterface = "1.0"
 GeometryBasics = "0.4"
+GeometryOps = "0.1.19"
 GeometryOpsCore = "0.1.1"
 Makie = "0.20, 0.21, 0.22"
 Missings = "0.4, 1"
@@ -83,6 +86,7 @@ Reexport = "0.2, 1.0"
 SafeTestsets = "0.1"
 Setfield = "0.6, 0.7, 0.8, 1"
 Shapefile = "0.10, 0.11"
+SortTileRecursiveTree = "0.1.1"
 Statistics = "1"
 StatsBase = "0.34"
 Test = "1"
@@ -112,3 +116,8 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Aqua", "ArchGDAL", "CFTime", "CoordinateTransformations", "DataFrames", "GeoDataFrames", "GeometryBasics", "GRIBDatasets", "Makie", "NCDatasets", "Plots", "Proj", "RasterDataSources", "SafeTestsets", "Shapefile", "StableRNGs", "StatsBase", "Test", "ZarrDatasets"]
+
+
+[sources]
+DimensionalData = {url = "https://github.com/rafaqz/DimensionalData.jl", rev = "as/individualindexing"}
+GeometryOps = {url = "https://github.com/JuliaGeo/GeometryOps.jl", rev = "as/extentforwarding_for_predicates"}

ext/RastersMakieExt/plotrecipes.jl

@@ -332,3 +332,34 @@ _prepare_dimarray(A) = DimArray(map(x -> _convert_with_missing(x, missingval(A))
 
 _convert_with_missing(x::Real, missingval) = isequal(x, missingval) || ismissing(x) ? NaN32 : Float32(x)
 _convert_with_missing(x, missingval) = isequal(x, missingval) ? missing : x
+
+Makie.expand_dimensions(::Type{<: Makie.Poly}, A::Raster{T, 1, <: Tuple{<: Dimension{<:GeometryLookup}}}) where {T} = nothing
+Makie.expand_dimensions(::Makie.NoConversion, A::Raster{T, 1, <: Tuple{<: Dimension{<:GeometryLookup}}}) where {T} = nothing
+
+function Makie.convert_arguments(::Type{<: Makie.Poly}, A::Raster{T, 1, <: Tuple{<: Dimension{<:GeometryLookup}}}) where {T}
+    geometries = lookup(only(dims(A))).data
+    color = replace_missing(A, NaN).data
+    label = string(Rasters.DD.name(A))
+    isempty(label) && (label = string(Rasters.DD.name(only(dims(A)))))
+
+    return Makie.SpecApi.Poly(
+        geometries; 
+        color = color, 
+        label = label
+    )
+end
+
+
+Makie.expand_dimensions(::Type{<: Makie.Poly}, A::Base.SkipMissing{<: Raster{T, 1, <: Tuple{<: Dimension{<:GeometryLookup}}}}) where {T} = nothing
+Makie.expand_dimensions(::Type{<: Makie.Poly}, A::Rasters.SkipMissingVal{<: Raster{T, 1, <: Tuple{<: Dimension{<:GeometryLookup}}}}) where {T} = nothing
+
+Makie.expand_dimensions(::Makie.NoConversion, A::Base.SkipMissing{<: Raster{T, 1, <: Tuple{<: Dimension{<:GeometryLookup}}}}) where {T} = nothing
+Makie.expand_dimensions(::Makie.NoConversion, A::Rasters.SkipMissingVal{<: Raster{T, 1, <: Tuple{<: Dimension{<:GeometryLookup}}}}) where {T} = nothing
+
+function Makie.convert_arguments(::Type{<: Makie.Poly}, A::Base.SkipMissing{<: Raster{T, 1, <: Tuple{<: Dimension{<:GeometryLookup}}}}) where {T}
+    return Makie.convert_arguments(Makie.Poly, A.x[ismissing.(A.x)])
+end
+
+function Makie.convert_arguments(::Type{<: Makie.Poly}, A::Rasters.SkipMissingVal{<: Raster{T, 1, <: Tuple{<: Dimension{<:GeometryLookup}}}}) where {T}
+    return Makie.convert_arguments(Makie.Poly, A.x[ismissing.(A.x)])
+end

src/Rasters.jl

@@ -21,6 +21,7 @@ import Adapt,
        FillArrays,
        Flatten,
        GeoInterface,
+       GeometryOps,
        GeometryOpsCore,
        OffsetArrays,
        ProgressMeter,
@@ -29,6 +30,7 @@ import Adapt,
        RecipesBase,
        Reexport,
        Setfield,
+       SortTileRecursiveTree,
        Statistics
 
 Reexport.@reexport using DimensionalData, GeoFormatTypes
@@ -62,8 +64,8 @@ export Planar, Spherical
 export AbstractRaster, Raster
 export AbstractRasterStack, RasterStack
 export AbstractRasterSeries, RasterSeries
-export Projected, Mapped
-export Band
+export Projected, Mapped, GeometryLookup
+export Band, Geometry
 export missingval, boolmask, missingmask, replace_missing, replace_missing!,
        aggregate, aggregate!, disaggregate, disaggregate!, mask, mask!,
        resample, warp, zonal, crop, extend, trim, slice, combine, points,
@@ -76,6 +78,7 @@ export Extent, extent
 const DD = DimensionalData
 const DA = DiskArrays
 const GI = GeoInterface
+const GO = GeometryOps
 const LA = Lookups
 
 # DimensionalData documentation urls
@@ -119,6 +122,11 @@ const RasterSeriesOrStack = Union{AbstractRasterSeries,AbstractRasterStack}
 include("utils.jl")
 include("skipmissing.jl")
 
+include("geometry_lookup/geometry_lookup.jl")
+include("geometry_lookup/lookups.jl")
+include("geometry_lookup/methods.jl")
+include("geometry_lookup/io.jl")
+
 include("table_ops.jl")
 include("create.jl")
 include("read.jl")

src/crs.jl

@@ -9,12 +9,26 @@ coordinate reference system at all.
 See [`setcrs`](@ref) to set it manually.
 """
 function GeoInterface.crs(obj::Union{<:AbstractRaster,<:AbstractRasterStack,<:AbstractRasterSeries, <:DimTuple})
-    if hasdim(obj, Y)
-        crs(dims(obj, Y))
-    elseif hasdim(obj, X)
-        crs(dims(obj, X))
+    each_dim_crs = map(crs, dims(obj))
+    firstcrs = findfirst(!isnothing, each_dim_crs)
+    if isnothing(firstcrs)
+        return nothing
     else
-        nothing
+        for (dim, crs) in zip(dims(obj), each_dim_crs)
+            if !isnothing(crs) && crs !== each_dim_crs[firstcrs]
+                throw(ArgumentError("""
+                All dimensions must have the same crs, but dims $(name(dim)) and $(name(dims(obj, firstcrs)))
+                have different CRS:
+                $(each_dim_crs[firstcrs])
+
+                and
+
+                $(crs)
+                """
+                ))
+            end
+        end
+        return each_dim_crs[firstcrs]
     end
 end
 GeoInterface.crs(dim::Dimension) = crs(lookup(dim))

src/dimensions.jl

@@ -8,7 +8,7 @@ const SpatialDim = Union{XDim,YDim,ZDim}
 
 Band [`Dimension`]($DDdimdocs) for multi-band rasters.
 
-## Example:
+## Example
 ```julia
 banddim = Band(10:10:100)
 # Or
@@ -18,3 +18,23 @@ mean(A; dims=Band)
 ```
 """
 @dim Band
+
+"""
+    Geometry <: Dimension
+
+    Geometry(geoms)
+
+Geometry [`Dimension`]($DDdimdocs) for vector data cubes.
+
+## Example
+```julia
+geomdim = Geometry(GeometryLookup(polygons))
+# Or
+val = A[Geometry(1)]
+# Or
+val = A[Geometry(Touches(other_geom))] # this is automatically accelerated by spatial trees!
+# Or
+mean(A; dims=Geometry)
+```
+"""
+@dim Geometry

src/geometry_lookup/geometry_lookup.jl

@@ -0,0 +1,169 @@
+"""
+    GeometryLookup(data, dims = (X(), Y()); geometrycolumn = nothing)
+
+
+A lookup type for geometry dimensions in vector data cubes.
+
+`GeometryLookup` provides efficient spatial indexing and lookup for geometries using an STRtree (Sort-Tile-Recursive tree).
+It is used as the lookup type for geometry dimensions in vector data cubes, enabling fast spatial queries and operations.
+
+It spans the dimensions given to it in `dims`, as well as the dimension it's wrapped in - you would construct a DimArray with a GeometryLookup
+like `DimArray(data, Geometry(GeometryLookup(data, dims)))`.  Here, `Geometry` is a dimension - but selectors in X and Y will also eventually work!
+
+
+# Examples
+
+```julia
+using Rasters
+
+using NaturalEarth
+import GeometryOps as GO
+
+# construct the polygon lookup
+polygons = NaturalEarth.naturalearth("admin_0_countries", 110).geometry
+polygon_lookup = GeometryLookup(polygons, (X(), Y()))
+
+# create a DimArray with the polygon lookup
+dv = rand(Geometry(polygon_lookup))
+
+# select the polygon with the centroid of the 88th polygon
+dv[Geometry(Contains(GO.centroid(polygons[88])))] == dv[Geometry(88)] # true
+```
+
+"""
+struct GeometryLookup{T, A <: AbstractVector{T}, D, M <: GO.Manifold, Tree, CRS} <: DD.Dimensions.MultiDimensionalLookup{T}
+    manifold::M
+    data::A
+    tree::Tree
+    dims::D
+    crs::CRS
+end
+
+function GeometryLookup(data, dims=(X(), Y()); geometrycolumn=nothing, crs=nokw, tree=nokw)
+
+    # First, retrieve the geometries - from a table, vector of geometries, etc.
+    geometries = _get_geometries(data, geometrycolumn)
+    geometries = Missings.disallowmissing(geometries)
+
+    if isnokw(crs)
+        crs = GI.crs(data)
+        if isnothing(crs)
+            crs = GI.crs(first(geometries))
+        end
+    end
+
+    # Check that the geometries are actually geometries
+    if any(!GI.isgeometry, geometries)
+        throw(ArgumentError("""
+        The collection passed in to `GeometryLookup` has some elements that are not geometries 
+        (`GI.isgeometry(x) == false` for some `x` in `data`).
+        """))
+    end
+    # Make sure there are only two dimensions
+    if length(dims) != 2
+        throw(ArgumentError("""
+        The `dims` argument to `GeometryLookup` must have two dimensions, but it has $(length(dims)) dimensions (`$(dims)`).
+        Please make sure that it has only two dimensions, like `(X(), Y())`.
+        """))
+    end
+    # Build the lookup accelerator tree
+    tree = if isnokw(tree)
+        SortTileRecursiveTree.STRtree(geometries)
+    elseif GO.SpatialTreeInterface.isspatialtree(tree)
+        if tree isa DataType
+            tree(geometries)
+        else
+            tree
+        end
+    elseif isnothing(tree)
+        nothing
+    else
+        throw(ArgumentError("""
+        Got an argument for `tree` which is not a valid spatial tree (according to `GeometryOps.SpatialTreeInterface`)
+        nor `nokw` or `nothing`
+
+        Type is $(typeof(tree))
+        """))
+    end
+    # TODO: auto manifold detection and best tree type for that manifold
+    GeometryLookup(GO.Planar(), geometries, tree, dims, crs)
+end
+
+GeoInterface.crs(l::GeometryLookup) = l.crs
+setcrs(l::GeometryLookup, crs) = rebuild(l; crs)
+
+#=
+
+## DD methods for the lookup
+
+Here we define DimensionalData's methods for the lookup.
+This is broadly standard except for the `rebuild` method, which is used to update the tree accelerator when the data changes.
+=#
+
+DD.dims(l::GeometryLookup) = l.dims
+DD.dims(d::DD.Dimension{<: GeometryLookup}) = val(d).dims
+DD.order(::GeometryLookup) = Lookups.Unordered()
+DD.parent(lookup::GeometryLookup) = lookup.data
+# TODO: format for geometry lookup
+DD.Dimensions.format(l::GeometryLookup, D::Type, values, axis::AbstractRange) = l
+
+# Make sure that the tree is rebuilt if the data changes
+function DD.rebuild(
+        lookup::GeometryLookup; 
+        data=lookup.data, tree=nokw, 
+        dims=lookup.dims, crs=nokw, 
+        manifold=nokw, metadata=nokw
+    )
+    # TODO: metadata support for geometry lookup
+    new_tree = if isnokw(tree)
+        if data == lookup.data
+            lookup.tree
+        elseif isempty(data)
+            nothing
+        else
+            SortTileRecursiveTree.STRtree(data)
+        end
+    elseif GO.SpatialTreeInterface.isspatialtree(tree)
+        if tree isa DataType
+            tree(data)
+        else
+            tree
+        end
+    else
+        SortTileRecursiveTree.STRtree(data)
+    end
+    new_crs = if isnokw(crs)
+        data_crs = GI.crs(data)
+        if isnothing(data_crs)
+            lookup.crs
+        else
+            data_crs
+        end
+    else
+        crs
+    end
+
+    new_manifold = if isnokw(manifold)
+        lookup.manifold
+    else
+        manifold
+    end
+
+    GeometryLookup(new_manifold, Missings.disallowmissing(data), new_tree, dims, new_crs)
+end
+
+
+# total_area_of_intersection = 0.0
+# current_area_of_intersection = 0.0
+# last_point = nothing
+# apply_with_signal(trait, geom) do subgeom, state
+#     if state == :start
+#         total_area_of_intersection += current_area_of_intersection
+#         current_area_of_intersection = 0.0
+#         last_point = nothing
+#     elseif state == :continue
+#         # shoelace formula for this point
+#     elseif state == :end
+#         # finish off the shoelace formula
+#     end
+# end