Breaking: Line and Polygon extract optimisation (#824)

* faster line extract

* performance

* extract in a separate file

* flatten

* bugfixes

* some more comments

* more perf tweaks

* bugfix

* fix extract ambiguities

* more ambiguity

* Bugfix geomtrait

* arg fixes

* add id column

* tests and bugfixes

* bugfix

* set id default to _False()

* fix tests

Author: rafaqz

Project.toml

@@ -60,7 +60,7 @@ ConstructionBase = "1"
 CoordinateTransformations = "0.6.2"
 DataFrames = "1"
 DimensionalData = "0.29.4"
-DiskArrays = "0.3, 0.4"
+DiskArrays = "0.4"
 Extents = "0.1"
 FillArrays = "0.12, 0.13, 1"
 Flatten = "0.4"

ext/RastersStatsBaseExt/sample.jl

@@ -52,6 +52,7 @@ function _getindex(::Type{T}, x::AbstractRasterStack{<:Any, NT}, dims, idx) wher
     RA._maybe_add_fields(
         T, 
         NT(x[RA.commondims(idx, x)]), 
+        nothing,
         DimPoints(dims)[RA.commondims(idx, dims)], 
         val(idx)
     )

src/methods/burning/allocs.jl

@@ -11,22 +11,36 @@ function Allocs(buffer)
     return Allocs(buffer, edges, scratch, crossings)
 end
 
-function _burning_allocs(x; nthreads=_nthreads(), threaded=true, kw...) 
-    dims = commondims(x, DEFAULT_POINT_ORDER)
+function _burning_allocs(x; 
+    nthreads=_nthreads(), 
+    threaded=true, 
+    burncheck_metadata=Metadata(),
+    kw...
+) 
     if threaded
-        [Allocs(_init_bools(dims; metadata=Metadata())) for _ in 1:nthreads]
+        if isnothing(x)
+            [Allocs(nothing) for _ in 1:nthreads]
+        else
+            dims = commondims(x, DEFAULT_POINT_ORDER)
+            [Allocs(_init_bools(dims; metadata=deepcopy(burncheck_metadata))) for _ in 1:nthreads]
+        end
     else
-        Allocs(_init_bools(dims; metadata=Metadata()))
+        if isnothing(x)
+            Allocs()
+        else
+            dims = commondims(x, DEFAULT_POINT_ORDER)
+            Allocs(_init_bools(dims; metadata=burncheck_metadata))
+        end
     end
 end
 
 _get_alloc(allocs::Vector{<:Allocs}) = _get_alloc(allocs[Threads.threadid()])
 _get_alloc(allocs::Allocs) = allocs
 
-
 # TODO include these in Allocs
 _alloc_burnchecks(n::Int) = fill(false, n)
 _alloc_burnchecks(x::AbstractArray) = _alloc_burnchecks(length(x))
+
 function _set_burnchecks(burnchecks, metadata::Metadata{<:Any,<:Dict}, verbose)
     metadata["missed_geometries"] = .!burnchecks
     verbose && _burncheck_info(burnchecks)

src/methods/burning/array_init.jl

@@ -1,14 +1,19 @@
 
-# Like `create` but without disk writes, mostly for Bool/Union{Missing,Boo},
+# Like `create` but without disk writes, mostly for Bool or Union{Missing,Bool},
 # and uses `similar` where possible
 # TODO merge this with `create` somehow
 _init_bools(to; kw...) = _init_bools(to, BitArray; kw...)
 _init_bools(to, T::Type; kw...) = _init_bools(to, T, nothing; kw...)
-_init_bools(to::AbstractRasterSeries, T::Type, data; kw...) = _init_bools(first(to), T, data; kw...)
-_init_bools(to::AbstractRasterStack, T::Type, data; kw...) = _init_bools(first(to), T, data; kw...)
-_init_bools(to::AbstractRaster, T::Type, data; kw...) = _init_bools(to, dims(to), T, data; kw...)
-_init_bools(to::Extents.Extent, T::Type, data; kw...) = _init_bools(to, _extent2dims(to; kw...), T, data; kw...)
-_init_bools(to::DimTuple, T::Type, data; kw...) = _init_bools(to, to, T, data; kw...)
+_init_bools(to::AbstractRasterSeries, T::Type, data; kw...) =
+    _init_bools(dims(first(to)), T, data; kw...)
+_init_bools(to::AbstractRasterStack, T::Type, data; kw...) =
+    _init_bools(dims(to), dims(to), T, data; kw...)
+_init_bools(to::AbstractRaster, T::Type, data; kw...) =
+    _init_bools(to, dims(to), T, data; kw...)
+_init_bools(to::Extents.Extent, T::Type, data; kw...) =
+    _init_bools(to, _extent2dims(to; kw...), T, data; kw...)
+_init_bools(to::DimTuple, T::Type, data; kw...) =
+    _init_bools(to, to, T, data; kw...)
 function _init_bools(to::Nothing, T::Type, data; geometrycolumn=nothing,kw...)
     # Get the extent of the geometries
     ext = _extent(data; geometrycolumn, kw...)
@@ -38,7 +43,7 @@ function _alloc_bools(to, dims::DimTuple, ::Type{BitArray}; missingval::Bool=fal
 end
 function _alloc_bools(to, dims::DimTuple, ::Type{<:Array{T}}; missingval=false, metadata=NoMetadata(), kw...) where T
     # Use an `Array`
-    data = fill!(Raster{T}(undef, dims), missingval) 
+    data = fill!(Raster{T}(undef, dims), missingval)
     return rebuild(data; missingval, metadata)
 end
 
@@ -53,7 +58,9 @@ function _prepare_for_burning(B, locus=Center())
 end
 
 # Convert to Array if its not one already
-_lookup_as_array(d::Dimension) = parent(lookup(d)) isa Array ? d : modify(Array, d) 
+_lookup_as_array(x) = setdims(x, _lookup_as_array(dims(x)))
+_lookup_as_array(dims::Tuple) = map(_lookup_as_array, dims)
+_lookup_as_array(d::Dimension) = parent(lookup(d)) isa Array ? d : modify(Array, d)
 
 function _forward_ordered(B)
     reduce(dims(B); init=B) do A, d

src/methods/burning/geometry.jl

@@ -22,10 +22,11 @@ function _burn_geometry!(B::AbstractRaster, trait::Nothing, data;
     lock=Threads.SpinLock(), 
     verbose=true, 
     progress=true, 
-    threaded=true,
+    threaded=false,
     fill=true,
-    allocs=_burning_allocs(B; threaded), 
     geometrycolumn=nothing,
+    burncheck_metadata=Metadata(),
+    allocs=_burning_allocs(B; threaded, burncheck_metadata), 
     kw...
 )::Bool
     geoms = _get_geometries(data, geometrycolumn)
@@ -36,8 +37,8 @@ function _burn_geometry!(B::AbstractRaster, trait::Nothing, data;
             geom = _getgeom(geoms, i)
             ismissing(geom) && return nothing
             a = _get_alloc(allocs)
-            B1 = a.buffer
-            burnchecks[i] = _burn_geometry!(B1, geom; fill, allocs=a, lock, kw...)
+            buffer = a.buffer
+            burnchecks[i] = _burn_geometry!(buffer, geom; fill, allocs=a, lock, kw...)
             return nothing
         end
         if fill
@@ -99,14 +100,17 @@ function _burn_geometry!(B::AbstractRaster, ::GI.AbstractGeometryTrait, geom;
             return false
         end
         # Take a view of the geometry extent
-        B1 = view(B, Touches(geomextent))
-        buf1 = _init_bools(B1; missingval=false)
+        V = view(B, Touches(geomextent))
+        buffer = _init_bools(V; missingval=false)
         # Burn the polygon into the buffer
-        allocs = isnothing(allocs) ? Allocs(B) : allocs
-        hasburned = _burn_polygon!(buf1, geom; shape, geomextent, allocs, boundary, kw...)
-        @inbounds for i in eachindex(B1)
-            if buf1[i]
-                B1[i] = fill
+        # We allocate a new bitarray for the view for performance
+        # and always fill with `true`. 
+        allocs = isnothing(allocs) ? Allocs(nothing) : allocs
+        hasburned = _burn_polygon!(buffer, geom; shape, geomextent, allocs, boundary, kw...)
+        # We then transfer burned `true` values to B via the V view
+        @inbounds for i in eachindex(V)
+            if buffer[i]
+                V[i] = fill
             end
         end
     else

src/methods/burning/line.jl

@@ -1,33 +1,53 @@
 # _burn_lines!
 # Fill a raster with `fill` where pixels touch lines in a geom
-# Separated for a type stability function barrier
-function _burn_lines!(B::AbstractRaster, geom; fill=true, kw...)
-    _check_intervals(B)
+# Usually `fill` is `true` of `false`
+function _burn_lines!(
+    B::AbstractRaster, geom; fill=true, verbose=false, kw...
+)
+    _check_intervals(B, verbose)
     B1 = _prepare_for_burning(B)
-    return _burn_lines!(B1, geom, fill)
+
+    xdim, ydim = dims(B, DEFAULT_POINT_ORDER)
+    regular = map((xdim, ydim)) do d
+        # @assert (parent(lookup(d)) isa Array)
+        lookup(d) isa AbstractSampled && span(d) isa Regular
+    end
+    msg = """
+        Can only fill lines where dimensions have `Regular` lookups.
+        Consider using `boundary=:center`, reprojecting the crs,
+        or make an issue in Rasters.jl on github if you need this to work.
+        """
+    all(regular) || throw(ArgumentError(msg))
+
+    # Set indices of B as `fill` when a cell is found to burn.
+    _burn_lines!(identity, dims(B1), geom) do D
+        @inbounds B1[D] = fill
+    end
 end
 
-_burn_lines!(B, geom, fill) =
-    _burn_lines!(B, GI.geomtrait(geom), geom, fill)
-function _burn_lines!(B::AbstractArray, ::Union{GI.MultiLineStringTrait}, geom, fill)
+_burn_lines!(f::F, c::C, dims::Tuple, geom) where {F<:Function,C<:Function} =
+    _burn_lines!(f, c, dims, GI.geomtrait(geom), geom)
+function _burn_lines!(
+    f::F, c::C, dims::Tuple, ::Union{GI.MultiLineStringTrait}, geom
+) where {F<:Function,C<:Function}
     n_on_line = 0
     for linestring in GI.getlinestring(geom)
-        n_on_line += _burn_lines!(B, linestring, fill)
+        n_on_line += _burn_lines!(f, c, dims, linestring)
     end
     return n_on_line
 end
 function _burn_lines!(
-    B::AbstractArray, ::Union{GI.MultiPolygonTrait,GI.PolygonTrait}, geom, fill
-)
+    f::F, c::C, dims::Tuple, ::Union{GI.MultiPolygonTrait,GI.PolygonTrait}, geom
+) where {F<:Function,C<:Function}
     n_on_line = 0
     for ring in GI.getring(geom)
-        n_on_line += _burn_lines!(B, ring, fill)
+        n_on_line += _burn_lines!(f, c, dims, ring)
     end
     return n_on_line
 end
 function _burn_lines!(
-    B::AbstractArray, ::GI.AbstractCurveTrait, linestring, fill
-)
+    f::F, c::C, dims::Tuple, ::GI.AbstractCurveTrait, linestring
+) where {F<:Function,C<:Function}
     isfirst = true
     local firstpoint, laststop
     n_on_line = 0
@@ -46,46 +66,42 @@ function _burn_lines!(
             stop=(x=GI.x(point), y=GI.y(point)),
         )
         laststop = line.stop
-        n_on_line += _burn_line!(B, line, fill)
+        n_on_line += _burn_line!(f, c, dims, line)
     end
     return n_on_line
 end
 function _burn_lines!(
-    B::AbstractArray, t::GI.LineTrait, line, fill
-)
+    f::F, c::C, dims::Tuple, t::GI.LineTrait, line
+) where {F<:Function,C<:Function}
     p1, p2 = GI.getpoint(t, line)
     line1 = (
         start=(x=GI.x(p1), y=GI.y(p1)),
         stop=(x=GI.x(p2), y=GI.y(p2)),
     )
-    return _burn_line!(B, line1, fill)
+    return _burn_line!(f, c, dims, line1)
 end
 
 # _burn_line!
 #
 # Line-burning algorithm
 # Burns a single line into a raster with value where pixels touch a line
 #
+# Function `f` does the actual work when passed a Dimension Tuple of a pixel to burn, 
+# and `c` is an initialisation callback that is passed the maximyum 
+# number of times `f` will be called. It may be called less than that.
+#
 # TODO: generalise to Irregular spans?
-function _burn_line!(A::AbstractRaster, line, fill)
-
-    xdim, ydim = dims(A, DEFAULT_POINT_ORDER)
-    regular = map((xdim, ydim)) do d
-        @assert (parent(lookup(d)) isa Array)
-        lookup(d) isa AbstractSampled && span(d) isa Regular
-    end
-    msg = """
-        Can only fill lines where dimensions have `Regular` lookups.
-        Consider using `boundary=:center`, reprojecting the crs,
-        or make an issue in Rasters.jl on github if you need this to work.
-        """
-    all(regular) || throw(ArgumentError(msg))
+function _burn_line!(f::Function, c::Function, dims::Tuple, line::NamedTuple)
+    xdim, ydim = DD.dims(dims, DEFAULT_POINT_ORDER)
+    di = DimIndices(dims)
 
+    @assert xdim isa XDim
+    @assert ydim isa YDim
     @assert order(xdim) == order(ydim) == Lookups.ForwardOrdered()
     @assert locus(xdim) == locus(ydim) == Lookups.Center()
 
-    raster_x_step = abs(step(span(A, X)))
-    raster_y_step = abs(step(span(A, Y)))
+    raster_x_step = abs(step(span(xdim)))
+    raster_y_step = abs(step(span(ydim)))
     raster_x_offset = @inbounds xdim[1] - raster_x_step / 2 # Shift from center to start of pixel
     raster_y_offset = @inbounds ydim[1] - raster_y_step / 2
 
@@ -117,15 +133,18 @@ function _burn_line!(A::AbstractRaster, line, fill)
     # Int starting points for the line. +1 converts to julia indexing
     j, i = trunc(Int, relstart.x) + 1, trunc(Int, relstart.y) + 1 # Int
 
-    # For arbitrary dimension indexing
-    dimconstructors = map(DD.basetypeof, (xdim, ydim))
+    n_on_line = 0
+    
+    # Pass of number of runs of `f` to callback `c`
+    # This can help with e.g. allocating a vector
+    c(manhattan_distance + 1)
 
     if manhattan_distance == 0
-        D = map((d, o) -> d(o), dimconstructors, (j, i))
-        if checkbounds(Bool, A, D...)
-            @inbounds A[D...] = fill
+        D = map(rebuild, dims, (j, i))
+        if checkbounds(Bool, di, D...)
+            f(D)
+            n_on_line += 1
         end
-        n_on_line = 1
         return n_on_line
     end
 
@@ -153,18 +172,17 @@ function _burn_line!(A::AbstractRaster, line, fill)
     n_on_line = 0
     countx = county = 0
 
-
     # Int steps to move allong the line
     step_j = signbit(diff_x) * -2 + 1
     step_i = signbit(diff_y) * -2 + 1
 
     # Travel one grid cell at a time. Start at zero for the current cell
     for _ in 0:manhattan_distance
-        D = map((d, o) -> d(o), dimconstructors, (j, i))
-        if checkbounds(Bool, A, D...)
-            @inbounds A[D...] = fill
+        D = map(rebuild, dims, (j, i))
+        if checkbounds(Bool, di, D...)
+            f(D)
+            n_on_line += 1
         end
-
         # Only move in either X or Y coordinates, not both.
         if abs(max_x) < abs(max_y)
             max_x += delta_x
@@ -176,12 +194,6 @@ function _burn_line!(A::AbstractRaster, line, fill)
             county +=1
         end
     end
-    return n_on_line
-end
-function _burn_line!(A::AbstractRaster, line, fill, order::Tuple{Vararg{Dimension}})
-    msg = """"
-        Converting a `:line` geometry to raster is currently only implemented for 2d lines.
-        Make a Rasters.jl github issue if you need this for more dimensions.
-        """
-    throw(ArgumentError(msg))
+
+    return n_on_line 
 end

src/methods/burning/polygon.jl

@@ -24,14 +24,14 @@ function _burn_polygon!(B::AbstractDimArray, trait, geom;
     # Lines
     n_on_line = 0
     if boundary !== :center
-        _check_intervals(B, boundary)
+        _check_intervals(B, boundary, verbose)
         if boundary === :touches 
-            if _check_intervals(B, boundary)
+            if _check_intervals(B, boundary, verbose)
                 # Add line pixels
                 n_on_line = _burn_lines!(B, geom; fill)::Int
             end
         elseif boundary === :inside 
-            if _check_intervals(B, boundary)
+            if _check_intervals(B, boundary, verbose)
                 # Remove line pixels
                 n_on_line = _burn_lines!(B, geom; fill=!fill)::Int
             end
@@ -135,9 +135,9 @@ end
 
 const INTERVALS_INFO = "makes more sense on `Intervals` than `Points` and will have more correct results. You can construct dimensions with a `X(values; sampling=Intervals(Center()))` to achieve this"
 
-@noinline _check_intervals(B) = 
-    _chki(B) ? true : (@info "burning lines $INTERVALS_INFO"; false)
-@noinline _check_intervals(B, boundary) =
-    _chki(B) ? true : (@info "`boundary=:$boundary` $INTERVALS_INFO"; false)
+@noinline _check_intervals(B, verbose::Bool) = 
+    _chki(B) ? true : (verbose && @info "burning lines $INTERVALS_INFO"; false)
+@noinline _check_intervals(B, boundary, verbose::Bool) =
+    _chki(B) ? true : (verbose && @info "`boundary=:$boundary` $INTERVALS_INFO"; false)
 
 _chki(B) = all(map(s -> s isa Intervals, sampling(dims(B))))