Merge a0fcd2b into b5dc705

src/OrthogonalSphericalShellGrids.jl

@@ -12,7 +12,8 @@ using Oceananigans.Grids: R_Earth,
                           halo_size, spherical_area_quadrilateral,
                           lat_lon_to_cartesian, generate_coordinate, topology
 using Oceananigans.Operators
-using Oceananigans.Utils: get_cartesian_nodes_and_vertices
+
+using Oceananigans.Utils: get_cartesian_nodes_and_vertices                       
 
 using Adapt 
 using JLD2
@@ -28,6 +29,8 @@ include("generate_tripolar_coordinates.jl")
 include("tripolar_grid.jl")
 include("grid_extensions.jl")
 include("distributed_tripolar_grid.jl")
+include("distributed_zipper.jl")
+include("distributed_zipper_north_tags.jl")
 include("with_halo.jl")
 include("split_explicit_free_surface.jl")
 

src/distributed_zipper.jl

@@ -0,0 +1,111 @@
+using Oceananigans.BoundaryConditions: fill_open_boundary_regions!, 
+                                       permute_boundary_conditions, 
+                                       fill_halo_event!,
+                                       DistributedCommunication
+
+using Oceananigans.DistributedComputations: cooperative_waitall!,
+                                            recv_from_buffers!,
+                                            fill_corners!,
+                                            loc_id, 
+                                            DCBCT
+
+import Oceananigans.BoundaryConditions: fill_halo_regions!
+import Oceananigans.DistributedComputations: synchronize_communication!
+
+@inline instantiate(T::DataType) = T()
+@inline instantiate(T) = T
+
+const DistributedZipper = BoundaryCondition{<:DistributedCommunication, <:ZipperHaloCommunicationRanks}
+
+switch_north_halos!(c, north_bc, grid, loc) = nothing
+
+function switch_north_halos!(c, north_bc::DistributedZipper, grid, loc) 
+    sign  = north_bc.condition.sign
+    hz = halo_size(grid)
+    sz = size(grid)
+
+    _switch_north_halos!(parent(c), loc, sign, sz, hz)
+
+    return nothing
+end
+
+@inline reversed_halos(::Tuple{<:Any, <:Center, <:Any}, Ny, Hy) = Ny+2Hy:-1:Ny+Hy+2
+@inline reversed_halos(::Tuple{<:Any, <:Face,   <:Any}, Ny, Hy) = Ny+2Hy-1:-1:Ny+Hy+1
+
+@inline west_corner_halos(::Tuple{<:Face,   <:Any, <:Any}, Hx) = 2:Hx 
+@inline west_corner_halos(::Tuple{<:Center, <:Any, <:Any}, Hx) = 1:Hx
+
+# We throw away the first point!
+@inline function _switch_north_halos!(c, loc, sign, (Nx, Ny, Nz), (Hx, Hy, Hz)) 
+
+    # Domain indices common for all locations
+    north_halos = Ny+Hy+1:Ny+2Hy-1
+    east_corner = Nx+Hx+1:Nx+2Hx
+    interior    = Hx+1:Nx+Hx
+
+    # Location - dependent halo indices 
+    reversed_north_halos = reversed_halos(loc, Ny, Hy)
+    west_corner = west_corner_halos(loc, Hx)
+
+    view(c, west_corner, north_halos, :) .= sign .* reverse(view(c, west_corner, reversed_north_halos, :), dims = 1) 
+    view(c, east_corner, north_halos, :) .= sign .* reverse(view(c, east_corner, reversed_north_halos, :), dims = 1) 
+    view(c, interior,    north_halos, :) .= sign .* reverse(view(c, interior,    reversed_north_halos, :), dims = 1) 
+
+    return nothing
+end
+
+function fill_halo_regions!(c::OffsetArray, bcs, indices, loc, grid::DTRG, buffers, args...; only_local_halos = false, fill_boundary_normal_velocities = true, kwargs...)
+    if fill_boundary_normal_velocities
+        fill_open_boundary_regions!(c, bcs, indices, loc, grid, args...; kwargs...)
+    end
+
+    north_bc = bcs.north
+
+    arch = architecture(grid)
+    fill_halos!, bcs = permute_boundary_conditions(bcs) 
+
+    number_of_tasks = length(fill_halos!)
+
+    for task = 1:number_of_tasks
+        fill_halo_event!(c, fill_halos![task], bcs[task], indices, loc, arch, grid, buffers, args...; only_local_halos, kwargs...)
+    end
+
+    fill_corners!(c, arch.connectivity, indices, loc, arch, grid, buffers, args...; only_local_halos, kwargs...)
+
+    # We increment the tag counter only if we have actually initiated the MPI communication.
+    # This is the case only if at least one of the boundary conditions is a distributed communication 
+    # boundary condition (DCBCT) _and_ the `only_local_halos` keyword argument is false.
+    increment_tag = any(isa.(bcs, DCBCT)) && !only_local_halos
+
+    if increment_tag 
+        arch.mpi_tag[] += 1
+    end
+
+    switch_north_halos!(c, north_bc, grid, loc)
+
+    return nothing
+end
+
+function synchronize_communication!(field::Field{<:Any, <:Any, <:Any, <:Any, <:DTRG})
+    arch = architecture(field.grid)
+
+    # Wait for outstanding requests
+    if !isempty(arch.mpi_requests) 
+        cooperative_waitall!(arch.mpi_requests)
+
+        # Reset MPI tag
+        arch.mpi_tag[] = 0
+
+        # Reset MPI requests
+        empty!(arch.mpi_requests)
+    end
+
+    recv_from_buffers!(field.data, field.boundary_buffers, field.grid)
+
+    north_bc = field.boundary_conditions.north
+    instantiated_location = map(instantiate, location(field))
+
+    switch_north_halos!(field, north_bc, field.grid, instantiated_location)
+
+    return nothing
+end

src/distributed_zipper_north_tags.jl

@@ -0,0 +1,60 @@
+import Oceananigans.DistributedComputations: north_recv_tag, 
+                                             north_send_tag,
+                                             northwest_recv_tag, 
+                                             northwest_send_tag,
+                                             northeast_recv_tag, 
+                                             northeast_send_tag
+
+ID_DIGITS = 2
+
+sides  = (:west, :east, :south, :north, :southwest, :southeast, :northwest, :northeast)
+side_id = Dict(side => n-1 for (n, side) in enumerate(sides))
+
+# Change these and we are golden!
+function north_recv_tag(arch, ::DTRG, location)
+    field_id   = string(arch.mpi_tag[], pad=ID_DIGITS)
+    loc_digit  = string(loc_id(location...), pad=ID_DIGITS)
+    last_rank  = arch.local_index[2] == ranks(arch)[2]
+    side_digit = last_rank ? "8" : string(side_id[:south])
+    return parse(Int, field_id * loc_digit * side_digit)
+end
+
+function north_send_tag(arch, ::DTRG, location)
+    field_id   = string(arch.mpi_tag[], pad=ID_DIGITS)
+    loc_digit  = string(loc_id(location...), pad=ID_DIGITS)
+    last_rank  = arch.local_index[2] == ranks(arch)[2]
+    side_digit = last_rank ? "8" : string(side_id[:north])
+    return parse(Int, field_id * loc_digit * side_digit)
+end
+
+function northwest_recv_tag(arch, ::DTRG, location)
+    field_id   = string(arch.mpi_tag[], pad=ID_DIGITS)
+    loc_digit  = string(loc_id(location...), pad=ID_DIGITS)
+    last_rank  = arch.local_index[2] == ranks(arch)[2]
+    side_digit = last_rank ? "9" : string(side_id[:southeast])
+    return parse(Int, field_id * loc_digit * side_digit)
+end
+
+function northwest_send_tag(arch, ::DTRG, location)
+    field_id   = string(arch.mpi_tag[], pad=ID_DIGITS)
+    loc_digit  = string(loc_id(location...), pad=ID_DIGITS)
+    last_rank  = arch.local_index[2] == ranks(arch)[2]
+    side_digit = last_rank ? "9" : string(side_id[:northwest])
+    return parse(Int, field_id * loc_digit * side_digit)
+end
+
+function northeast_recv_tag(arch, ::DTRG, location)
+    field_id   = string(arch.mpi_tag[], pad=ID_DIGITS)
+    loc_digit  = string(loc_id(location...), pad=ID_DIGITS)
+    last_rank  = arch.local_index[2] == ranks(arch)[2]
+    side_digit = last_rank ? "10" : string(side_id[:southwest])
+    return parse(Int, field_id * loc_digit * side_digit)
+end
+
+function northeast_send_tag(arch, ::DTRG, location)
+    field_id   = string(arch.mpi_tag[], pad=ID_DIGITS)
+    loc_digit  = string(loc_id(location...), pad=ID_DIGITS)
+    last_rank  = arch.local_index[2] == ranks(arch)[2]
+    side_digit = last_rank ? "10" : string(side_id[:northeast])
+    return parse(Int, field_id * loc_digit * side_digit)
+end

src/grid_utils.jl

@@ -38,26 +38,8 @@ end
         d = lat_lon_to_cartesian(φᶠᶠᵃ[ i , j+1], λᶠᶠᵃ[ i , j+1], 1)
 
         Azᶜᶜᵃ[i, j] = spherical_area_quadrilateral(a, b, c, d) * radius^2
-
-        a = lat_lon_to_cartesian(φᶜᶠᵃ[i-1,  j ], λᶜᶠᵃ[i-1,  j ], 1)
-        b = lat_lon_to_cartesian(φᶜᶠᵃ[ i ,  j ], λᶜᶠᵃ[ i ,  j ], 1)
-        c = lat_lon_to_cartesian(φᶜᶠᵃ[ i , j+1], λᶜᶠᵃ[ i , j+1], 1)
-        d = lat_lon_to_cartesian(φᶜᶠᵃ[i-1, j+1], λᶜᶠᵃ[i-1, j+1], 1)
-
-        Azᶠᶜᵃ[i, j] = spherical_area_quadrilateral(a, b, c, d) * radius^2 
-
-        a = lat_lon_to_cartesian(φᶠᶜᵃ[ i , j-1], λᶠᶜᵃ[ i , j-1], 1)
-        b = lat_lon_to_cartesian(φᶠᶜᵃ[i+1, j-1], λᶠᶜᵃ[i+1, j-1], 1)
-        c = lat_lon_to_cartesian(φᶠᶜᵃ[i+1,  j ], λᶠᶜᵃ[i+1,  j ], 1)
-        d = lat_lon_to_cartesian(φᶠᶜᵃ[ i ,  j ], λᶠᶜᵃ[ i ,  j ], 1)
-
-        Azᶜᶠᵃ[i, j] = spherical_area_quadrilateral(a, b, c, d) * radius^2 
-
-        a = lat_lon_to_cartesian(φᶜᶜᵃ[i-1, j-1], λᶜᶜᵃ[i-1, j-1], 1)
-        b = lat_lon_to_cartesian(φᶜᶜᵃ[ i , j-1], λᶜᶜᵃ[ i , j-1], 1)
-        c = lat_lon_to_cartesian(φᶜᶜᵃ[ i ,  j ], λᶜᶜᵃ[ i ,  j ], 1)
-        d = lat_lon_to_cartesian(φᶜᶜᵃ[i-1,  j ], λᶜᶜᵃ[i-1,  j ], 1)
-
-        Azᶠᶠᵃ[i, j] = spherical_area_quadrilateral(a, b, c, d) * radius^2 
+        Azᶠᶜᵃ[i, j] = Δyᶠᶜᵃ[i, j] * Δxᶠᶜᵃ[i, j]
+        Azᶜᶠᵃ[i, j] = Δyᶜᶠᵃ[i, j] * Δxᶜᶠᵃ[i, j]
+        Azᶠᶠᵃ[i, j] = Δyᶠᶠᵃ[i, j] * Δxᶠᶠᵃ[i, j]
     end
 end

src/split_explicit_free_surface.jl

@@ -53,7 +53,7 @@ function positive_zipper_boundary(default_field, grid::DTRG)
         arch = architecture(grid)
         workers = ranks(arch.partition)
 
-        if arch.local_rank == workers[2] - 1
+        if arch.local_index[2] == workers[2]
                 return  FieldBoundaryConditions(
                                 top    = nothing,
                                 bottom = nothing,
@@ -76,7 +76,16 @@ end
 
 # We play the same trick as in the Distributed implementation and we extend the halos for
 # a split explicit barotropic solver on a tripolar grid. Only on the North boundary though!
-@inline tripolar_split_explicit_halos(old_halos, step_halo) = old_halos[1], max(step_halo, old_halos[2]), old_halos[3]
+@inline tripolar_split_explicit_halos(old_halos, step_halo, grid) = old_halos[1], max(step_halo, old_halos[2]), old_halos[3]
+
+@inline function tripolar_split_explicit_halos(old_halos, step_halo, grid::DTRG) 
+    Rx, Ry, _ = architecture(grid).ranks
+
+    Hx = Rx == 1 ? old_halos[1] : max(step_halo, old_halos[1])
+    Hy = max(step_halo, old_halos[2]) # Always!
+
+    return Hx, Hy, old_halos[3]
+end
 
 # Internal function for HydrostaticFreeSurfaceModel
 function materialize_free_surface(free_surface::SplitExplicitFreeSurface, velocities, grid::TRG)
@@ -86,7 +95,9 @@ function materialize_free_surface(free_surface::SplitExplicitFreeSurface, veloci
         old_halos  = halo_size(grid)
         Nsubsteps  = length(settings.substepping.averaging_weights)
 
-        extended_halos = tripolar_split_explicit_halos(old_halos, Nsubsteps+1)
+        # We need 1 additional halos in both directions because of the shifting
+        # caused by by the fill halo of the horizontal velocity.
+        extended_halos = tripolar_split_explicit_halos(old_halos, Nsubsteps+3, grid)
         extended_grid  = with_halo(extended_halos, grid)
 
         Nze = size(extended_grid, 3)

src/zipper_boundary_condition.jl

@@ -74,9 +74,7 @@ validate_boundary_condition_location(bc::Zipper, loc::Face, side) =
     Hy = grid.Hy
 
     for j = 1 : Hy
-        @inbounds begin
-            c[i, Ny + j, k] = s * c[i′, Ny - j + 1, k] 
-        end
+        @inbounds c[i, Ny + j, k] = s * c[i′, Ny - j + 1, k] 
     end
 
     return nothing
@@ -91,9 +89,7 @@ end
     Hy = grid.Hy
 
     for j = 1 : Hy
-        @inbounds begin
-            c[i, Ny + j, k] = s * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1
-        end
+        @inbounds c[i, Ny + j, k] = s * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1
     end
 
     return nothing
@@ -106,9 +102,7 @@ end
     Hy = grid.Hy
 
     for j = 1 : Hy
-        @inbounds begin
-            c[i, Ny + j, k] = sign * c[i′, Ny - j + 1, k] 
-        end
+        @inbounds c[i, Ny + j, k] = sign * c[i′, Ny - j + 1, k] 
     end
 
     return nothing
@@ -121,9 +115,7 @@ end
     Hy = grid.Hy
 
     for j = 1 : Hy
-        @inbounds begin
-            c[i, Ny + j, k] = sign * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1
-        end
+        @inbounds c[i, Ny + j, k] = sign * c[i′, Ny - j, k] # The Ny line is duplicated so we substitute starting Ny-1
     end
 
     return nothing

test/dependencies_for_runtests.jl

@@ -0,0 +1,28 @@
+using OrthogonalSphericalShellGrids
+using Oceananigans
+using Oceananigans.Grids: halo_size
+using Oceananigans.Utils
+using Oceananigans.Units
+using Oceananigans.BoundaryConditions
+using OrthogonalSphericalShellGrids: get_cartesian_nodes_and_vertices
+using Oceananigans.CUDA
+using Test
+
+using KernelAbstractions: @kernel, @index
+
+arch = CUDA.has_cuda_gpu() ? GPU() : CPU()
+
+# Mask the singularity of the grid in a region of 
+# 5 degrees radius around the singularities
+function mask_singularities(underlying_grid::TripolarGrid)
+    λp = underlying_grid.conformal_mapping.first_pole_longitude
+    φp = underlying_grid.conformal_mapping.north_poles_latitude
+
+    # We need a bottom height field that ``masks'' the singularities
+    bottom_height(λ, φ) = ((abs(λ - λp) < 5)       & (abs(φp - φ) < 5)) |
+                          ((abs(λ - λp - 180) < 5) & (abs(φp - φ) < 5)) | (φ < -80) ? 0 : - 1000
+
+    grid = ImmersedBoundaryGrid(underlying_grid, GridFittedBottom(bottom_height))
+
+    return grid
+end

test/distributed_tests_utils.jl

@@ -0,0 +1,58 @@
+using JLD2
+using MPI
+using Oceananigans.DistributedComputations: reconstruct_global_field
+
+include("dependencies_for_runtests.jl")
+
+# Run the distributed grid simulation and save down reconstructed results
+function run_distributed_tripolar_grid(arch, filename)
+    distributed_grid = TripolarGrid(arch; size = (100, 100, 1), z = (-1000, 0), halo = (5, 5, 5))
+    distributed_grid = mask_singularities(distributed_grid)
+    simulation       = run_tripolar_simulation(distributed_grid)
+
+    η = reconstruct_global_field(simulation.model.free_surface.η)
+    u = reconstruct_global_field(simulation.model.velocities.u)
+    v = reconstruct_global_field(simulation.model.velocities.v)
+    c = reconstruct_global_field(simulation.model.tracers.c)
+
+    fill_halo_regions!(η)
+    fill_halo_regions!(u)
+    fill_halo_regions!(v)
+    fill_halo_regions!(c)
+
+    if arch.local_rank == 0
+        jldsave(filename; η = interior(η, :, :, 1), 
+                          u = interior(u, :, :, 1),
+                          v = interior(v, :, :, 1), 
+                          c = interior(c, :, :, 1)) 
+    end
+
+    MPI.Barrier(MPI.COMM_WORLD)
+    MPI.Finalize()
+
+    return nothing
+end
+
+# Just a random simulation on a tripolar grid
+function run_tripolar_simulation(grid)
+
+    model = HydrostaticFreeSurfaceModel(; grid = grid,
+                                          free_surface = SplitExplicitFreeSurface(grid; substeps = 20),
+                                          tracers = :c,
+                                          buoyancy = nothing, 
+                                          tracer_advection = WENO(),
+                                          momentum_advection = VectorInvariant(),
+                                          coriolis = HydrostaticSphericalCoriolis())
+
+    # Setup the model with a gaussian sea surface height
+    # near the physical north poles and one near the equator
+    ηᵢ(λ, φ, z) = exp(- (φ - 90)^2 / 10^2) + exp(- φ^2 / 10^2)
+
+    set!(model, η = ηᵢ, c = ηᵢ)
+
+    simulation = Simulation(model, Δt = 5minutes, stop_iteration = 100)
+
+    run!(simulation)
+
+    return simulation
+end