Use TiledIteration.TileIterator in RegularCoarsening (#1123)

* Use 'TiledIteration.TileIterator' in 'RegularCoarsening'

* Add tests

* Update comments

* Update src/coarsening/regular.jl

Co-authored-by: Júlio Hoffimann <[email protected]>

* Add more tests

* Apply suggestions

---------

Co-authored-by: Júlio Hoffimann <[email protected]>

Author: 2 people

Project.toml

@@ -18,6 +18,7 @@ ScopedValues = "7e506255-f358-4e82-b7e4-beb19740aa63"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+TiledIteration = "06e1c1a7-607b-532d-9fad-de7d9aa2abac"
 TransformsBase = "28dd2a49-a57a-4bfb-84ca-1a49db9b96b8"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
@@ -37,6 +38,7 @@ ScopedValues = "1.2"
 SparseArrays = "1.9"
 StaticArrays = "1.0"
 StatsBase = "0.33, 0.34"
+TiledIteration = "0.5"
 TransformsBase = "1.6"
 Unitful = "1.17"
 julia = "1.9"

src/Meshes.jl

@@ -20,6 +20,7 @@ using Distances: Haversine, SphericalAngle
 using Distances: evaluate, result_type
 using Rotations: Rotation, QuatRotation, Angle2d
 using Rotations: rotation_between
+using TiledIteration: TileIterator
 using CoordRefSystems: Basic, Projected, Geographic
 using NearestNeighbors: KDTree, BallTree
 using NearestNeighbors: knn, inrange

src/coarsening/regular.jl

@@ -22,26 +22,43 @@ RegularCoarsening(factors::Vararg{Int,N}) where {N} = RegularCoarsening(factors)
 
 function coarsen(grid::OrthoRegularGrid, method::RegularCoarsening)
   factors = fitdims(method.factors, paramdim(grid))
-  RegularGrid(minimum(grid), maximum(grid), dims=size(grid) .÷ factors)
+  dims = _coarsesize(grid, factors)
+  RegularGrid(minimum(grid), maximum(grid), dims=dims)
 end
 
 function coarsen(grid::RectilinearGrid, method::RegularCoarsening)
   factors = fitdims(method.factors, paramdim(grid))
-  dims = vsize(grid)
-  rngs = ntuple(i -> 1:factors[i]:dims[i], paramdim(grid))
+  inds = _coarseinds(grid, factors)
   xyzₛ = xyz(grid)
-  xyzₜ = ntuple(i -> xyzₛ[i][rngs[i]], paramdim(grid))
+  xyzₜ = ntuple(i -> xyzₛ[i][inds[i]], paramdim(grid))
   RectilinearGrid{manifold(grid),crs(grid)}(xyzₜ)
 end
 
 function coarsen(grid::StructuredGrid, method::RegularCoarsening)
   factors = fitdims(method.factors, paramdim(grid))
-  dims = vsize(grid)
-  rngs = ntuple(i -> 1:factors[i]:dims[i], paramdim(grid))
+  inds = _coarseinds(grid, factors)
   XYZₛ = XYZ(grid)
-  XYZₜ = ntuple(i -> XYZₛ[i][rngs...], paramdim(grid))
+  XYZₜ = ntuple(i -> XYZₛ[i][inds...], paramdim(grid))
   StructuredGrid{manifold(grid),crs(grid)}(XYZₜ)
 end
 
 coarsen(grid::TransformedGrid, method::RegularCoarsening) =
   TransformedGrid(coarsen(parent(grid), method), transform(grid))
+
+# -----------------
+# HELPER FUNCTIONS
+# -----------------
+
+function _coarsesize(grid, factors)
+  dims = size(grid)
+  axes = ntuple(i -> 1:dims[i], paramdim(grid))
+  size(TileIterator(axes, factors))
+end
+
+_coarsevsize(grid, factors) = _coarsesize(grid, factors) .+ .!isperiodic(grid)
+
+function _coarseinds(grid, factors)
+  dims = vsize(grid)
+  tdims = _coarsevsize(grid, factors)
+  ntuple(i -> floor.(Int, range(start=1, stop=dims[i], length=tdims[i])), paramdim(grid))
+end

test/coarsening.jl

@@ -16,6 +16,22 @@
   tgrid = CartesianGrid(cart(0.0, 0.0), cart(10.0, 10.0), dims=(10, 5))
   @test coarsen(grid, RegularCoarsening(2, 4)) == tgrid
 
+  # non-multiple dimensions
+  grid = CartesianGrid(cart(0, 0), cart(13, 17), dims=(13, 17))
+  tgrid = CartesianGrid(cart(0, 0), cart(13, 17), dims=(3, 6))
+  @test coarsen(grid, RegularCoarsening(5, 3)) == tgrid
+  rgrid = convert(RectilinearGrid, grid)
+  @test size(coarsen(rgrid, RegularCoarsening(5, 3))) == (3, 6)
+  sgrid = convert(StructuredGrid, grid)
+  @test size(coarsen(sgrid, RegularCoarsening(5, 3))) == (3, 6)
+  tfgrid = TransformedGrid(grid, Identity())
+  @test size(coarsen(tfgrid, RegularCoarsening(5, 3))) == (3, 6)
+
+  # large grid
+  grid = CartesianGrid(cart(0, 0), cart(16200, 8100), dims=(16200, 8100))
+  tgrid = CartesianGrid(cart(0, 0), cart(16200, 8100), dims=(203, 203))
+  @test coarsen(grid, RegularCoarsening(80, 40)) == tgrid
+
   # 3D grids
   grid = cartgrid(100, 100, 100)
   tgrid = CartesianGrid(minimum(grid), maximum(grid), dims=(50, 25, 20))
@@ -28,4 +44,15 @@
   @test coarsen(sgrid, RegularCoarsening(2, 4, 5)) == tsgrid
   tfgrid = TransformedGrid(grid, Identity())
   @test coarsen(tfgrid, RegularCoarsening(2, 4, 5)) == coarsen(grid, RegularCoarsening(2, 4, 5))
+
+  # non-multiple dimensions
+  grid = CartesianGrid(cart(0, 0, 0), cart(13, 17, 23), dims=(13, 17, 23))
+  tgrid = CartesianGrid(cart(0, 0, 0), cart(13, 17, 23), dims=(2, 4, 8))
+  @test coarsen(grid, RegularCoarsening(7, 5, 3)) == tgrid
+  rgrid = convert(RectilinearGrid, grid)
+  @test size(coarsen(rgrid, RegularCoarsening(7, 5, 3))) == (2, 4, 8)
+  sgrid = convert(StructuredGrid, grid)
+  @test size(coarsen(sgrid, RegularCoarsening(7, 5, 3))) == (2, 4, 8)
+  tfgrid = TransformedGrid(grid, Identity())
+  @test size(coarsen(tfgrid, RegularCoarsening(7, 5, 3))) == (2, 4, 8)
 end