rasterize GeometryCollection (#531)

* rasterize GeometryCollection

* more testing of rasterize GeometryCollection

* bugfix

* fix tests

* move code to burning

src/Rasters.jl

@@ -119,6 +119,8 @@ include("methods/burning/geometry.jl")
 include("methods/burning/point.jl")
 include("methods/burning/line.jl")
 include("methods/burning/polygon.jl")
+include("methods/burning/extents.jl")
+include("methods/burning/utils.jl")
 
 include("methods/shared_docstrings.jl")
 include("methods/mask.jl")

src/methods/burning/extents.jl

@@ -0,0 +1,83 @@
+const XYExtent = Extents.Extent{(:X,:Y),Tuple{Tuple{Float64,Float64},Tuple{Float64,Float64}}}
+
+# Get the bounds of a geometry
+_extent(geom; kw...)::XYExtent = _extent(GI.trait(geom), geom; kw...)
+function _extent(::Nothing, data::AbstractVector; kw...)::XYExtent
+    g1 = first(data)
+    if GI.trait(g1) isa GI.PointTrait 
+        xs = extrema(p -> GI.x(p), data)
+        ys = extrema(p -> GI.y(p), data)
+        return _float64_xy_extent(Extents.Extent(X=xs, Y=ys))
+    else
+        ext = reduce(data; init=_extent(first(data))) do ext, geom
+            Extents.union(ext, _extent(geom))
+        end
+        return _float64_xy_extent(ext)
+    end
+end
+_extent(::Nothing, data::RasterStackOrArray; kw...)::XYExtent = _float64_xy_extent(Extents.extent(data))
+function _extent(::Nothing, data::T; geometrycolumn=nothing)::XYExtent where T
+    if Tables.istable(T)
+        singlecolumn = isnothing(geometrycolumn) ? first(GI.geometrycolumns(data)) : geometrycolumn
+        cols = Tables.columns(data)
+        if singlecolumn isa Symbol && singlecolumn in Tables.columnnames(cols)
+            # Table of geometries
+            geoms = Tables.getcolumn(data, singlecolumn)
+            return _extent(nothing, geoms)
+        else
+            multicolumn = isnothing(geometrycolumn) ? DEFAULT_POINT_ORDER : geometrycolumn 
+            # TODO: test this branch
+            # Table of points with dimension columns
+            bounds = reduce(multicolumn; init=(;)) do acc, key
+                if key in Tables.columnnames(cols)
+                    merge(acc, (; key=extrema(cols[key])))
+                else
+                    acc
+                end
+            end
+            return _float64_xy_extent(Extensts.Extent(bounds))
+        end
+    else
+        ext = Extents.extent(data)
+        ext isa Extents.Extent || throw(ArgumentError("object returns `nothing` from `Extents.extent`."))
+        return _float64_xy_extent(ext)
+    end
+end
+function _extent(::GI.AbstractPointTrait, point; kw...)::XYExtent
+    x, y = Float64(GI.x(point)), Float64(GI.y(point))
+    Extents.Extent(X=(x, x), Y=(y, y))
+end
+function _extent(::GI.AbstractGeometryTrait, geom; kw...)::XYExtent
+    geomextent = GI.extent(geom; fallback=false)
+    if isnothing(geomextent)
+        points = GI.getpoint(geom)
+        xbounds = extrema(GI.x(p) for p in points)
+        ybounds = extrema(GI.y(p) for p in points)
+        return _float64_xy_extent(Extents.Extent(X=xbounds, Y=ybounds))
+    else
+        return _float64_xy_extent(geomextent)
+    end
+end
+_extent(::GI.AbstractFeatureTrait, feature; kw...)::XYExtent = _extent(GI.geometry(feature))
+function _extent(::GI.GeometryCollectionTrait, collection; kw...)::XYExtent
+    geometries = GI.getgeom(collection)
+    init = _float64_xy_extent(_extent(first(geometries)))
+    ext = reduce(geometries; init) do acc, g
+        Extents.union(acc, _extent(g))
+    end
+    return _float64_xy_extent(ext)
+end
+function _extent(::GI.AbstractFeatureCollectionTrait, features; kw...)::XYExtent
+    features = GI.getfeature(features)
+    init = _float64_xy_extent(_extent(first(features)))
+    ext = reduce(features; init) do acc, f
+        Extents.union(acc, _extent(f))
+    end
+    return _float64_xy_extent(ext)
+end
+
+function _float64_xy_extent(ext::Extents.Extent)
+    xbounds = map(Float64, ext.X)
+    ybounds = map(Float64, ext.Y)
+    return Extents.Extent(X=xbounds, Y=ybounds)
+end

src/methods/burning/point.jl

@@ -1,7 +1,16 @@
+# TODO use rasterize_points for this instead
 
 # _fill_point!
 # Fill a raster with `fill` where points are inside raster pixels
 @noinline _fill_point!(x::RasterStackOrArray, geom; kw...) = _fill_point!(x, GI.geomtrait(geom), geom; kw...)
+@noinline function _fill_point!(x::RasterStackOrArray, ::GI.GeometryCollectionTrait, geom; kw...)
+    _without_mapped_crs(x) do x1
+        for geom in GI.getgeom(geom)
+            _fill_point!(x, geom; kw...)
+        end
+    end
+    return true
+end
 @noinline function _fill_point!(x::RasterStackOrArray, ::GI.AbstractGeometryTrait, geom; kw...)
     # Just find which pixels contain the points, and set them to true
     _without_mapped_crs(x) do x1

src/methods/burning/utils.jl

@@ -0,0 +1,19 @@
+_geomindices(geoms) = _geomindices(GI.trait(geoms), geoms) 
+_geomindices(::Nothing, geoms) = eachindex(geoms)
+_geomindices(::GI.FeatureCollectionTrait, geoms) = 1:GI.nfeature(geoms)
+_geomindices(::GI.FeatureTrait, geoms) = _geomindices(GI.geometry(geoms))
+_geomindices(::GI.AbstractGeometryTrait, geoms) = 1:GI.ngeom(geoms)
+
+_getgeom(geoms, i::Integer) = _getgeom(GI.trait(geoms), geoms, i)
+_getgeom(::GI.FeatureCollectionTrait, geoms, i::Integer) = GI.geometry(GI.getfeature(geoms, i))
+_getgeom(::GI.FeatureTrait, geoms, i::Integer) = GI.getgeom(GI.geometry(geoms), i)
+_getgeom(::GI.AbstractGeometryTrait, geoms, i::Integer) = GI.getgeom(geom, i)
+_getgeom(::GI.PointTrait, geom, i::Integer) = error("PointTrait should not be reached")
+_getgeom(::Nothing, geoms, i::Integer) = geoms[i] # Otherwise we can probably just index?
+
+_getgeom(geoms) = _getgeom(GI.trait(geoms), geoms)
+_getgeom(::GI.FeatureCollectionTrait, geoms) = (GI.geometry(f) for f in GI.getfeature(geoms))
+_getgeom(::GI.FeatureTrait, geoms) = GI.getgeom(GI.geometry(geoms))
+_getgeom(::GI.AbstractGeometryTrait, geoms) = GI.getgeom(geoms)
+_getgeom(::GI.PointTrait, geom) = error("PointTrait should not be reached")
+_getgeom(::Nothing, geoms) = geoms

src/methods/rasterize.jl

@@ -182,6 +182,9 @@ function Rasterizer(::GI.AbstractFeatureTrait, feature; fill, kw...)
     # fillval = _featurefillval(feature, fill)
     Rasterizer(GI.geometry(feature), fill, fillitr; kw...)
 end
+function Rasterizer(::GI.GeometryCollectionTrait, collection; kw...)
+    Rasterizer(collect(GI.getgeom(collection)); kw...)
+end
 function Rasterizer(::Nothing, geoms; fill, kw...)
     fillitr = _iterable_fill(geoms, fill)
     Rasterizer(geoms, fill, fillitr; kw...)
@@ -629,11 +632,29 @@ function _rasterize_points!(A, ::GI.AbstractGeometryTrait, geom, fill, r::Raster
     fill1 =_iterable_fill(points, fill)
     _rasterize_points!(A, nothing, points, fill1, r)
 end
+function _rasterize_points!(A, ::GI.GeometryCollectionTrait, collection, fill, r::Rasterizer)
+    # TODO How to handle fill when there is another level of nesting
+    hasburned = false
+    for geom in _getgeom(collection)
+        hasburned |= _rasterize_points!(A, geom, fill, r)
+    end
+    return hasburned
+end
 function _rasterize_points!(A, ::Nothing, geoms, fillitr, r::Rasterizer)
     (; reducer, op, missingval, init) = r
     t1 = GI.trait(first(skipmissing(geoms)))
     hasburned = false
-    if !(t1 isa GI.PointTrait)
+    if t1 isa GI.PointTrait
+        # Get extent information to properly shift the points
+        # to the region of the array during rounding
+        ext = Extents.extent(A)
+        xrange = ext.X[2] - ext.X[1]
+        yrange = ext.Y[2] - ext.Y[1]
+        xsize = size(A, X)
+        ysize = size(A, Y)
+        s = (; ext, xrange, yrange, xsize, ysize)
+        return _rasterize_points_inner!(A, geoms, fillitr, s, reducer, op, missingval, init)
+    else
         # Recurse down until we hit points
         if r.fillitr isa Function
             for geom in _getgeom(geoms)
@@ -653,15 +674,6 @@ function _rasterize_points!(A, ::Nothing, geoms, fillitr, r::Rasterizer)
         end
         return hasburned
     end
-    # Get extent information to properly shift the points
-    # to the region of the array during rounding
-    ext = Extents.extent(A)
-    xrange = ext.X[2] - ext.X[1]
-    yrange = ext.Y[2] - ext.Y[1]
-    xsize = size(A, X)
-    ysize = size(A, Y)
-    s = (; ext, xrange, yrange, xsize, ysize)
-    _rasterize_points_inner!(A, geoms, fillitr, s, reducer, op, missingval, init)
 end
 
 @noinline function _rasterize_points_inner!(A, geoms, fillitr::F, s, reducer::R, op::O, missingval, init)::Bool where {F,O,R}