Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Faster mosaic #839

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Faster mosaic #839

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a5bdd95
faster mosaic
rafaqz Dec 24, 2024
03cac09
bugfix
rafaqz Dec 24, 2024
f4afb50
remove comment
rafaqz Dec 24, 2024
daa3d6b
tweaks
rafaqz Dec 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
151 changes: 119 additions & 32 deletions src/methods/mosaic.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,26 +65,31 @@ function mosaic(f::Function, r1::RasterStackOrArray, rs::RasterStackOrArray...;
mosaic(f, (r1, rs...); kw...)
end
mosaic(f::Function, regions; kw...) = _mosaic(f, first(regions), regions; kw...)
function _mosaic(f::Function, ::AbstractRaster, regions;
missingval=missingval(first(regions)), filename=nothing, suffix=nothing, kw...
function _mosaic(f::Function, r1::AbstractRaster, regions;
missingval=missing, filename=nothing, suffix=nothing, kw...
)
missingval = missingval isa Nothing ? missing : missingval
T = Base.promote_type(typeof(missingval), Base.promote_eltype(regions...))
V = Vector{promote_type(map(Missings.nonmissingtype ∘ eltype, regions)...)}
T = Base.promote_op(f, V)
dims = _mosaic(Tuple(map(DD.dims, regions)))
l1 = first(regions)
A = create(filename, T, dims; name=name(l1), missingval, metadata=metadata(l1))
A = create(filename, T, dims; name=name(r1), missingval)
if isnothing(missingval)
A .= zero(eltype(A))
else
@show Rasters.missingval(A)
A .= Rasters.missingval(A)
end
open(A; write=true) do a
_mosaic!(f, a, regions; missingval, kw...)
mosaic!(f, a, regions; missingval=Rasters.missingval(A), kw...)
end
return A
end
function _mosaic(f::Function, ::AbstractRasterStack, regions;
filename=nothing, suffix=keys(first(regions)), kw...
function _mosaic(f::Function, r1::AbstractRasterStack, regions;
filename=nothing, suffix=keys(r1), kw...
)
layers = map(suffix, map(values, regions)...) do s, A...
mosaic(f, A...; filename, suffix=s, kw...)
end
return DD.rebuild_from_arrays(first(regions), Tuple(layers))
return DD.rebuild_from_arrays(r1, layers)
end

"""
Expand Down Expand Up @@ -137,47 +142,129 @@ $EXPERIMENTAL
"""
mosaic!(f::Function, dest::RasterStackOrArray, regions::RasterStackOrArray...; kw...) =
_mosaic!(f, dest, regions; kw...)

function _mosaic!(f::Function, A::AbstractRaster{T}, regions::Union{Tuple,AbstractArray};
function mosaic!(
f::Function, dest::RasterStackOrArray,
regions::Union{Tuple,AbstractArray{<:RasterStackOrArray}};
op=_reduce_op(f, missingval(dest)),
kw...
)
_mosaic!(f, op, dest, regions; kw...)
end
function mosaic!(
f::typeof(mean), op::Nothing, dest::RasterStackOrArray,
regions::Union{Tuple,AbstractArray{<:RasterStackOrArray}};
kw...
)
if length(regions) < 256
_mosaic_mean!(dest, UInt8, regions; kw...)
else
_mosiac_mean!(dest, UInt16, regions; kw...)
end
end
function mosaic!(
f::typeof(length), op::Nothing, dest::RasterStackOrArray,
regions::Union{Tuple,AbstractArray{<:RasterStackOrArray}};
kw...
)
for region in regions
_count_region!(dest, region; kw...)
end
return dest
end
# Where there is a known reduction operator we can apply each region as a whole
function _mosaic!(
f::Function, op::Function, dest::RasterStackOrArray, regions::Union{Tuple,AbstractArray};
kw...
)
for region in regions
_mosaic_region!(op, dest, region; kw...)
end
return dest
end
# Generic unknown functions
function _mosaic!(
f::Function, op::Nothing, A::AbstractRaster{T}, regions::Union{Tuple,AbstractArray};
missingval=missingval(A), atol=maybe_eps(T)
) where T
R = promote_type(map(Missings.nonmissingtype ∘ eltype, regions)...)
buffer = Vector{R}(undef, length(regions))
_without_mapped_crs(A) do A1
broadcast!(A1, DimKeys(A1; atol)) do ds
broadcast!(A1, DimSelectors(A1; atol)) do ds
# Get all the regions that have this point
ls = foldl(regions; init=()) do acc, l
if DD.hasselection(l, ds)
v = l[ds...]
(acc..., l)
else
acc
end
end
values = foldl(ls; init=()) do acc, l
v = l[ds...]
if isnothing(Rasters.missingval(l))
(acc..., v)
elseif ismissing(Rasters.missingval(l))
ismissing(v) ? acc : (acc..., v)
else
v === Rasters.missingval(l) ? acc : (acc..., v)
i = 0
for r in regions
if DD.hasselection(r, ds)
x = r[ds...]
if x !== Rasters.missingval(r)
i += 1
buffer[i] = x
end
end
end
if length(values) === 0
if i === 0
missingval
else
f(values)
f(view(buffer, 1:i))
end
end
end
return A
end
function _mosaic!(f::Function, st::AbstractRasterStack, regions::Union{Tuple,AbstractArray}; kw...)
function _mosaic!(f::Function, op::Nothing, st::AbstractRasterStack, regions::Union{Tuple,AbstractArray}; kw...)
map(values(st), map(values, regions)...) do A, r...
mosaic!(f, A, r...; kw...)
end
return st
end

function _mosaic_mean!(dest, ::Type{T}, regions; kw...) where T
# Note: sum and count are separate broadcasts because
# most disk formats don't support writing a tuple

# Define a Raster to count into
counts = create(nothing, T, dest; missingval=zero(T))
counts .= zero(T)
for region in regions
# Add region to dest
_mosaic_region!(Base.add_sum, dest, region; kw...)
# Count region
_count_region!(counts, region; kw...)
end
# Divide dest by counts
@show parent(counts)
# Avoid divide by zero for missing values
dest .= ((d, c) -> d === missingval(dest) ? missingval(dest) : d / c).(dest, counts)
return dest
end
function _mosaic_region!(op, dest, region; kw...)
function skip_or_op(a, b)
if b === missingval(region)
a
elseif a === missingval(dest)
b
else
op(a, b)
end
end
ds = DimSelectors(view(dest, extent(region)))
dest[extent(region)] .= skip_or_op.(view(dest, extent(region)), view(region, ds))
end
function _count_region!(count::AbstractRaster{T}, region::AbstractRaster; kw...) where T
@show missingval(count)
function skip_or_count(a, b)
if b === missingval(region)
a
elseif a === missingval(count)
oneunit(Missings.nonmissingtype(T))
else
a + oneunit(a)
end
end
ext = extent(region)
ds = DimSelectors(view(count, ext))
view(count, ext) .= skip_or_count.(view(count, ext), view(region, ds))
end

_mosaic(alldims::Tuple{<:DimTuple,Vararg{DimTuple}}) = map(_mosaic, alldims...)
function _mosaic(dims::Dimension...)
map(dims) do d
Expand Down
2 changes: 1 addition & 1 deletion src/methods/rasterize.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
struct _TakeFirst{MV}
struct _TakeFirst{MV} <: Function
missingval::MV
end
(tf::_TakeFirst)(a, b) = a === tf.missingval ? b : a
Expand Down
61 changes: 41 additions & 20 deletions test/methods.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ pointvec = [(-20.0, 30.0),
(0.0, 10.0),
(0.0, 30.0),
(-20.0, 30.0)]


vals = [1, 2, 3, 4, 5]
polygon = ArchGDAL.createpolygon(pointvec)
multi_polygon = ArchGDAL.createmultipolygon([[pointvec]])
Expand Down Expand Up @@ -409,7 +411,7 @@ createpoint(args...) = ArchGDAL.createpoint(args...)
@test all(extract(rast, [(Y=0.1, X=9.0), (Y=0.2, X=10.0), (Y=0.3, X=10.0)]) .=== T[
(geometry = (Y = 0.1, X = 9.0), test = 1)
(geometry = (Y = 0.2, X = 10.0), test = 4)
(geometry = (Y = 0.3, X = 10.0), test = missing)
(geometry = (Y = .3, X = 10.0), test = missing)
])
# Vector points
@test all(extract(rast, [[9.0, 0.1], [10.0, 0.2]]) .== [
Expand Down Expand Up @@ -662,25 +664,44 @@ end
missing 1.3 1.4])

@test all(mosaic(first, [reverse(reg2; dims=Y), reverse(reg1; dims=Y)]) .===
[missing 0.2 0.1;
1.2 1.1 0.3;
1.4 1.3 missing]
)

# 3 dimensions
A1 = Raster(ones(2, 2, 2), (X(2.0:-1.0:1.0), Y(5.0:1.0:6.0), Ti(DateTime(2001):Year(1):DateTime(2002))))
A2 = Raster(zeros(2, 2, 2), (X(3.0:-1.0:2.0), Y(4.0:1.0:5.0), Ti(DateTime(2002):Year(1):DateTime(2003))))
@test all(mosaic(mean, A1, A2) |> parent .===
mosaic(mean, RasterStack(A1), RasterStack(A2)).layer1 .===
cat([missing missing missing
missing 1.0 1.0
missing 1.0 1.0 ],
[0.0 0.0 missing
0.0 0.5 1.0
missing 1.0 1.0 ],
[0.0 0.0 missing
0.0 0.0 missing
missing missing missing], dims=3))
[missing 0.2 0.1;
1.2 1.1 0.3;
1.4 1.3 missing]
)

@testset "Generic functions" begin
@test all(mosaic(xs -> count(x -> x > 0, xs), reg1, reg2) .===
[1 1 missing
1 2 1
missing 1 1]
)
end

@testset "3 dimensions" begin
A1 = Raster(ones(2, 2, 2), (X(2.0:-1.0:1.0), Y(5.0:1.0:6.0), Ti(DateTime(2001):Year(1):DateTime(2002))))
A2 = Raster(zeros(2, 2, 2), (X(3.0:-1.0:2.0), Y(4.0:1.0:5.0), Ti(DateTime(2002):Year(1):DateTime(2003))))
mean_mos = cat([missing missing missing
missing 1.0 1.0
missing 1.0 1.0 ],
[0.0 0.0 missing
0.0 0.5 1.0
missing 1.0 1.0 ],
[0.0 0.0 missing
0.0 0.0 missing
missing missing missing], dims=3)
@test all(mosaic(mean, A1, A2) .===
mosaic(mean, RasterStack(A1), RasterStack(A2)).layer1 .===
mean_mos)
@test mosaic(length, A1, A2; missingval=0) == cat([0 0 0
0 1 1
0 1 1],
[1 1 0
1 2 1
0 1 1],
[1 1 0
1 1 0
0 0 0], dims=3)
end
end

using StableRNGs, StatsBase
Expand Down
Loading