Throw error if nodes given by function/array generators are not sorted (

#3331) * Check that coordinates are sorted * Fix error for Flat * Fix distributed grids --------- Co-authored-by: Navid C. Constantinou <[email protected]>
CliMA · Oct 12, 2023 · 8dddf6f · 8dddf6f
1 parent dcbdc04
commit 8dddf6f
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 38 deletions.
diff --git a/src/DistributedComputations/distributed_grids.jl b/src/DistributedComputations/distributed_grids.jl
@@ -70,9 +70,9 @@ function RectilinearGrid(arch::Distributed,
     yl = partition(y, ny, arch, 2)
     zl = partition(z, nz, arch, 3)
 
-    Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, topology[1](), nx, Hx, xl, child_architecture(arch))
-    Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, topology[2](), ny, Hy, yl, child_architecture(arch))
-    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topology[3](), nz, Hz, zl, child_architecture(arch))
+    Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, topology[1](), nx, Hx, xl, :x, child_architecture(arch))
+    Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, topology[2](), ny, Hy, yl, :y, child_architecture(arch))
+    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topology[3](), nz, Hz, zl, :z, child_architecture(arch))
 
     return RectilinearGrid{TX, TY, TZ}(arch,
                                        nx, ny, nz,
@@ -117,15 +117,16 @@ function LatitudeLongitudeGrid(arch::Distributed,
     # Calculate all direction (which might be stretched)
     # A direction is regular if the domain passed is a Tuple{<:Real, <:Real}, 
     # it is stretched if being passed is a function or vector (as for the VerticallyStretchedRectilinearGrid)
-    Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), nλ, Hλ, λl, arch.child_architecture)
-    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), nz, Hz, zl, arch.child_architecture)
-    # The Latitudinal direction is _special_ :
-    # Preconmpute metrics assumes that `length(φᵃᶠᵃ) = length(φᵃᶜᵃ) + 1`, which is always the case in a 
+    Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), nλ, Hλ, λl, :longitude, arch.child_architecture)
+    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), nz, Hz, zl, :z,         arch.child_architecture)
+
+    # The Latitudinal direction is _special_:
+    # precompute_metrics assumes that `length(φᵃᶠᵃ) = length(φᵃᶜᵃ) + 1`, which is always the case in a 
     # serial grid because `LatitudeLongitudeGrid` should be always `Bounded`, but it is not true for a
     # partitioned `DistributedGrid` with Ry > 1 (one rank will hold a `RightConnected` topology)
     # But we need an extra point to precompute the Y direction in case of only one halo so we disregard the topology
     # when constructing the metrics!
-    Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, Bounded(), nφ, Hφ, φl, arch.child_architecture)
+    Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, Bounded(), nφ, Hφ, φl, :latitude, arch.child_architecture)
 
     preliminary_grid = LatitudeLongitudeGrid{TX, TY, TZ}(arch,
                                                          nλ, nφ, nz,
@@ -174,9 +175,9 @@ function reconstruct_global_grid(grid::DistributedRectilinearGrid)
 
     FT = eltype(grid)
 
-    Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, TX(), Nx, Hx, xG, child_arch)
-    Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, TY(), Ny, Hy, yG, child_arch)
-    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, zG, child_arch)
+    Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, TX(), Nx, Hx, xG, :x, child_arch)
+    Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, TY(), Ny, Hy, yG, :y, child_arch)
+    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, zG, :z, child_arch)
 
     return RectilinearGrid{TX, TY, TZ}(child_arch,
                                        Nx, Ny, Nz,
@@ -220,9 +221,9 @@ function reconstruct_global_grid(grid::DistributedLatitudeLongitudeGrid)
     # Calculate all direction (which might be stretched)
     # A direction is regular if the domain passed is a Tuple{<:Real, <:Real}, 
     # it is stretched if being passed is a function or vector (as for the VerticallyStretchedRectilinearGrid)
-    Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), Nλ, Hλ, λG, child_arch)
-    Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, TY(), Nφ, Hφ, φG, child_arch)
-    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, zG, child_arch)
+    Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), Nλ, Hλ, λG, :longitude, child_arch)
+    Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, TY(), Nφ, Hφ, φG, :latitude,  child_arch)
+    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, zG, :z,         child_arch)
 
     precompute_metrics = metrics_precomputed(grid)
 

diff --git a/src/Grids/grid_generation.jl b/src/Grids/grid_generation.jl
@@ -8,9 +8,9 @@ get_domain_extent(coord, N)                 = (coord[1], coord[2])
 get_domain_extent(coord::Function, N)       = (coord(1), coord(N+1))
 get_domain_extent(coord::AbstractVector, N) = CUDA.@allowscalar (coord[1], coord[N+1])
 
-get_coord_face(coord::Nothing, i) = 1
-get_coord_face(coord::Function, i) = coord(i)
-get_coord_face(coord::AbstractVector, i) = CUDA.@allowscalar coord[i]
+get_face_node(coord::Nothing, i) = 1
+get_face_node(coord::Function, i) = coord(i)
+get_face_node(coord::AbstractVector, i) = CUDA.@allowscalar coord[i]
 
 const AT = AbstractTopology
 lower_exterior_Δcoordᶠ(::AT,              Fi, Hcoord) = [Fi[end - Hcoord + i] - Fi[end - Hcoord + i - 1] for i = 1:Hcoord]
@@ -25,28 +25,39 @@ upper_interior_F(::BoundedTopology, coord) = coord
 total_interior_length(::AT, N)              = N
 total_interior_length(::BoundedTopology, N) = N + 1
 
+bad_coordinate_message(ξ::Function, name) = "The values of $name(index) must increase as the index increases!"
+bad_coordinate_message(ξ::AbstractArray, name) = "The elements of $name must be increasing!"
+
 # generate a variably-spaced coordinate passing the explicit coord faces as vector or function
-function generate_coordinate(FT, topo::AT, N, H, coord, arch)
+function generate_coordinate(FT, topo::AT, N, H, node_generator, coordinate_name, arch)
 
     # Ensure correct type for F and derived quantities
-    interiorF = zeros(FT, N+1)
+    interior_face_nodes = zeros(FT, N+1)
+
+    # Use the user-supplied "generator" to build the interior nodes
+    for idx = 1:N+1
+        interior_face_nodes[idx] = get_face_node(node_generator, idx)
+    end
 
-    for i = 1:N+1
-        interiorF[i] = get_coord_face(coord, i)
+    # Check that the interior nodes are increasing
+    if !issorted(interior_face_nodes)
+        msg = bad_coordinate_message(node_generator, coordinate_name)
+        throw(ArgumentError(msg))
     end
 
-    L = interiorF[N+1] - interiorF[1]
+    # Get domain extent
+    L = interior_face_nodes[N+1] - interior_face_nodes[1]
 
-    # Build halo regions
-    Δᶠ₋ = lower_exterior_Δcoordᶠ(topo, interiorF, H)
-    Δᶠ₊ = reverse(upper_exterior_Δcoordᶠ(topo, interiorF, H))
+    # Build halo regions: spacings first
+    Δᶠ₋ = lower_exterior_Δcoordᶠ(topo, interior_face_nodes, H)
+    Δᶠ₊ = reverse(upper_exterior_Δcoordᶠ(topo, interior_face_nodes, H))
 
-    c¹, cᴺ⁺¹ = interiorF[1], interiorF[N+1]
+    c¹, cᴺ⁺¹ = interior_face_nodes[1], interior_face_nodes[N+1]
 
     F₋ =         [c¹   - sum(Δᶠ₋[i:H]) for i = 1:H]  # locations of faces in lower halo
     F₊ = reverse([cᴺ⁺¹ + sum(Δᶠ₊[i:H]) for i = 1:H]) # locations of faces in top halo
 
-    F = vcat(F₋, interiorF, F₊)
+    F = vcat(F₋, interior_face_nodes, F₊)
 
     # Build cell centers, cell center spacings, and cell interface spacings
     TC = total_length(Center(), topo, N, H)
@@ -77,12 +88,15 @@ function generate_coordinate(FT, topo::AT, N, H, coord, arch)
     return L, F, C, Δᶠ, Δᶜ
 end
 
-# generate a regularly-spaced coordinate passing the domain extent (2-tuple) and number of points
-function generate_coordinate(FT, topo::AT, N, H, coord::Tuple{<:Number, <:Number}, arch)
+# Generate a regularly-spaced coordinate passing the domain extent (2-tuple) and number of points
+function generate_coordinate(FT, topo::AT, N, H, node_interval::Tuple{<:Number, <:Number}, coordinate_name, arch)
 
-    @assert length(coord) == 2
+    if node_interval[2] < node_interval[1]
+        msg = "$coordinate_name must be an increasing interval!"
+        throw(ArgumentError(msg))
+    end
 
-    c₁, c₂ = @. BigFloat(coord)
+    c₁, c₂ = @. BigFloat(node_interval)
     @assert c₁ < c₂
     L = c₂ - c₁
 
@@ -108,5 +122,5 @@ function generate_coordinate(FT, topo::AT, N, H, coord::Tuple{<:Number, <:Number
 end
 
 # Flat domains
-generate_coordinate(FT, ::Flat, N, H, coord::Tuple{<:Number, <:Number}, arch) =
+generate_coordinate(FT, ::Flat, N, H, coord::Tuple{<:Number, <:Number}, coordinate_name, arch) =
     FT(1), range(1, 1, length=N), range(1, 1, length=N), FT(1), FT(1)
diff --git a/src/Grids/latitude_longitude_grid.jl b/src/Grids/latitude_longitude_grid.jl
@@ -194,9 +194,9 @@ function LatitudeLongitudeGrid(architecture::AbstractArchitecture = CPU(),
     # it is stretched if being passed is a function or vector (as for the VerticallyStretchedRectilinearGrid)
     TX, TY, TZ = topology
 
-    Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), Nλ, Hλ, longitude, architecture)
-    Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, TY(), Nφ, Hφ, latitude,  architecture)
-    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, z,         architecture)
+    Lλ, λᶠᵃᵃ, λᶜᵃᵃ, Δλᶠᵃᵃ, Δλᶜᵃᵃ = generate_coordinate(FT, TX(), Nλ, Hλ, longitude, :longitude, architecture)
+    Lφ, φᵃᶠᵃ, φᵃᶜᵃ, Δφᵃᶠᵃ, Δφᵃᶜᵃ = generate_coordinate(FT, TY(), Nφ, Hφ, latitude,  :latitude,  architecture)
+    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, TZ(), Nz, Hz, z,         :z,         architecture)
 
     preliminary_grid = LatitudeLongitudeGrid{TX, TY, TZ}(architecture,
                                                          Nλ, Nφ, Nz,

diff --git a/src/Grids/rectilinear_grid.jl b/src/Grids/rectilinear_grid.jl
@@ -270,9 +270,9 @@ function RectilinearGrid(architecture::AbstractArchitecture = CPU(),
     Nx, Ny, Nz = size
     Hx, Hy, Hz = halo
 
-    Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, topology[1](), Nx, Hx, x, architecture)
-    Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, topology[2](), Ny, Hy, y, architecture)
-    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topology[3](), Nz, Hz, z, architecture)
+    Lx, xᶠᵃᵃ, xᶜᵃᵃ, Δxᶠᵃᵃ, Δxᶜᵃᵃ = generate_coordinate(FT, topology[1](), Nx, Hx, x, :x, architecture)
+    Ly, yᵃᶠᵃ, yᵃᶜᵃ, Δyᵃᶠᵃ, Δyᵃᶜᵃ = generate_coordinate(FT, topology[2](), Ny, Hy, y, :y, architecture)
+    Lz, zᵃᵃᶠ, zᵃᵃᶜ, Δzᵃᵃᶠ, Δzᵃᵃᶜ = generate_coordinate(FT, topology[3](), Nz, Hz, z, :z, architecture)
 
     return RectilinearGrid{TX, TY, TZ}(architecture,
                                        Nx, Ny, Nz,