Faster mosaic (#839)

* faster mosaic * bugfix * remove comment * tweaks * cleanup and reorganise * add missing deps * 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 * bugfix * bugfix * half fix ncdatasets mosaic * working LinerSolver with tests * fix tests * bugfix doctests * bugfix atol * bugfix create * bugfix tests * pad float intervals before view for fp error * bugfix lazy reorder * remove a # * bugfixes * remove show * fix example * fix examples * fix tests again * fix ambiguities
rafaqz · Jan 21, 2025 · aa2fd19 · aa2fd19
1 parent 991a147
commit aa2fd19
Show file tree

Hide file tree

Showing 26 changed files with 1,383 additions and 655 deletions.
diff --git a/.github/dependabot.yml b/.github/dependabot.yml
@@ -0,0 +1,11 @@
+# To get started with Dependabot version updates, you'll need to specify which
+# package ecosystems to update and where the package manifests are located.
+# Please see the documentation for all configuration options:
+# https://docs.github.com/code-security/dependabot/dependabot-version-updates/configuration-options-for-the-dependabot.yml-file
+
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
diff --git a/Project.toml b/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"

diff --git a/ext/RastersStatsBaseExt/sample.jl b/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)
     )

diff --git a/src/create.jl b/src/create.jl
@@ -291,15 +291,15 @@ function create(filename::AbstractString, source::Source, ::Type{T}, dims::Tuple
     end
     # Create layers of zero arrays
     rast = Raster(A, dims; name, missingval=mv_inner)
+    # Write to disk
     Rasters.write(f, filename, source, rast;
-        eltype, chunks, metadata, scale, offset, missingval, verbose, force, coerce, write, kw...
+        eltype, chunks, metadata, scale, offset, missingval=mv_inner, verbose, force, coerce, write, kw...
     ) do W
-        # write returns a variable, wrap it as a Raster
+        # `write`` returns a variable, wrap it as a Raster and run f
         f(rebuild(rast; data=W))
     end
     # Don't pass in `missingval`, read it again from disk in case it changed
-    r = Raster(filename; source, lazy, metadata, dropband, coerce, missingval=mv_outer)
-    return r
+    return Raster(filename; source, lazy, metadata, dropband, coerce, missingval=mv_outer)
 end
 # Create on-disk RasterStack from filename, source, type and dims
 function create(filename::AbstractString, source::Source, layertypes::NamedTuple, dims::Tuple;

diff --git a/src/lookup.jl b/src/lookup.jl
@@ -95,15 +95,15 @@ dim(lookup::Projected) = lookup.dim
 end
 @inline function LA.selectindices(l::Projected, sel::LA.Selector{<:AbstractVector}; kw...)
     selval = reproject(mappedcrs(l), crs(l), dim(l), val(sel))
-    _selectvec(l, rebuild(sel; val=selval)sel; kw...)
+    LA._selectvec(l, rebuild(sel; val=selval); kw...)
 end
 @inline function LA.selectindices(l::Projected, sel::LA.IntSelector{<:Tuple}; kw...)
     selval = reproject(mappedcrs(l), crs(l), dim(l), val(sel))
-    _selecttuple(l, rebuild(sel; val=selval); kw...)
+    LA._selecttuple(l, rebuild(sel; val=selval); kw...)
 end
 @inline LA.selectindices(l::Projected{<:Tuple}, sel::LA.IntSelector{<:Tuple}; kw...) = LA._selectindices(l, sel; kw...)
 @inline LA.selectindices(l::Projected{<:Tuple}, sel::LA.IntSelector{<:Tuple{<:Tuple,<:Tuple}}; kw...) = 
-    _selecttuple(l, sel; kw...)
+    LA._selecttuple(l, sel; kw...)
 
 function LA.selectindices(l::Projected, sel::Between{<:Tuple})
     selval = map(v -> reproject(mappedcrs(l), crs(l), dim(l), v), val(sel))

diff --git a/src/methods/burning/allocs.jl b/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)

diff --git a/src/methods/burning/array_init.jl b/src/methods/burning/array_init.jl
@@ -4,29 +4,22 @@
 # 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::DimTuple, T::Type, data; kw...) = _init_bools(to, to, T, data; kw...)
-function _init_bools(to::Nothing, T::Type, data; 
-    geometrycolumn=nothing, 
-    collapse=nokw, 
-    res=nokw,
-    size=nokw,
-    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::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)
     isnothing(ext) && throw(ArgumentError("no recognised dimensions, extent or geometry"))
-    return _init_bools(ext, T, data; collapse, res, size, kw...)
-end
-function _init_bools(to::Extents.Extent, T::Type, data;
-    collapse=nokw, size=nokw, res=nokw, sampling=nokw, kw...
-)
-    # Convert the extent to dims (there must be `res` or `size` in `kw`)
-    dims = _extent2dims(to; size, res, sampling, kw...)
-    _init_bools(to, dims, T, data; collapse, kw...)
+    return _init_bools(ext, T, data; kw...)
 end
+_init_bools(to::Extents.Extent, T::Type, data; kw...) =
+    _init_bools(to, _extent2dims(to; kw...), T, data; kw...)
 function _init_bools(to, dims::DimTuple, T::Type, data; 
     collapse::Union{Bool,Nothing,NoKW}=nokw, kw...
 )
@@ -54,12 +47,12 @@ 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
 
-function _prepare_for_burning(B, locus=Center())
-    B1 = _forward_ordered(B)
+function _prepare_for_burning(B; locus=Center(), order=ForwardOrdered())
+    B1 = _maybe_lazy_reorder(order, B)
     start_dims = map(dims(B1, DEFAULT_POINT_ORDER)) do d
         # Shift lookup values to center of pixels
         d = DD.maybeshiftlocus(locus, d)
@@ -69,9 +62,12 @@ 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)
+_maybe_lazy_reorder(::Nothing, B) = B
+function _maybe_lazy_reorder(::ForwardOrdered, B)
     reduce(dims(B); init=B) do A, d
         if DD.order(d) isa ReverseOrdered
             A = view(A, rebuild(d, lastindex(d):-1:firstindex(d)))

diff --git a/src/methods/burning/geometry.jl b/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