mostly working faster line zonal

src/Rasters.jl

@@ -131,6 +131,7 @@ include("methods/burning/line.jl")
 include("methods/burning/polygon.jl")
 include("methods/burning/extents.jl")
 include("methods/burning/utils.jl")
+include("methods/burning/reductions.jl")
 
 include("methods/mask.jl")
 include("methods/rasterize.jl")

src/methods/burning/allocs.jl

@@ -11,30 +11,35 @@ function Allocs(buffer)
     return Allocs(buffer, edges, scratch, crossings)
 end
 
+function _burning_allocs(x::Nothing; 
+    nthreads=_nthreads(), 
+    threaded=true, 
+    burncheck_metadata=nothing,
+    kw...
+) 
+    threaded ? [Allocs(nothing) for _ in 1:nthreads] : [Allocs(nothing)]
+end
 function _burning_allocs(x; 
     nthreads=_nthreads(), 
     threaded=true, 
     burncheck_metadata=Metadata(),
     kw...
 ) 
+    dims = commondims(x, (XDim, YDim))
     if threaded
-        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
+        [Allocs(_init_bools(dims; metadata=deepcopy(burncheck_metadata))) for _ in 1:nthreads]
     else
-        if isnothing(x)
-            Allocs()
-        else
-            dims = commondims(x, DEFAULT_POINT_ORDER)
-            Allocs(_init_bools(dims; metadata=burncheck_metadata))
-        end
+        [Allocs(_init_bools(dims; metadata=burncheck_metadata))]
     end
 end
 
-_get_alloc(allocs::Vector{<:Allocs}) = _get_alloc(allocs[Threads.threadid()])
+function _get_alloc(allocs::Vector{<:Allocs}) 
+    if length(allocs) == 1
+        only(allocs)
+    else
+        allocs[Threads.threadid()]
+    end
+end
 _get_alloc(allocs::Allocs) = allocs
 
 # TODO include these in Allocs

src/methods/burning/line.jl

@@ -1,3 +1,44 @@
+#= This mutable object is passed into line burning closures,
+as updating object fields is type-stable in closures but 
+updating variables is not
+i: number of valid pixels found while line burning so far
+prev: previous pixel index to avoid repeats at start/end of lines
+vals: vector for extracting values to
+acc: an optional accumulator for reductions
+=#
+mutable struct LineRefs{T}
+    i::Int
+    prev::CartesianIndex{2}
+    vals::Vector{T}
+    acc::T
+    function LineRefs{T}(; init::Union{T,Nothing}=nothing) where T
+        i = 1
+        prev = CartesianIndex((typemin(Int), typemin(Int)))
+        vals = Vector{T}()
+        if isnothing(init)
+            new{T}(i, prev, vals)
+        else
+            new{T}(i, prev, vals, init)
+        end
+    end
+end
+
+# Re-initialise line refs
+function _initialise!(lr::LineRefs{T}; 
+    vals=true, init::Union{T,Nothing}=nothing
+) where T
+    lr.i = 1
+    lr.prev = CartesianIndex((typemin(Int), typemin(Int)))
+    if !isnothing(init)
+        lr.acc = init
+    end
+    if vals 
+        lr.vals = Vector{T}()
+    end
+    return lr
+end
+
+
 # _burn_lines!
 # Fill a raster with `fill` where pixels touch lines in a geom
 # Usually `fill` is `true` of `false`

src/methods/burning/reductions.jl

@@ -0,0 +1,39 @@
+# _TakeFirst: `first` with a missing value check
+# TODO: clarify why we need this  for `first` and not `last`
+struct _TakeFirst{MV}
+    missingval::MV
+end
+(tf::_TakeFirst)(a, b) = a === tf.missingval ? b : a
+_take_last(a, b) = b
+
+# Get the operator for a reduction. 
+# This can be used for online stats style optimizations
+_reduce_op(::typeof(sum)) = Base.add_sum
+_reduce_op(::typeof(prod)) = Base.mul_prod
+_reduce_op(::typeof(minimum)) = min
+_reduce_op(::typeof(maximum)) = max
+_reduce_op(::typeof(last)) = _take_last
+_reduce_op(f, missingval) = _reduce_op(f)
+_reduce_op(::typeof(first), missingval) = _TakeFirst(missingval)
+_reduce_op(x) = nothing
+
+# Can we split the reducion and combine later
+_is_op_threadsafe(::typeof(sum)) = true
+_is_op_threadsafe(::typeof(prod)) = true
+_is_op_threadsafe(::typeof(minimum)) = true
+_is_op_threadsafe(::typeof(maximum)) = true
+_is_op_threadsafe(f) = false
+
+# Get an initial value for the reduction
+# TODO: clarify the difference between returning missing and zero
+_reduce_init(reducer, st::AbstractRasterStack, missingval) = map(A -> _reduce_init(reducer, A, missingval), st)
+_reduce_init(reducer, ::AbstractRaster{T}, missingval) where T = _reduce_init(reducer, T, missingval)
+_reduce_init(reducer, nt::NamedTuple, missingval) = map(x -> _reduce_init(reducer, x, missingval), nt)
+_reduce_init(f, x, missingval) = _reduce_init(f, typeof(x), missingval)
+_reduce_init(::Nothing, x::Type{T}, missingval) where T = zero(T)
+_reduce_init(f::Function, ::Type{T}, missingval) where T = zero(f((zero(nonmissingtype(T)), zero(nonmissingtype(T)))))
+_reduce_init(::typeof(sum), ::Type{T}, missingval) where T = zero(nonmissingtype(T))
+_reduce_init(::typeof(prod), ::Type{T}, missingval) where T = oneunit(nonmissingtype(T))
+_reduce_init(::typeof(minimum), ::Type{T}, missingval) where T = typemax(nonmissingtype(T))
+_reduce_init(::typeof(maximum), ::Type{T}, missingval) where T = typemin(nonmissingtype(T))
+_reduce_init(::typeof(last), ::Type{T}, missingval) where T = _maybe_nothing_to_missing(missingval)

src/methods/extract.jl

@@ -34,25 +34,6 @@ end
 
 Base.eltype(::Extractor{T}) where T = T
 
-# This mutable object is passed into closures as 
-# fields are type-stable in clusores but variables are not
-mutable struct LineRefs{T}
-    i::Int
-    prev::CartesianIndex{2}
-    rows::Vector{T}
-    function LineRefs{T}() where T
-        i = 1
-        prev = CartesianIndex((typemin(Int), typemin(Int)))
-        new{T}(i, prev, Vector{T}())
-    end
-end
-
-function _initialise!(lr::LineRefs{T}) where T
-    lr.i = 1
-    lr.prev = CartesianIndex((typemin(Int), typemin(Int)))
-    lr.rows = Vector{T}()
-end
-
 _geom_rowtype(A, geom; kw...) =
     _geom_rowtype(A, GI.trait(geom), geom; kw...)
 _geom_rowtype(A, ::Nothing, geoms; kw...) =
@@ -297,9 +278,7 @@ end
     # places when they are right on the line.
     dims, offset = _template(A, geom)
     dp = DimPoints(A)
-    function _length_callback(n)
-        resize!(line_refs.rows, n)
-    end
+    _length_callback(n) = resize!(line_refs.vals, n)
 
     _burn_lines!(_length_callback, dims, geom) do D
         I = CartesianIndex(map(val, D))
@@ -308,11 +287,11 @@ end
         line_refs.prev = I
         # Make sure we only hit this pixel once
         # D is always inbounds
-        line_refs.i += _maybe_set_row!(line_refs.rows, e.skipmissing, e, id, A, dp, offset, I, line_refs.i)
+        line_refs.i += _maybe_set_row!(line_refs.vals, e.skipmissing, e, id, A, dp, offset, I, line_refs.i)
         return nothing
     end
-    deleteat!(line_refs.rows, line_refs.i:length(line_refs.rows))
-    return line_refs.rows
+    deleteat!(line_refs.vals, line_refs.i:length(line_refs.vals))
+    return line_refs.vals
 end
 @noinline function _extract(
     A::RasterStackOrArray, e::Extractor{T}, id::Int, ::GI.AbstractGeometryTrait, geom; kw...

src/methods/rasterize.jl

@@ -1,37 +1,3 @@
-struct _TakeFirst{MV}
-    missingval::MV
-end
-(tf::_TakeFirst)(a, b) = a === tf.missingval ? b : a
-_take_last(a, b) = b
-
-_reduce_op(::typeof(sum)) = Base.add_sum
-_reduce_op(::typeof(prod)) = Base.mul_prod
-_reduce_op(::typeof(minimum)) = min
-_reduce_op(::typeof(maximum)) = max
-_reduce_op(::typeof(last)) = _take_last
-_reduce_op(f, missingval) = _reduce_op(f)
-_reduce_op(::typeof(first), missingval) = _TakeFirst(missingval)
-_reduce_op(x) = nothing
-
-_is_op_threadsafe(::typeof(sum)) = true
-_is_op_threadsafe(::typeof(prod)) = true
-_is_op_threadsafe(::typeof(minimum)) = true
-_is_op_threadsafe(::typeof(maximum)) = true
-_is_op_threadsafe(f) = false
-
-_reduce_init(reducer, st::AbstractRasterStack, missingval) = map(A -> _reduce_init(reducer, A, missingval), st)
-_reduce_init(reducer, ::AbstractRaster{T}, missingval) where T = _reduce_init(reducer, T, missingval)
-_reduce_init(reducer, nt::NamedTuple, missingval) = map(x -> _reduce_init(reducer, x, missingval), nt)
-_reduce_init(f, x, missingval) = _reduce_init(f, typeof(x), missingval)
-
-_reduce_init(::Nothing, x::Type{T}, missingval) where T = zero(T)
-_reduce_init(f::Function, ::Type{T}, missingval) where T = zero(f((zero(nonmissingtype(T)), zero(nonmissingtype(T)))))
-_reduce_init(::typeof(sum), ::Type{T}, missingval) where T = zero(nonmissingtype(T))
-_reduce_init(::typeof(prod), ::Type{T}, missingval) where T = oneunit(nonmissingtype(T))
-_reduce_init(::typeof(minimum), ::Type{T}, missingval) where T = typemax(nonmissingtype(T))
-_reduce_init(::typeof(maximum), ::Type{T}, missingval) where T = typemin(nonmissingtype(T))
-_reduce_init(::typeof(last), ::Type{T}, missingval) where T = _maybe_nothing_to_missing(missingval)
-
 struct FillChooser{F,I,M}
     fill::F
     init::I