Zonal improvements without geometry lookup #956

asinghvi17 · 2025-04-27T14:24:32Z

spatialslices = false: if true, run f on the slices of the datacube in X, Y. You can also pass a tuple of dimensions to specify which dims to slice over. If false, run zonal on the whole cube cropped to each geom. If spatialslices is not false, then a data cube will be returned, with a Geometry(1:ngeom) axis.
bylayer = true: if true, run f on each layer of the stack individually, then recombine results into a NamedTuple. If false, then run f on the whole stack.
missingval = missingval(x): specify the missing value that is to be pushed into the resulting raster.

Explanatory script

using Rasters, RasterDataSources, ArchGDAL
using NaturalEarth
import Proj

precip = cat((Raster(WorldClim{Climate}, :prec; month = i) for i in 1:12)...; dims = Ti)
all_countries = naturalearth("admin_0_countries", 10)

precip_cellarea_stack = RasterStack((; precip = precip, cellarea = cellarea(precip)))

# First, we have standard zonal.  This returns what it always used to return.
function f(ras)
    sum(ras)
end

@time zonal(f, precip_cellarea_stack[Ti(1)]; of = all_countries, skipmissing = true, threaded = false)

# Then, we have zonal with `bylayer = false`.  This gives the zonal function the entire stack,
# which means that it can operate on all layers of the stack at once.  This allows you to e.g.
# compute statistics which require the values of multiple layers, simultaneously.
# TODO: for best use, this still needs a good implementation of skipmissing on stacks.
function f(stack)
    isempty(stack) || all(ismissing, stack.precip) && return missing
    prec_net = sum(skipmissing(stack.precip .* stack.cellarea))
    avg_precip = prec_net ./ sum(skipmissing(stack.cellarea))
    return avg_precip
end

@time zonal(f, precip_cellarea_stack[Ti(1)]; of = all_countries, bylayer = false, skipmissing = false, threaded = false)

# Next, we have `spatialslices = true`.  This allows you to pass a data cube (not just a 2d raster),
# and have zonal automatically slice the cube into 2D (X, Y) slices, such that your function only needs 
# to know how to operate on 2D slices.
zonal(f, precip_cellarea_stack; of = all_countries, bylayer = false, spatialslices = true, skipmissing = false, threaded = false)

# Finally, if you return a named tuple from your zonal function, zonal will automatically
# construct a RasterStack for you.
@time zonal(precip_cellarea_stack; of = all_countries, bylayer = false, spatialslices = true, skipmissing = false, threaded = false) do stack
    return (; precip = sum(skipmissing(stack.precip)), cellarea = sum(skipmissing(stack.cellarea)))
end

f1(stack) = (; precip = sum(skipmissing(stack.precip)), cellarea = sum(skipmissing(stack.cellarea)))
s1 = @time zonal(f1, precip_cellarea_stack; of = all_countries, bylayer = false, spatialslices = true, skipmissing = false, threaded = true)
s2 = @time zonal(sum ∘ skipmissing, precip_cellarea_stack; of = all_countries, bylayer = true, spatialslices = true, skipmissing = false, threaded = true)

all(map(==, s1, s2)) # true - => bylayer is not changing the result

this will be quite useful later.

rafaqz · 2025-04-27T15:28:18Z

If we use spatial slices and bylayer together does it return a RasterStack?

asinghvi17 · 2025-04-27T15:35:16Z

That's the plan! It's almost there but there are some indexing order issues.

asinghvi17 · 2025-04-27T15:40:35Z

ah fixed it, was just overcomplicating it in my head!

@time zonal(precip_cellarea_stack; of = all_countries, bylayer = true, spatialslices = true, skipmissing = false, threaded = false) do stack
           return sum(skipmissing(stack))
       end
Applying #107 to each geometry... 100%|██████████████████████████████████████████████████| Time: 0:00:00
  0.755803 seconds (2.27 M allocations: 336.739 MiB, 5.28% gc time, 57.64% compilation time: 6% of which was recompilation)
┌ 12×258 RasterStack ┐
├────────────────────┴────────────────────────────────────────────────────────────────────────────────────────────────────────────────── dims ┐
  ↓ Ti      ,
  → Geometry Sampled{Int64} Base.OneTo(258) ForwardOrdered Regular Points
├───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────── layers ┤
  :precip   eltype: Int64 dims: Ti, Geometry size: 12×258
  :cellarea eltype: Float64 dims: Geometry size: 258
├───────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────── raster ┤
  missingval: (precip = missing, cellarea = nothing)
  extent: Extent(Ti = (1, 12), Geometry = (1, 258))
  filename: (precip = nothing, cellarea = nothing)
└─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘

lazarusA · 2025-04-27T15:43:05Z

Btw, make sure 'show progress ' false works. It is a pain in CI to deal with those prints.

asinghvi17 · 2025-04-27T15:46:46Z

Yeah, I'm actually thinking maybe we should have that be a preferences.jl thing you can configure, but that's a later thing

asinghvi17 · 2025-04-27T15:52:28Z

But there's a lot of recompilation time here (50%) and it's probably pretty slow in general, this could use a pass by you @rafaqz to clean up the code and get this fast.

asinghvi17 · 2025-04-27T15:53:07Z

here's a test script I wrote btw, some of the things still don't work but this is the idea

using Rasters, RasterDataSources, ArchGDAL
using NaturalEarth
import Proj

precip = cat((Raster(WorldClim{Climate}, :prec; month = i) for i in 1:12)...; dims = Ti)
all_countries = naturalearth("admin_0_countries", 10)

precip_cellarea_stack = RasterStack((; precip = precip, cellarea = cellarea(precip)))

zonal(precip_cellarea_stack[Ti(1)]; of = all_countries, bylayer = false, skipmissing = false, threaded = false) do stack
    isempty(stack) || all(ismissing, stack.precip) && return missing
    prec_net = sum(skipmissing(stack.precip .* stack.cellarea))
    avg_precip = prec_net ./ sum(skipmissing(stack.cellarea))
    return avg_precip
end

zonal(precip_cellarea_stack; of = all_countries, bylayer = false, spatialslices = true, skipmissing = false, threaded = false) do stack
    isempty(stack) || all(ismissing, stack.precip) && return missing
    prec_net = sum(skipmissing(stack.precip .* stack.cellarea))
    avg_precip = prec_net ./ sum(skipmissing(stack.cellarea))
    return avg_precip
end


@time zonal(precip_cellarea_stack; of = all_countries, bylayer = false, spatialslices = true, skipmissing = false, threaded = false) do stack
    return (; precip = sum(skipmissing(stack.precip)), cellarea = sum(skipmissing(stack.cellarea)))
end


s1 = @time zonal(precip_cellarea_stack; of = all_countries, bylayer = false, spatialslices = true, skipmissing = false, threaded = true) do stack
    return (; precip = sum(skipmissing(stack.precip)), cellarea = sum(skipmissing(stack.cellarea)))
end

s2 = @time zonal(precip_cellarea_stack; of = all_countries, bylayer = true, spatialslices = true, skipmissing = false, threaded = true) do stack
    return sum(skipmissing(stack))
end

asinghvi17 · 2025-04-27T23:51:22Z

Added a new test / example script to the top comment. But I'm worried runtimes will have dropped significantly with all this complexity.

asinghvi17 · 2025-04-28T00:01:14Z

Some interesting benchmarks and timer runs:

@be zonal($sum, $(precip); of = $all_countries, spatialslices = $false, progress = $false, threaded = $false) seconds=5
──────────────────────────────────────────────────────────────────────────────────────────────
                                                     Time                    Allocations      
                                            ───────────────────────   ────────────────────────
             Tot / % measured:                   7.95s /  60.1%           4.52GiB /  97.4%    

Section                             ncalls     time    %tot     avg     alloc    %tot      avg
──────────────────────────────────────────────────────────────────────────────────────────────
crop + mask                          6.19k    4.17s   87.2%   673μs   4.40GiB  100.0%   744KiB
skipmissing + spatialsliceify + f    6.19k    613ms   12.8%  99.0μs   1.22MiB    0.0%     207B
build result                            24    351μs    0.0%  14.6μs   9.75KiB    0.0%     416B
──────────────────────────────────────────────────────────────────────────────────────────────

@be zonal($sum, $(precip); of = $all_countries, spatialslices = $true, progress = $false, threaded = $false) seconds=5
──────────────────────────────────────────────────────────────────────────────────────────────
                                                     Time                    Allocations      
                                            ───────────────────────   ────────────────────────
             Tot / % measured:                   12.3s /  39.2%           4.38GiB /  97.4%    

Section                             ncalls     time    %tot     avg     alloc    %tot      avg
──────────────────────────────────────────────────────────────────────────────────────────────
crop + mask                          5.93k    4.05s   84.1%   682μs   4.21GiB   98.8%   744KiB
skipmissing + spatialsliceify + f    5.93k    760ms   15.8%   128μs   47.7MiB    1.1%  8.24KiB
build result                            23   5.03ms    0.1%   219μs   4.19MiB    0.1%   186KiB
──────────────────────────────────────────────────────────────────────────────────────────────

so it seems that spatialslices does add some (but really quite negligible) amount of time to the final answer. I'm not surprised that it allocates more during spatialsliceify though, since it is constructing an array (there is "more" output).

So I'm satisfied enough with spatialslices.

rafaqz · 2025-04-28T08:00:30Z

Yeah, larger allocations will add time. How does the base case compare to before this PR?

asinghvi17 · 2025-04-28T10:17:25Z

Same timing, I checked

rafaqz · 2025-04-28T12:42:28Z

src/methods/zonal.jl

+
+function _determine_missingval(x, bylayer)
+    raw_missingval = isnothing(missingval(x)) ? missing : missingval(x)
+    if isfalse(bylayer) && x isa AbstractRasterStack


What is all of this for? It seems type unstable

Is raw_missingval a NamedTuple?

yes, that's why I am trying to extract it.

The missingval on stacks can be like (; layer1=missing, layer2=nothing, layer3=NaN)

which I can't necessarily use as the missing value - I need a single thing, whether that's missing or nan.

Thats not really possible... Each layer can have its own missing value and it will depend on the type. It could be 0xFF for UInt8.

The issue at hand is that if bylayer = true, ok fine, I know that I can use each. But what do I use if bylayer = false? Missing is the only option barring user input.

also, is this the correct way to do it? If I do

julia> r1 = Raster(1:2, (X(1:2),),; missingval = 1) ┌ 2-element Raster{Int64, 1} ┐ ├────────────────────────────┴─────────────────────────────────────────────────────────────────────────────────────────────────────── dims ┐ ↓ X Sampled{Int64} 1:2 ForwardOrdered Regular Points ├────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────── raster ┤ missingval: 1 extent: Extent(X = (1, 2),) └──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────┘ 1 1 2 2

then none of them are actually missing...

I think you need to map over all values and missingvals and check if any match. There should be a helper function somewhere that does this.

But yeah, I guess missing is what to use?

asinghvi17 added 8 commits April 26, 2025 21:42

all zonal improvements so far

e868288

implement isfalse in utils

8f432e2

implement bylayer

2d62e96

for bylayer=false, spatialslices=true, run zonal via map instead

2f44efc

add a Geometry dimension

a0584fd

remove infiltrate

fcbd3af

add methods for istrue and isfalse on Vals

8973f74

rebuild with Geometry dim indicating source geom

51db2c8

this will be quite useful later.

reduce complication and make rasterstack work again

28446bf

asinghvi17 mentioned this pull request Apr 27, 2025

Benchmark using map in zonal #957

Closed

asinghvi17 added 2 commits April 27, 2025 19:30

replace alloc_zonal with map / maptasks

43ae2b2

inefficient, allocating way to get returned vecs of namedtuples to stack

70e0fac

fix for actual rasters

ba99500

rafaqz reviewed Apr 28, 2025

View reviewed changes

Zonal improvements without geometry lookup #956

Are you sure you want to change the base?

Zonal improvements without geometry lookup #956

Uh oh!

Conversation

asinghvi17 commented Apr 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Explanatory script

Uh oh!

rafaqz commented Apr 27, 2025

Uh oh!

asinghvi17 commented Apr 27, 2025

Uh oh!

asinghvi17 commented Apr 27, 2025

Uh oh!

lazarusA commented Apr 27, 2025

Uh oh!

asinghvi17 commented Apr 27, 2025

Uh oh!

asinghvi17 commented Apr 27, 2025

Uh oh!

asinghvi17 commented Apr 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

asinghvi17 commented Apr 27, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

asinghvi17 commented Apr 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rafaqz commented Apr 28, 2025

Uh oh!

asinghvi17 commented Apr 28, 2025

Uh oh!

rafaqz Apr 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

asinghvi17 Apr 28, 2025

Choose a reason for hiding this comment

Uh oh!

rafaqz Apr 28, 2025

Choose a reason for hiding this comment

Uh oh!

asinghvi17 Apr 28, 2025

Choose a reason for hiding this comment

Uh oh!

asinghvi17 Apr 28, 2025

Choose a reason for hiding this comment

Uh oh!

rafaqz Apr 29, 2025

Choose a reason for hiding this comment

Uh oh!

rafaqz Apr 30, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

asinghvi17 commented Apr 27, 2025 •

edited

Loading

asinghvi17 commented Apr 27, 2025 •

edited

Loading

asinghvi17 commented Apr 27, 2025 •

edited

Loading

asinghvi17 commented Apr 28, 2025 •

edited

Loading

rafaqz Apr 28, 2025 •

edited

Loading