DGGSpace as a regridding source descends its own subtree (42x)

Project.toml

@@ -4,7 +4,7 @@ version = "0.1.0"
 authors = ["Anshul Singhvi <anshulsinghvi@gmail.com>", "Fabian Gans <fabian.gans@bgc-jena.mpi.de>", "Maarten Pronk <git@evetion.nl>", "Alexander Kmoch <alexander.kmoch@ut.ee>", "and contributors"]
 
 [workspace]
-projects = ["test", "docs", "lib/DiscreteGlobalGridsConformanceTesting", "lib/GlobalRegridding", "lib/DiscreteGlobalGridVisualizations", "lib/DiscreteGlobalGridVisualizations/test"]
+projects = ["test", "docs", "bench", "lib/DiscreteGlobalGridsConformanceTesting", "lib/GlobalRegridding", "lib/DiscreteGlobalGridVisualizations", "lib/DiscreteGlobalGridVisualizations/test"]
 
 [deps]
 ConservativeRegridding = "8e50ac2c-eb48-49bc-a402-07c87b949343"
@@ -28,8 +28,8 @@ MakieExt = ["Makie"]
 DiscreteGlobalGridsZarrExt = ["Zarr"]
 
 [sources]
-ConservativeRegridding = {rev = "claude/budget-frontier", url = "https://github.com/JuliaGeo/ConservativeRegridding.jl"}
-GeometryOps = {rev = "as/intersection_area", url = "https://github.com/JuliaGeo/GeometryOps.jl.git"}
+ConservativeRegridding = {rev = "6a4b997ab2e66dea45b5b62b5b2f32f0a3b279b0", url = "https://github.com/JuliaGeo/ConservativeRegridding.jl"}
+GeometryOps = {rev = "2825c476647cfe6791dbb89fffb8592697b46995", url = "https://github.com/JuliaGeo/GeometryOps.jl.git"}
 GeometryOpsCore = {rev = "main", subdir = "GeometryOpsCore", url = "https://github.com/JuliaGeo/GeometryOps.jl.git"}
 # Not a workspace member: it resolves its own Manifest, so
 # `julia --project=lib/GlobalRegridding` stays self-contained.

bench/Project.toml

@@ -0,0 +1,33 @@
+# Benchmarking environment for DiscreteGlobalGrids.jl.
+#
+# A workspace member, so it shares the root Manifest and the same resolved
+# versions the package and its tests use. Adds only measurement tools and the
+# raster I/O needed to feed real Copernicus DEM data through the regridder.
+#
+#     julia -t auto --project=bench bench/regrid_bench.jl
+
+[deps]
+ArchGDAL = "c9ce4bd3-c3d5-55b8-8973-c0e20141b8c3"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+ConservativeRegridding = "8e50ac2c-eb48-49bc-a402-07c87b949343"
+DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
+DiscreteGlobalGrids = "443be784-8f3c-41ab-9ff0-108bddc7e690"
+DiskArrays = "3c3547ce-8d99-4f5e-a174-61eb10b00ae3"
+Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
+Extents = "411431e0-e8b7-467b-b5e0-f676ba4f2910"
+FlameGraphs = "08572546-2f56-4bcf-ba4e-bab62c3a3f89"
+GeoInterface = "cf35fbd7-0cd7-5166-be24-54bfbe79505f"
+GeometryOps = "3251bfac-6a57-4b6d-aa61-ac1fef2975ab"
+GeometryOpsCore = "05efe853-fabf-41c8-927e-7063c8b9f013"
+GlobalRegridding = "76847d5f-3a3e-4aa5-90cc-2a3a47dd1c11"
+JSON3 = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
+RasterDataSources = "3cb90ccd-e1b6-4867-9617-4276c8b2ca36"
+Rasters = "a3a2b9e3-a471-40c9-b274-f788e487c689"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99"
+
+[sources]
+DiscreteGlobalGrids = {path = ".."}

docs/Project.toml

@@ -27,11 +27,11 @@ WGLMakie = "276b4fcb-3e11-5398-bf8b-a0c2d153d008"
 Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99"
 
 [sources]
-ConservativeRegridding = {rev = "claude/budget-frontier", url = "https://github.com/JuliaGeo/ConservativeRegridding.jl"}
+ConservativeRegridding = {rev = "6a4b997ab2e66dea45b5b62b5b2f32f0a3b279b0", url = "https://github.com/JuliaGeo/ConservativeRegridding.jl"}
 DiscreteGlobalGridVisualizations = {path = "../lib/DiscreteGlobalGridVisualizations"}
 DiscreteGlobalGrids = {path = ".."}
 DocumenterVitepress = {rev = "as/bonito-plugin", url = "https://github.com/LuxDL/DocumenterVitepress.jl.git"}
-GeometryOps = {rev = "as/intersection_area", url = "https://github.com/JuliaGeo/GeometryOps.jl.git"}
+GeometryOps = {rev = "2825c476647cfe6791dbb89fffb8592697b46995", url = "https://github.com/JuliaGeo/GeometryOps.jl.git"}
 GeometryOpsCore = {rev = "main", subdir = "GeometryOpsCore", url = "https://github.com/JuliaGeo/GeometryOps.jl.git"}
 Geomorphometry = {rev = "clipped-neighbors", url = "https://github.com/asinghvi17/GeoArrayOps.jl.git"}
 

lib/DiscreteGlobalGridsConformanceTesting/Project.toml

@@ -12,7 +12,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [sources]
 DiscreteGlobalGrids = {path = "../.."}
-GeometryOps = {url = "https://github.com/JuliaGeo/GeometryOps.jl.git", rev = "as/intersection_area"}
+GeometryOps = {url = "https://github.com/JuliaGeo/GeometryOps.jl.git", rev = "2825c476647cfe6791dbb89fffb8592697b46995"}
 GeometryOpsCore = {rev = "main", subdir = "GeometryOpsCore", url = "https://github.com/JuliaGeo/GeometryOps.jl.git"}
 
 [compat]

lib/GlobalRegridding/Project.toml

@@ -19,8 +19,8 @@ Rasters = "a3a2b9e3-a471-40c9-b274-f788e487c689"
 GlobalRegriddingRastersExt = "Rasters"
 
 [sources]
-ConservativeRegridding = {rev = "claude/budget-frontier", url = "https://github.com/JuliaGeo/ConservativeRegridding.jl"}
-GeometryOps = {rev = "as/intersection_area", url = "https://github.com/JuliaGeo/GeometryOps.jl.git"}
+ConservativeRegridding = {rev = "6a4b997ab2e66dea45b5b62b5b2f32f0a3b279b0", url = "https://github.com/JuliaGeo/ConservativeRegridding.jl"}
+GeometryOps = {rev = "2825c476647cfe6791dbb89fffb8592697b46995", url = "https://github.com/JuliaGeo/GeometryOps.jl.git"}
 GeometryOpsCore = {rev = "main", subdir = "GeometryOpsCore", url = "https://github.com/JuliaGeo/GeometryOps.jl.git"}
 
 [compat]

lib/GlobalRegridding/src/conservative.jl

@@ -40,6 +40,12 @@ function CellCapTree(space::S, inds) where {S<:RegridSpace}
     return _cellcapnode(space, ix, caps, 1, length(ix))
 end
 
+# The split is on the linear index range, not on space. Where cell positions
+# run in raster row order, a shallow node is a band of complete rows spanning
+# the source's whole longitude range, and no descent can prune on longitude
+# above about the last third of the depth. A median split of the node's cap
+# centres on their widest axis would fix it; nothing downstream reads `ix` in
+# order, so permuting `ix` and `caps` together is enough.
 function _cellcapnode(space::S, ix::Vector{Int}, caps::Vector{Cap},
     lo::Int, hi::Int) where {S}
     children = CellCapTree{S}[]
@@ -57,26 +63,33 @@ end
 Base.show(io::IO, tree::CellCapTree) =
     print(io, "CellCapTree(", tree.hi - tree.lo + 1, " cells)")
 
-# Merge caps around their mean centre. Use the full sphere beyond the convex range.
+# A cap containing all of `caps`, folded pairwise. Centring on the mean of the
+# centres instead — which is what this did — inflates a node by the spread of
+# its children and then compounds that inflation at every level above: measured
+# 4x on the root of a 1200x1200 tile, which is about two levels of lost pruning
+# depth. Use the full sphere beyond the convex range.
 function _mergecaps(caps)
-    sx = sy = sz = 0.0
-    n = 0
+    acc = _WHOLE_SPHERE
+    seen = false
     for c in caps
-        sx += c.point[1]
-        sy += c.point[2]
-        sz += c.point[3]
-        n += 1
-    end
-    n == 0 && return _WHOLE_SPHERE
-    norm = sqrt(sx^2 + sy^2 + sz^2)
-    norm <= eps(Float64) && return _WHOLE_SPHERE
-    centre = USPoint(sx / norm, sy / norm, sz / norm)
-    radius = 0.0
-    for c in caps
-        radius = max(radius, US.spherical_distance(centre, c.point) + c.radius)
+        acc = seen ? _mergecap(acc, c) : c
+        seen = true
+        acc.radius > Float64(pi) / 2 && return _WHOLE_SPHERE
     end
+    seen || return _WHOLE_SPHERE
+    return SphericalCap(acc.point, _padcap(acc.radius))
+end
+
+# The smallest cap containing both. Neither containment case leaves `d == 0`,
+# so the centre interpolation never divides by zero.
+function _mergecap(x::Cap, y::Cap)
+    d = US.spherical_distance(x.point, y.point)
+    d + y.radius <= x.radius && return x
+    d + x.radius <= y.radius && return y
+    radius = (x.radius + y.radius + d) / 2
     radius > Float64(pi) / 2 && return _WHOLE_SPHERE
-    return SphericalCap(centre, _padcap(radius))
+    return SphericalCap(US.slerp(x.point, y.point, 0.5 * (1 - (x.radius - y.radius) / d)),
+        radius)
 end
 
 _cellcap(space::RegridSpace, i::Int) = _mergecaps(

lib/GlobalRegridding/test/test_conservative.jl

@@ -281,6 +281,37 @@ cellareas(space, inds) = [GO.area(manifold(space), getcell(space, i)) for i in i
         @test all(c -> CR.Trees.split_weight(c) < CR.Trees.split_weight(tree), STI.getchild(tree))
     end
 
+    @testset "cap tree nodes bound their subtrees, and closely" begin
+        # Soundness: a node's cap must contain every cell below it, or the dual
+        # descent prunes a pair that does intersect.
+        space = ToyLonLatSpace(24, 12)
+        inside(cap, p) = US.spherical_distance(cap.point, p) <= cap.radius
+        function covers(node)
+            cap = STI.node_extent(node)
+            for i in view(node.inds, node.lo:node.hi)
+                all(inside(cap, USPoint(GI.x(p), GI.y(p), GI.z(p)))
+                    for p in GI.getpoint(getcell(space, i))) || return false
+            end
+            return all(covers, STI.getchild(node))
+        end
+        @test covers(GR.CellCapTree(space, 1:ncells(space)))
+
+        # Tightness: any cap covering these cells has a radius of at least half
+        # the widest distance between two of their vertices, and Jung's bound
+        # puts the smallest one below 1.16 of that. Centring on the mean of the
+        # child centres instead — which is what this did — inflated a node by
+        # its children's spread and compounded that at every level above,
+        # measured at 3-4x on a Copernicus tile, or two levels of lost pruning.
+        for inds in (1:6, 25:30, [cellposition(space, ix, iy) for ix in 3:6, iy in 4:7])
+            points = [USPoint(GI.x(p), GI.y(p), GI.z(p))
+                      for i in vec(collect(inds)) for p in GI.getpoint(getcell(space, i))]
+            widest = maximum(US.spherical_distance(a, b) for a in points, b in points)
+            root = GR.CellCapTree(space, vec(collect(inds)))
+            @test root.extent.radius >= widest / 2
+            @test root.extent.radius < 1.3 * widest / 2
+        end
+    end
+
     @testset "one tile's restricted tree, built once" begin
         # A non-chunk range on a fallback space forces the CellCapTree path.
         xd = DD.X(-168.75:22.5:168.75)

src/cap_cached_tree.jl

@@ -1,18 +1,26 @@
 # Whole-space destination tree with per-cell work done once, not once per visit.
 
 """
-    CapCachedTree(node, caps)
+    CapCachedTree(node, caps, offset = 0)
 
-Wrap a grid cursor so leaf extents read a position-indexed cap vector computed
-once per build. The dual-tree search otherwise re-derives each leaf's cap — an
-inverse projection over its boundary — once per opposing leaf, and then a
-second time in `child_indices_extents` for the same visit.
+Wrap a grid cursor so leaf extents read a cap vector computed once per build:
+grid position `p` is `caps[p - offset]`. The dual-tree search otherwise
+re-derives each leaf's cap — an inverse projection over its boundary — once per
+opposing leaf, and then a second time in `child_indices_extents` for the same
+visit.
+
+`offset` is `0` for a whole-space tree and `first(inds) - 1` for a chunk's,
+whose cap vector covers only that chunk's positions.
 """
 struct CapCachedTree{C<:HierarchicalGridCursor,E}
     node::C
     caps::Vector{E}
+    offset::Int
 end
 
+CapCachedTree(node::HierarchicalGridCursor, caps::Vector) =
+    CapCachedTree(node, caps, 0)
+
 Base.show(io::IO, t::CapCachedTree) =
     print(io, "CapCachedTree(", t.node, ", ", length(t.caps), " caps)")
 
@@ -22,15 +30,16 @@ STI.node_extent_is_expensive(::Type{<:CapCachedTree{C}}) where {C} =
 STI.isleaf(t::CapCachedTree) = STI.isleaf(t.node)
 STI.nchild(t::CapCachedTree) = STI.nchild(t.node)
 STI.getchild(t::CapCachedTree) =
-    Iterators.map(Base.Fix2(CapCachedTree, t.caps), STI.getchild(t.node))
-STI.getchild(t::CapCachedTree, i::Int) = CapCachedTree(STI.getchild(t.node, i), t.caps)
+    Iterators.map(n -> CapCachedTree(n, t.caps, t.offset), STI.getchild(t.node))
+STI.getchild(t::CapCachedTree, i::Int) =
+    CapCachedTree(STI.getchild(t.node, i), t.caps, t.offset)
 
 # A window node at the leaf level is exactly one stored cell; its position
 # indexes the cache. Everything else keeps the cursor's own extent logic.
 function STI.node_extent(t::CapCachedTree)
     c = t.node
     if !Engine._issynthetic(c) && c.level >= c.leaf_level
-        return @inbounds t.caps[c.first_index]
+        return @inbounds t.caps[c.first_index - t.offset]
     end
     return STI.node_extent(c)
 end
@@ -43,7 +52,7 @@ function STI.child_indices_extents(t::CapCachedTree{<:Any,E}) where {E}
     entries = Vector{Tuple{Int,E}}(undef, count)
     for k in 1:count
         index = Engine._stored_index(c, k)
-        entries[k] = (index, @inbounds t.caps[index])
+        entries[k] = (index, @inbounds t.caps[index - t.offset])
     end
     return entries
 end
@@ -59,7 +68,21 @@ function _cachedcelltree(space::DGGSpace)
     root = treeify(_decodedgrid(space.grid))
     (root isa HierarchicalGridCursor && root.selection === nothing) ||
         return GR.celltree(space)
-    return CapCachedTree(root, _leafcaps(root.grid))
+    return CapCachedTree(root, _leafcaps(root.grid, 1:ncells(root.grid)))
+end
+
+# Above this a chunk's cap vector costs more to fill than the revisits it saves.
+# A chunk at the limit holds 2 MiB of caps.
+const _CHUNK_CAP_CACHE_MAX = 1 << 16
+
+# A chunk's cursor with its own cap vector. `subtree` calls this once per block
+# build, and the destination side reaches it through `GR.TileCells`, which keeps
+# one tree per tile — so the caps are filled once and read by every source chunk
+# paired with that tile.
+function _cachedchunktree(cursor::HierarchicalGridCursor,
+        inds::AbstractUnitRange{<:Integer})
+    length(inds) > _CHUNK_CAP_CACHE_MAX && return cursor
+    return CapCachedTree(cursor, _leafcaps(cursor.grid, inds), Int(first(inds)) - 1)
 end
 
 # Decode a compressed id vector once so descent and geometry read O(1) ids.
@@ -69,32 +92,34 @@ _decodedgrid(grid::PartialGrid{<:AbstractHierarchicalGridSystem,<:CellVector}) =
         root = Engine._is_rooted(grid) ? grid.root_id : nothing)
 _decodedgrid(grid::AbstractGrid) = grid
 
-# One tight cap per cell, in parallel at top level ("outer parallelism wins").
-function _leafcaps(grid::AbstractGrid)
-    n = ncells(grid)
-    n == 0 && return [_cellcap(grid, i) for i in 1:n]
-    caps = Vector{typeof(_cellcap(grid, 1))}(undef, n)
+# One tight cap per cell of `inds`, in parallel at top level ("outer
+# parallelism wins"). Entry `k` holds the cap of position `first(inds) + k - 1`.
+function _leafcaps(grid::AbstractGrid, inds::AbstractUnitRange{<:Integer})
+    n = length(inds)
+    n == 0 && return [_cellcap(grid, Int(i)) for i in inds]
+    off = Int(first(inds)) - 1
+    caps = Vector{typeof(_cellcap(grid, off + 1))}(undef, n)
     nt = GR._innerthreaded() isa GOCore.True ?
         min(Threads.nthreads(), max(1, n >> 14)) : 1
     if nt > 1
         tasks = Vector{Task}(undef, nt)
         for t in 1:nt
             lo = (n * (t - 1)) ÷ nt + 1
             hi = (n * t) ÷ nt
-            tasks[t] = Threads.@spawn _fillcaps!(caps, grid, lo, hi)
+            tasks[t] = Threads.@spawn _fillcaps!(caps, grid, off, lo, hi)
         end
         foreach(wait, tasks)
     else
-        _fillcaps!(caps, grid, 1, n)
+        _fillcaps!(caps, grid, off, 1, n)
     end
     return caps
 end
 
 _cellcap(grid::AbstractGrid, i::Int) = Fallbacks.cell_cap(grid, cellindex(grid, i))
 
-function _fillcaps!(caps::Vector, grid::AbstractGrid, lo::Int, hi::Int)
-    for i in lo:hi
-        @inbounds caps[i] = _cellcap(grid, i)
+function _fillcaps!(caps::Vector, grid::AbstractGrid, off::Int, lo::Int, hi::Int)
+    for k in lo:hi
+        @inbounds caps[k] = _cellcap(grid, off + k)
     end
     return nothing
 end

src/interface/grid.jl

@@ -554,6 +554,25 @@ The one-argument form picks the manifold with `best_manifold(grid)`.
 """
 function treeify end
 
+"""
+    subcursor(grid::AbstractGrid, inds::AbstractUnitRange) -> tree or `nothing`
+
+The [`treeify`](@ref) tree restricted to the grid positions `inds`, with leaf
+indices still in `grid`'s own position space, or `nothing` when this grid cannot
+express that restriction.
+
+The default is `nothing`. A grid whose tree is a lattice rectangle rather than a
+cell hierarchy — a raster-backed system's — implements this so that a chunk of
+it keeps a prunable tree instead of the regridder's bounding-cap fallback. A
+hierarchical grid needs no method: regridding descends its own hierarchy to the
+chunk's ancestor.
+
+The result must cover exactly the cells at `inds`, no more.
+"""
+function subcursor end
+
+subcursor(::AbstractGrid, ::AbstractUnitRange{<:Integer}) = nothing
+
 # ===========================================================================
 # Queries
 # ===========================================================================

src/regridding.jl

@@ -148,14 +148,19 @@ GR.chunkranges(space::DGGSpace, chunk::Integer, ::NTuple{1,Int}) =
     GlobalRegridding.subtree(space::DGGSpace, inds)
 
 Return the cell tree restricted to `inds`, preserving global cell positions.
-The whole space gets a cursor with decoded ids and precomputed leaf caps;
-exact chunk ranges reuse the grid hierarchy in `O(1)`; other ranges use a
-bounding-cap tree.
+The whole space gets a cursor with decoded ids and precomputed leaf caps; a
+grid that can window its own tree ([`subcursor`](@ref)) does so; exact chunk
+ranges reuse the grid hierarchy in `O(1)`, with the same cap precomputation
+where the chunk is small enough to pay for it; other ranges use a bounding-cap
+tree.
 """
 function GR.subtree(space::DGGSpace, inds::AbstractUnitRange{<:Integer})
     GR._iswholespace(space, inds) && return _cachedcelltree(space)
+    # Dispatch-only, so it costs nothing for the grids that have no method.
+    window = subcursor(space.grid, inds)
+    window === nothing || return window
     cursor = _chunkcursor(space, inds)
-    cursor === nothing || return cursor
+    cursor === nothing || return _cachedchunktree(cursor, inds)
     return GR.CellCapTree(space, inds)
 end