Fix reconstruct_global_grid (#4157)

Author: glwagner

src/DistributedComputations/distributed_grids.jl

@@ -16,8 +16,21 @@ const DistributedGrid{FT, TX, TY, TZ} = AbstractGrid{FT, TX, TY, TZ, <:Distribut
 const DistributedRectilinearGrid{FT, TX, TY, TZ, CZ, FX, FY, VX, VY} =
     RectilinearGrid{FT, TX, TY, TZ, CZ, FX, FY, VX, VY, <:Distributed} where {FT, TX, TY, TZ, CZ, FX, FY, VX, VY}
 
-const DistributedLatitudeLongitudeGrid{FT, TX, TY, TZ, CZ, M, MY, FX, FY, VX, VY} = 
-    LatitudeLongitudeGrid{FT, TX, TY, TZ, CZ, M, MY, FX, FY, VX, VY, <:Distributed} where {FT, TX, TY, TZ, CZ, M, MY, FX, FY, VX, VY}
+const DistributedLatitudeLongitudeGrid{FT, TX, TY, TZ, Z,
+                                       DXF, DXC, XF, XC,
+                                       DYF, DYC, YF, YC,
+                                       DXFC, DXCF, DXFF,
+                                       DXCC, DYFC, DYCF} =
+    LatitudeLongitudeGrid{FT, TX, TY, TZ, Z,
+                          DXF, DXC, XF, XC,
+                          DYF, DYC, YF, YC,
+                          DXFC, DXCF, DXFF,
+                          DXCC, DYFC, DYCF,
+                          <:Distributed} where {FT, TX, TY, TZ, Z,
+                                                DXF, DXC, XF, XC,
+                                                DYF, DYC, YF, YC,
+                                                DXFC, DXCF, DXFF,
+                                                DXCC, DYFC, DYCF}
 
 # Local size from global size and architecture
 local_size(arch::Distributed, global_sz) = (local_size(global_sz[1], arch.partition.x, arch.local_index[1]),

src/Grids/input_validation.jl

@@ -84,7 +84,11 @@ function validate_halo(TX, TY, TZ, size, halo)
     halo = inflate_tuple(TX, TY, TZ, halo, default=0)
 
     for i in 1:2
-        !(halo[i] ≤ size[i]) && throw(ArgumentError("halo must be ≤ size for coordinate $(coordinate_name(i))"))
+        H = halo[i]
+        N = size[i]
+        if !(H ≤ N)
+            throw(ArgumentError("halo=$H must be ≤ size=$N for coordinate $(coordinate_name(i))"))
+        end
     end
 
     return halo

test/test_distributed_hydrostatic_model.jl

@@ -7,7 +7,7 @@ using MPI
 # These tests are meant to be run on 4 ranks. This script may be run
 # stand-alone (outside the test environment) via
 #
-# mpiexec -n 4 julia --project test_distributed_models.jl
+# $ MPI_TEST=true mpiexec -n 4 julia --project test_distributed_hydrostatic_model.jl
 #
 # provided that a few packages (like TimesDates.jl) are in your global environment.
 #
@@ -17,14 +17,10 @@ using MPI
 #
 # then later:
 # 
-# julia> include("test_distributed_models.jl")
-#
-# When running the tests this way, uncomment the following line
+# julia> include("test_distributed_hydrostatic_model.jl")
 
 MPI.Initialized() || MPI.Init()
 
-# to initialize MPI.
-
 using Oceananigans.Operators: hack_cosd
 using Oceananigans.DistributedComputations: ranks, partition, all_reduce, cpu_architecture, reconstruct_global_grid, synchronized
 using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity
@@ -86,7 +82,7 @@ Nx = 32
 Ny = 32 
 
 for arch in archs
-    
+
     # We do not test on `Fractional` partitions where we cannot easily ensure that H ≤ N 
     # which would lead to different advection schemes for partitioned and non-partitioned grids.
     # `Fractional` is, however, tested in regression tests where the horizontal dimensions are larger.
@@ -96,23 +92,25 @@ for arch in archs
 
     if valid_x_partition & valid_y_partition & valid_z_partition
         @testset "Testing distributed solid body rotation" begin
-            underlying_grid = LatitudeLongitudeGrid(arch, size = (Nx, Ny, 3),
+            underlying_grid = LatitudeLongitudeGrid(arch,
+                                                    size = (Nx, Ny, 3),
                                                     halo = (4, 4, 3),
                                                     latitude = (-80, 80),
                                                     longitude = (-160, 160),
                                                     z = (-1, 0),
                                                     radius = 1,
-                                                    topology=(Bounded, Bounded, Bounded))
+                                                    topology = (Bounded, Bounded, Bounded))
 
             bottom(λ, φ) = -30 < λ < 30 && -40 < φ < 20 ? 0 : - 1
 
             immersed_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom))
-            immersed_active_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom); active_cells_map = true)
+            immersed_active_grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom); active_cells_map=true)
 
             global_underlying_grid = reconstruct_global_grid(underlying_grid)
             global_immersed_grid   = ImmersedBoundaryGrid(global_underlying_grid, GridFittedBottom(bottom))
 
-            for (grid, global_grid) in zip((underlying_grid, immersed_grid, immersed_active_grid), (global_underlying_grid, global_immersed_grid, global_immersed_grid))
+            for (grid, global_grid) in zip((underlying_grid, immersed_grid, immersed_active_grid),
+                                           (global_underlying_grid, global_immersed_grid, global_immersed_grid))
                 if arch.local_rank == 0
                     @info "  Testing distributed solid body rotation with $(ranks(arch)) ranks on $(typeof(grid).name.wrapper)"
                 end
@@ -190,3 +188,4 @@ for arch in archs
         end
     end
 end
+

test/test_distributed_models.jl

@@ -18,13 +18,9 @@ using MPI
 # then later:
 # 
 # julia> include("test_distributed_models.jl")
-#
-# When running the tests this way, uncomment the following line
 
 MPI.Init()
 
-# to initialize MPI.
-
 using Oceananigans.BoundaryConditions: fill_halo_regions!, DCBC
 using Oceananigans.DistributedComputations: Distributed, index2rank
 using Oceananigans.Fields: AbstractField