Fix distributed tripolar fold

src/Advection/topologically_conditional_interpolation.jl

@@ -20,7 +20,9 @@ using Oceananigans.Grids: AbstractGrid,
 const AG = AbstractGrid
 
 # topologies bounded at least on one side
-const BT = Union{Bounded, RightConnected, LeftConnected}
+const BT = Union{Bounded, RightConnected, LeftConnected,
+                  LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded,
+                  LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected}
 
 # Bounded underlying Grids
 const AGX   = AG{<:Any, <:BT}

src/BoundaryConditions/field_boundary_conditions.jl

@@ -21,13 +21,18 @@ default_prognostic_bc(::RightConnected, ::Center, default)  = default.boundary_c
 default_prognostic_bc(::RightFaceFolded, ::Center, default) = default.boundary_condition
 default_prognostic_bc(::RightCenterFolded, ::Center, default) = default.boundary_condition
 
+const DistributedFoldTopology = Union{LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded,
+                                      LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected}
+
+default_prognostic_bc(::DistributedFoldTopology, ::Center, default) = default.boundary_condition
 
 # TODO: make model constructors enforce impenetrability on velocity components to simplify this code
 default_prognostic_bc(::Bounded,        ::Face, default) = ImpenetrableBoundaryCondition()
 default_prognostic_bc(::LeftConnected,  ::Face, default) = ImpenetrableBoundaryCondition()
 default_prognostic_bc(::RightConnected, ::Face, default) = ImpenetrableBoundaryCondition()
 default_prognostic_bc(::RightFaceFolded, ::Face, default) = ImpenetrableBoundaryCondition()
 default_prognostic_bc(::RightCenterFolded, ::Face, default) = ImpenetrableBoundaryCondition()
+default_prognostic_bc(::DistributedFoldTopology, ::Face, default) = ImpenetrableBoundaryCondition()
 
 default_prognostic_bc(::Bounded,        ::Nothing, default) = nothing
 default_prognostic_bc(::Flat,           ::Nothing, default) = nothing
@@ -37,13 +42,15 @@ default_prognostic_bc(::LeftConnected,  ::Nothing, default) = nothing
 default_prognostic_bc(::RightConnected, ::Nothing, default) = nothing
 default_prognostic_bc(::RightFaceFolded, ::Nothing, default) = nothing
 default_prognostic_bc(::RightCenterFolded, ::Nothing, default) = nothing
+default_prognostic_bc(::DistributedFoldTopology, ::Nothing, default) = nothing
 
 _default_auxiliary_bc(topo, loc) = default_prognostic_bc(topo, loc, DefaultBoundaryCondition())
 _default_auxiliary_bc(::Bounded, ::Face)        = nothing
 _default_auxiliary_bc(::RightConnected, ::Face) = nothing
 _default_auxiliary_bc(::LeftConnected,  ::Face) = nothing
 _default_auxiliary_bc(::RightFaceFolded, ::Face) = nothing
 _default_auxiliary_bc(::RightCenterFolded, ::Face) = nothing
+_default_auxiliary_bc(::DistributedFoldTopology, ::Face) = nothing
 
 default_auxiliary_bc(grid, ::Val{:east}, loc)   = _default_auxiliary_bc(topology(grid, 1)(), loc[1])
 default_auxiliary_bc(grid, ::Val{:west}, loc)   = _default_auxiliary_bc(topology(grid, 1)(), loc[1])

src/DistributedComputations/distributed_grids.jl

@@ -349,6 +349,26 @@ insert_connected_topology(::Type{Bounded}, R, r) = ifelse(R == 1, Bounded,
 
 insert_connected_topology(::Type{Periodic}, R, r) = ifelse(R == 1, Periodic, FullyConnected)
 
+# Fold-aware topology insertion for distributed tripolar grids.
+# These take 5 arguments: (global_y_topology, Ry, ry, Rx, rx) where
+# Ry/ry are y-rank count/index and Rx/rx are x-rank count/index (all 1-based).
+
+local_fold_topology(::Type{RightCenterFolded}, Rx, rx) =
+    Rx == 1 ? LeftConnectedRightCenterFolded : LeftConnectedRightCenterConnected
+
+local_fold_topology(::Type{RightFaceFolded}, Rx, rx) =
+    Rx == 1 ? LeftConnectedRightFaceFolded : LeftConnectedRightFaceConnected
+
+function insert_connected_topology(T::Type{<:Union{RightCenterFolded, RightFaceFolded}}, Ry, ry, Rx, rx)
+    if ry == 1
+        return RightConnected
+    elseif ry == Ry
+        return local_fold_topology(T, Rx, rx)
+    else
+        return FullyConnected
+    end
+end
+
 """
     reconstruct_global_topology(T, R, r, r1, r2, arch)
 

src/Grids/Grids.jl

@@ -4,6 +4,8 @@ export Center, Face
 export AbstractTopology, topology
 export Periodic, Bounded, Flat, FullyConnected, LeftConnected, RightConnected
 export RightFaceFolded, RightCenterFolded
+export LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded
+export LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected
 export AbstractGrid, AbstractUnderlyingGrid, halo_size, total_size
 export RectilinearGrid
 export AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid
@@ -117,6 +119,41 @@ Grid topology for tripolar U-point pivot connection.
 """
 struct RightCenterFolded <: AbstractTopology end
 
+"""
+    LeftConnectedRightCenterFolded
+
+Local grid topology for the northernmost y-rank of a 1×N distributed tripolar grid
+with U-point pivot (serial fold). Connected to the south neighbor on the left,
+center-folded on the right (north).
+"""
+struct LeftConnectedRightCenterFolded <: AbstractTopology end
+
+"""
+    LeftConnectedRightFaceFolded
+
+Local grid y topology for the northernmost y-rank of a 1×N distributed tripolar grid
+with F-point pivot (serial fold). Connected to the south neighbor on the left,
+face-folded on the right (north). Face-extended (Ny+1 Face points in y).
+"""
+struct LeftConnectedRightFaceFolded <: AbstractTopology end
+
+"""
+    LeftConnectedRightCenterConnected
+
+Local grid y topology for the northernmost y-rank of an M×N distributed tripolar grid
+with U-point pivot (distributed zipper).
+"""
+struct LeftConnectedRightCenterConnected <: AbstractTopology end
+
+"""
+    LeftConnectedRightFaceConnected
+
+Local grid y topology for the northernmost y-rank of an M×N distributed tripolar grid
+with F-point pivot (distributed zipper).
+Face-extended (Ny+1 Face points in y).
+"""
+struct LeftConnectedRightFaceConnected <: AbstractTopology end
+
 #####
 ##### Directions (for tilted domains)
 #####

src/Grids/grid_utils.jl

@@ -40,7 +40,9 @@ end
     end
 end
 
-const BoundedTopology = Union{Bounded, LeftConnected, RightFaceFolded}
+const BoundedTopology = Union{Bounded, LeftConnected, RightFaceFolded,
+                              LeftConnectedRightFaceFolded,
+                              LeftConnectedRightFaceConnected}
 const AT = AbstractTopology
 
 Base.length(::Face,    ::BoundedTopology, N) = N + 1

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl

@@ -10,7 +10,11 @@ using Oceananigans.ImmersedBoundaries: column_depthTᶠᶜᵃ, column_depthTᶜ
 using Oceananigans.Operators: ∂xᵣTᶠᶜᶠ, ∂xᵣᶠᶜᶠ, ∂yᵣTᶜᶠᶠ, ∂yᵣᶜᶠᶠ, δxTᶜᵃᵃ, δxᶜᵃᵃ, δyTᵃᶜᵃ, δyᵃᶜᵃ
 using Oceananigans.BoundaryConditions: FieldBoundaryConditions, fill_halo_regions!
 using Oceananigans.Fields: Field
-using Oceananigans.Grids: Center, Face, get_active_column_map, topology
+using Oceananigans.Grids: Center, Face, get_active_column_map, topology,
+                          LeftConnected, RightConnected, FullyConnected,
+                          RightCenterFolded, RightFaceFolded,
+                          LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded,
+                          LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected
 using Oceananigans.ImmersedBoundaries: mask_immersed_field!
 using Oceananigans.Models.HydrostaticFreeSurfaceModels: AbstractFreeSurface,
                                                         free_surface_displacement_field,

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl

@@ -225,7 +225,10 @@ function hydrostatic_tendency_fields(velocities, free_surface::SplitExplicitFree
     return merge((u=u, v=v, U=U, V=V), tracers)
 end
 
-const ConnectedTopology = Union{LeftConnected, RightConnected, FullyConnected, RightCenterFolded, RightFaceFolded}
+const ConnectedTopology = Union{LeftConnected, RightConnected, FullyConnected,
+                                RightCenterFolded, RightFaceFolded,
+                                LeftConnectedRightCenterFolded, LeftConnectedRightFaceFolded,
+                                LeftConnectedRightCenterConnected, LeftConnectedRightFaceConnected}
 
 # Internal function for HydrostaticFreeSurfaceModel
 function materialize_free_surface(free_surface::SplitExplicitFreeSurface{extend_halos}, velocities, grid) where {extend_halos}
@@ -403,6 +406,12 @@ split_explicit_kernel_size(::Type{LeftConnected},  N, H) = -H+2:N
 split_explicit_kernel_size(::Type{RightCenterFolded}, N, H) = 1:N+H-1
 split_explicit_kernel_size(::Type{RightFaceFolded}, N, H)   = 1:N+H-1
 
+# Distributed fold topologies: connected on both sides (left=MPI, right=fold/zipper)
+split_explicit_kernel_size(::Type{LeftConnectedRightCenterFolded},    N, H) = -H+2:N+H-1
+split_explicit_kernel_size(::Type{LeftConnectedRightFaceFolded},      N, H) = -H+2:N+H-1
+split_explicit_kernel_size(::Type{LeftConnectedRightCenterConnected}, N, H) = -H+2:N+H-1
+split_explicit_kernel_size(::Type{LeftConnectedRightFaceConnected},   N, H) = -H+2:N+H-1
+
 # Adapt
 Adapt.adapt_structure(to, free_surface::SplitExplicitFreeSurface{extend_halos}) where {extend_halos} =
     SplitExplicitFreeSurface{extend_halos}(Adapt.adapt(to, free_surface.displacement),

src/Models/NonhydrostaticModels/compute_nonhydrostatic_buffer_tendencies.jl

@@ -2,6 +2,7 @@ import Oceananigans.Models: compute_buffer_tendencies!
 
 using Oceananigans.TurbulenceClosures: required_halo_size_x, required_halo_size_y
 using Oceananigans.Grids: XFlatGrid, YFlatGrid
+using Oceananigans.Utils: worksize
 
 # TODO: the code in this file is difficult to understand.
 # Rewriting it may be helpful.
@@ -24,44 +25,48 @@ function compute_buffer_tendencies!(model::NonhydrostaticModel)
     return nothing
 end
 
-# tendencies need computing in the range 1 : H and N - H + 1 : N
+# Tendencies need computing in the range 1 : H and W - H + 1 : W where W = worksize.
+# These buffer regions complement the halo-independent interior kernel (H+1:W-H).
 function buffer_tendency_kernel_parameters(grid, arch)
     Nx, Ny, Nz = size(grid)
+    Wx, Wy, _  = worksize(grid)
     Hx, Hy, _  = halo_size(grid)
 
-    param_west  = (1:Hx,       1:Ny,       1:Nz)
-    param_east  = (Nx-Hx+1:Nx, 1:Ny,       1:Nz)
-    param_south = (1:Nx,       1:Hy,       1:Nz)
-    param_north = (1:Nx,       Ny-Hy+1:Ny, 1:Nz)
+    param_west  = (1:Hx,       1:Wy,       1:Nz)
+    param_east  = (Wx-Hx+1:Wx, 1:Wy,       1:Nz)
+    param_south = (1:Wx,       1:Hy,       1:Nz)
+    param_north = (1:Wx,       Wy-Hy+1:Wy, 1:Nz)
 
     params = (param_west, param_east, param_south, param_north)
     return buffer_parameters(params, grid, arch)
 end
 
-# p needs computing in the range  0 : 0 and N + 1 : N + 1
+# p needs computing in the range  0 : 0 and W + 1 : W + 1 where W = worksize
 function buffer_p_kernel_parameters(grid, arch)
     Nx, Ny, _ = size(grid)
+    Wx, Wy, _ = worksize(grid)
 
-    param_west  = (0:0,       1:Ny)
-    param_east  = (Nx+1:Nx+1, 1:Ny)
-    param_south = (1:Nx,      0:0)
-    param_north = (1:Nx,      Ny+1:Ny+1)
+    param_west  = (0:0,       1:Wy)
+    param_east  = (Wx+1:Wx+1, 1:Wy)
+    param_south = (1:Wx,      0:0)
+    param_north = (1:Wx,      Wy+1:Wy+1)
 
     params = (param_west, param_east, param_south, param_north)
     return buffer_parameters(params, grid, arch)
 end
 
-# closure_fields need recomputing in the range 0 : B and N - B + 1 : N + 1
+# closure_fields need recomputing in the range 0 : B and W - B + 1 : W + 1 where W = worksize
 function buffer_κ_kernel_parameters(grid, closure, arch)
     Nx, Ny, Nz = size(grid)
+    Wx, Wy, _  = worksize(grid)
 
     Bx = required_halo_size_x(closure)
     By = required_halo_size_y(closure)
 
-    param_west  = (0:Bx,         1:Ny,         1:Nz)
-    param_east  = (Nx-Bx+1:Nx+1, 1:Ny,         1:Nz)
-    param_south = (1:Nx,         0:By,         1:Nz)
-    param_north = (1:Nx,         Ny-By+1:Ny+1, 1:Nz)
+    param_west  = (0:Bx,         1:Wy,         1:Nz)
+    param_east  = (Wx-Bx+1:Wx+1, 1:Wy,         1:Nz)
+    param_south = (1:Wx,         0:By,         1:Nz)
+    param_north = (1:Wx,         Wy-By+1:Wy+1, 1:Nz)
 
     params = (param_west, param_east, param_south, param_north)
     return buffer_parameters(params, grid, arch)

src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl

@@ -7,10 +7,11 @@ import Oceananigans
 import Oceananigans.Architectures: on_architecture
 
 using Oceananigans.Architectures: on_architecture, AbstractArchitecture, CPU, GPU
-using Oceananigans.BoundaryConditions: BoundaryCondition, UZBC
-using Oceananigans.Grids: AbstractTopology, RightConnected
+using Oceananigans.BoundaryConditions: BoundaryCondition
+using Oceananigans.Grids: AbstractTopology
 using Oceananigans.Grids: halo_size, generate_coordinate, topology
 using Oceananigans.Grids: total_length, add_halos, fill_metric_halo_regions!
+using Oceananigans.BoundaryConditions: fill_halo_regions!
 
 using Distances: haversine
 using Adapt: Adapt, adapt

src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl

@@ -1,4 +1,4 @@
-using Oceananigans.BoundaryConditions: DistributedCommunicationBoundaryCondition
+using Oceananigans.BoundaryConditions: DistributedCommunicationBoundaryCondition, ZBC
 using Oceananigans.Fields: validate_indices, validate_field_data
 using Oceananigans.DistributedComputations:
     DistributedComputations,
@@ -11,7 +11,10 @@ using Oceananigans.DistributedComputations:
     concatenate_local_sizes,
     communication_buffers
 
-using Oceananigans.Grids: topology, RightConnected, FullyConnected
+using Oceananigans.Grids: topology, FullyConnected,
+    LeftConnectedRightFaceFolded, LeftConnectedRightFaceConnected
+using Oceananigans.DistributedComputations: insert_connected_topology
+using Oceananigans.Utils: Utils
 
 import Oceananigans.Fields: Field, validate_indices, validate_boundary_conditions
 
@@ -109,8 +112,14 @@ function TripolarGrid(arch::Distributed, FT::DataType=Float64;
     Azᶜᶠᵃ = partition_tripolar_metric(global_grid, :Azᶜᶠᵃ, irange, jrange)
     Azᶠᶠᵃ = partition_tripolar_metric(global_grid, :Azᶠᶠᵃ, irange, jrange)
 
-    LY = yrank == 0 ? RightConnected : FullyConnected
     LX = workers[1] == 1 ? Periodic : FullyConnected
+
+    global_fold_topology = topology(global_grid, 2)
+
+    # 1-based indices for insert_connected_topology
+    Rx, Ry = workers[1], workers[2]
+    rx, ry = xrank + 1, yrank + 1
+    LY = insert_connected_topology(global_fold_topology, Ry, ry, Rx, rx)
     ny = nylocal[yrank+1]
     nx = nxlocal[xrank+1]
 
@@ -228,9 +237,6 @@ function receiving_rank(arch; receive_idx_x = ranks(arch)[1] - arch.local_index[
     return receive_rank
 end
 
-# a distributed `TripolarGrid` needs a `UPivotZipperBoundaryCondition` for the north boundary
-# only on the last rank
-# TODO: generalize to any ZipperBoundaryCondition
 function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions,
                                               grid::MPITripolarGridOfSomeKind,
                                               field_name::Symbol,
@@ -248,7 +254,7 @@ function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions,
     south = regularize_boundary_condition(bcs.south, grid, loc, 2, LeftBoundary,  prognostic_names)
 
     north = if yrank == processor_size[2] - 1 && processor_size[1] == 1
-        UPivotZipperBoundaryCondition(sign)
+        north_fold_boundary_condition(fold_topology(grid.conformal_mapping))(sign)
 
     elseif yrank == processor_size[2] - 1 && processor_size[1] != 1
         from = arch.local_rank
@@ -280,26 +286,20 @@ function Field(loc::Tuple{<:LX, <:LY, <:LZ}, grid::MPITripolarGridOfSomeKind, da
     validate_field_data(loc, data, grid, indices)
     validate_boundary_conditions(loc, grid, old_bcs)
 
-    default_zipper = UPivotZipperBoundaryCondition(sign(LX, LY))
-
     if isnothing(old_bcs) || ismissing(old_bcs)
         new_bcs = old_bcs
     else
         new_bcs = inject_halo_communication_boundary_conditions(old_bcs, loc, arch.local_rank, arch.connectivity, topology(grid))
 
-        # North boundary conditions are "special". If we are at the top of the domain, i.e.
-        # the last rank, then we need to substitute the BC only if the old one is not already
-        # a zipper boundary condition. Otherwise we always substitute because we need to
-        # inject the halo boundary conditions.
         if yrank == processor_size[2] - 1 && processor_size[1] == 1
-            north_bc = if !(old_bcs.north isa UZBC)
-                default_zipper
+            north_bc = if !(old_bcs.north isa ZBC)
+                north_fold_boundary_condition(fold_topology(grid.conformal_mapping))(sign(LX, LY))
             else
                 old_bcs.north
             end
 
         elseif yrank == processor_size[2] - 1 && processor_size[1] != 1
-            sgn  = old_bcs.north isa UZBC ? old_bcs.north.condition : sign(LX, LY)
+            sgn  = old_bcs.north isa ZBC ? old_bcs.north.condition : sign(LX, LY)
             from = arch.local_rank
             to   = arch.connectivity.north
             halo_communication = ZipperHaloCommunicationRanks(sgn; from, to)
@@ -317,7 +317,7 @@ function Field(loc::Tuple{<:LX, <:LY, <:LZ}, grid::MPITripolarGridOfSomeKind, da
                                             bottom=new_bcs.bottom)
     end
 
-    buffers = communication_buffers(grid, data, new_bcs)
+    buffers = communication_buffers(grid, data, new_bcs, (LX(), LY(), LZ()))
 
     return Field{LX, LY, LZ}(grid, data, new_bcs, indices, op, status, buffers)
 end
@@ -347,7 +347,8 @@ function DistributedComputations.reconstruct_global_grid(grid::MPITripolarGrid)
                         north_poles_latitude,
                         first_pole_longitude,
                         southernmost_latitude,
-                        z)
+                        z,
+                        fold_topology = fold_topology(grid.conformal_mapping))
 end
 
 function Grids.with_halo(new_halo, old_grid::MPITripolarGrid)
@@ -361,12 +362,22 @@ function Grids.with_halo(new_halo, old_grid::MPITripolarGrid)
     north_poles_latitude = old_grid.conformal_mapping.north_poles_latitude
     first_pole_longitude = old_grid.conformal_mapping.first_pole_longitude
     southernmost_latitude = old_grid.conformal_mapping.southernmost_latitude
-
     return TripolarGrid(arch, eltype(old_grid);
                         halo = new_halo,
                         size = N,
                         north_poles_latitude,
                         first_pole_longitude,
                         southernmost_latitude,
-                        z)
+                        z,
+                        fold_topology = fold_topology(old_grid.conformal_mapping))
 end
+
+#####
+##### Extend worksize for distributed FPivot grids (matches RFTRG worksize for serial FPivot)
+#####
+
+const DistributedFPivotTopology = Union{LeftConnectedRightFaceFolded, LeftConnectedRightFaceConnected}
+const DRFTRG = Union{MPITripolarGrid{<:Any, <:Any, <:DistributedFPivotTopology},
+                     ImmersedBoundaryGrid{<:Any, <:Any, <:DistributedFPivotTopology, <:Any, <:MPITripolarGrid}}
+
+Utils.worksize(grid::DRFTRG) = grid.Nx, grid.Ny+1, grid.Nz