Add a RightFolded topology

ext/OceananigansReactantExt/Grids/Grids.jl

@@ -9,7 +9,7 @@ using Oceananigans
 using Oceananigans: Distributed
 using Oceananigans.Architectures: ReactantState, CPU
 using Oceananigans.Grids: AbstractGrid, AbstractUnderlyingGrid, StaticVerticalDiscretization, MutableVerticalDiscretization
-using Oceananigans.Grids: Center, Face, RightConnected, LeftConnected, Periodic, Bounded, Flat, BoundedTopology
+using Oceananigans.Grids: Center, Face, RightConnected, LeftConnected, Periodic, Bounded, Flat, RightFolded, BoundedTopology
 using Oceananigans.Fields: Field
 using Oceananigans.ImmersedBoundaries: GridFittedBottom, AbstractImmersedBoundary
 

ext/OceananigansReactantExt/Grids/sharded_grids.jl

@@ -226,7 +226,7 @@ function TripolarGrid(arch::ShardedDistributed,
     # ``1'' here is the maximum number of dimensions of the fields of ``z''
     replicate = Sharding.Replicated(arch.connectivity)
 
-    grid = OrthogonalSphericalShellGrid{Periodic, RightConnected, Bounded}(arch,
+    grid = OrthogonalSphericalShellGrid{Periodic, RightFolded, Bounded}(arch,
         global_size...,
         halo...,
         convert(FT, global_grid.Lz),

ext/OceananigansReactantExt/OceananigansReactantExt.jl

@@ -7,7 +7,7 @@ using OffsetArrays
 using Oceananigans: Distributed, DistributedComputations, ReactantState, CPU,
                     OrthogonalSphericalShellGrids
 using Oceananigans.Architectures: on_architecture
-using Oceananigans.Grids: Bounded, Periodic, RightConnected
+using Oceananigans.Grids: Bounded, Periodic, RightConnected, RightFolded
 
 deconcretize(obj) = obj # fallback
 deconcretize(a::OffsetArray) = OffsetArray(Array(a.parent), a.offsets...)

src/Advection/topologically_conditional_interpolation.jl

@@ -10,17 +10,19 @@
 #####        close to the boundary, or a second-order interpolation if i is close to a boundary.
 #####
 
-using Oceananigans.Grids: AbstractGrid,
-                          Bounded,
-                          RightConnected,
-                          LeftConnected,
+using Oceananigans.Grids: AbstractUnderlyingGrid, 
+                          Bounded, 
+                          RightConnected, 
+                          LeftConnected, 
+                          RightFolded,
                           topology,
                           architecture
 
 const AG = AbstractGrid
 
 # topologies bounded at least on one side
-const BT = Union{Bounded, RightConnected, LeftConnected}
+const BT = Union{Bounded, RightConnected, LeftConnected, RightFolded}
+const RightTopology = Union{RightConnected, RightFolded}
 
 # Bounded underlying Grids
 const AGX   = AG{<:Any, <:BT}
@@ -43,31 +45,31 @@ for dir in (:x, :y, :z)
 
     @eval begin
         # Bounded topologies
-        @inline $outside_symmetric_haloᶠ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - $required_halo_size(adv))
-        @inline $outside_symmetric_haloᶜ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv))     & (i <= N + 1 - $required_halo_size(adv))
+        @inline $outside_symmetric_haloᶠ(i, ::Bounded, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - $required_halo_size(adv))
+        @inline $outside_symmetric_haloᶜ(i, ::Bounded, N, adv) = (i >= $required_halo_size(adv))     & (i <= N + 1 - $required_halo_size(adv))
 
-        @inline $outside_biased_haloᶠ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - ($required_halo_size(adv) - 1)) &  # Left bias
-                                                                    (i >= $required_halo_size(adv))     & (i <= N + 1 - $required_halo_size(adv))          # Right bias
-        @inline $outside_biased_haloᶜ(i, ::Type{Bounded}, N, adv) = (i >= $required_halo_size(adv))     & (i <= N + 1 - ($required_halo_size(adv) - 1)) &  # Left bias
-                                                                    (i >= $required_halo_size(adv) - 1) & (i <= N + 1 - $required_halo_size(adv))          # Right bias
+        @inline $outside_biased_haloᶠ(i, ::Bounded, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - ($required_halo_size(adv) - 1)) &  # Left bias
+                                                              (i >= $required_halo_size(adv))     & (i <= N + 1 - $required_halo_size(adv))          # Right bias
+        @inline $outside_biased_haloᶜ(i, ::Bounded, N, adv) = (i >= $required_halo_size(adv))     & (i <= N + 1 - ($required_halo_size(adv) - 1)) &  # Left bias
+                                                              (i >= $required_halo_size(adv) - 1) & (i <= N + 1 - $required_halo_size(adv))          # Right bias
 
         # Right connected topologies (only test the left side, i.e. the bounded side)
-        @inline $outside_symmetric_haloᶠ(i, ::Type{RightConnected}, N, adv) = i >= $required_halo_size(adv) + 1
-        @inline $outside_symmetric_haloᶜ(i, ::Type{RightConnected}, N, adv) = i >= $required_halo_size(adv)
+        @inline $outside_symmetric_haloᶠ(i, ::RightTopology, N, adv) = i >= $required_halo_size(adv) + 1
+        @inline $outside_symmetric_haloᶜ(i, ::RightTopology, N, adv) = i >= $required_halo_size(adv)
 
-        @inline $outside_biased_haloᶠ(i, ::Type{RightConnected}, N, adv) = (i >= $required_halo_size(adv) + 1) &  # Left bias
-                                                                           (i >= $required_halo_size(adv))        # Right bias
-        @inline $outside_biased_haloᶜ(i, ::Type{RightConnected}, N, adv) = (i >= $required_halo_size(adv))     &  # Left bias
-                                                                           (i >= $required_halo_size(adv) - 1)    # Right bias
+        @inline $outside_biased_haloᶠ(i, ::RightTopology, N, adv) = (i >= $required_halo_size(adv) + 1) &  # Left bias
+                                                                    (i >= $required_halo_size(adv))        # Right bias
+        @inline $outside_biased_haloᶜ(i, ::RightTopology, N, adv) = (i >= $required_halo_size(adv))     &  # Left bias
+                                                                    (i >= $required_halo_size(adv) - 1)    # Right bias
 
         # Left bounded topologies (only test the right side, i.e. the bounded side)
-        @inline $outside_symmetric_haloᶠ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - $required_halo_size(adv))
-        @inline $outside_symmetric_haloᶜ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - $required_halo_size(adv))
+        @inline $outside_symmetric_haloᶠ(i, ::LeftConnected, N, adv) = (i <= N + 1 - $required_halo_size(adv))
+        @inline $outside_symmetric_haloᶜ(i, ::LeftConnected, N, adv) = (i <= N + 1 - $required_halo_size(adv))
 
-        @inline $outside_biased_haloᶠ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - ($required_halo_size(adv) - 1)) &  # Left bias
-                                                                          (i <= N + 1 - $required_halo_size(adv))          # Right bias
-        @inline $outside_biased_haloᶜ(i, ::Type{LeftConnected}, N, adv) = (i <= N + 1 - ($required_halo_size(adv) - 1)) &  # Left bias
-                                                                          (i <= N + 1 - $required_halo_size(adv))          # Right bias
+        @inline $outside_biased_haloᶠ(i, ::LeftConnected, N, adv) = (i <= N + 1 - ($required_halo_size(adv) - 1)) &  # Left bias
+                                                                    (i <= N + 1 - $required_halo_size(adv))          # Right bias
+        @inline $outside_biased_haloᶜ(i, ::LeftConnected, N, adv) = (i <= N + 1 - ($required_halo_size(adv) - 1)) &  # Left bias
+                                                                    (i <= N + 1 - $required_halo_size(adv))          # Right bias
     end
 end
 

src/BoundaryConditions/BoundaryConditions.jl

@@ -13,10 +13,11 @@ export
 using CUDA, Adapt
 using KernelAbstractions: @index, @kernel
 
+using Oceananigans.Grids
 using Oceananigans.Architectures: CPU, GPU, device
 using Oceananigans.Utils: work_layout, launch!
 using Oceananigans.Operators: Ax, Ay, Az, volume
-using Oceananigans.Grids
+using Oceananigans.Grids: RightFolded
 
 import Adapt: adapt_structure
 

src/BoundaryConditions/field_boundary_conditions.jl

@@ -18,22 +18,24 @@ default_prognostic_bc(::Flat,           loc,      default)  = nothing
 default_prognostic_bc(::Bounded,        ::Center, default)  = default.boundary_condition
 default_prognostic_bc(::LeftConnected,  ::Center, default)  = default.boundary_condition
 default_prognostic_bc(::RightConnected, ::Center, default)  = default.boundary_condition
+default_prognostic_bc(::RightFolded,    ::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(::RightFolded,    ::Face, default)    = ImpenetrableBoundaryCondition()
 
 default_prognostic_bc(::Bounded,        ::Nothing, default) = nothing
 default_prognostic_bc(::Flat,           ::Nothing, default) = nothing
 default_prognostic_bc(::Grids.Periodic, ::Nothing, default) = nothing
 default_prognostic_bc(::FullyConnected, ::Nothing, default) = nothing
 default_prognostic_bc(::LeftConnected,  ::Nothing, default) = nothing
-default_prognostic_bc(::RightConnected, ::Nothing, default) = nothing
+default_prognostic_bc(::RightFolded,    ::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(::RightFolded, ::Face)    = nothing
 default_auxiliary_bc(::LeftConnected,  ::Face) = nothing
 
 #####

src/DistributedComputations/distributed_architectures.jl

@@ -381,7 +381,7 @@ Generate a `NeighboringRanks` object that holds the MPI ranks of the neighboring
 NeighboringRanks(; east, west, north, south, southwest, southeast, northwest, northeast) =
     NeighboringRanks(east, west, north, south, southwest, southeast, northwest, northeast)
 
-# The "Periodic" topologies are `Periodic`, `FullyConnected` and `RightConnected`
+# The "Periodic" topologies are `Periodic`, `RightFolded` and `FullyConnected`, and `RightConnected`
 # The "Bounded" topologies are `Bounded` and `LeftConnected`
 function increment_index(i, R)
     R == 1 && return nothing

src/Grids/Grids.jl

@@ -107,6 +107,13 @@ Grid topology for dimensions that are connected to other models or domains only
 """
 struct RightConnected <: AbstractTopology end
 
+"""
+    RightFolded
+
+Grid topology for dimensions that are folded onto each other on the right (the other direction is bounded).
+"""
+struct RightFolded <: AbstractTopology end
+
 #####
 ##### Directions (for tilted domains)
 #####

src/Grids/grid_utils.jl

@@ -286,6 +286,7 @@ function domain_summary(topo, name, (left, right))
 
     topo_string = topo isa Periodic ? "Periodic " :
                   topo isa Bounded ? "Bounded  " :
+                  topo isa RightFolded ? "RightFolded " :
                   topo isa FullyConnected ? "FullyConnected " :
                   topo isa LeftConnected ? "LeftConnected  " :
                   topo isa RightConnected ? "RightConnected  " :

src/Grids/inactive_node.jl

@@ -17,7 +17,7 @@ function build_condition(Topo, side, dim, array::Bool)
         else
             return :(($side > grid.$dim))
         end
-    else # RightConnected
+    else # RightConnected and RightFolded
         if array
             return :((ReactantCore.materialize_traced_array($side) .< 1))
         else
@@ -36,7 +36,7 @@ end
 Return `true` when the tracer cell at `i, j, k` is "external" to the domain boundary.
 
 `inactive_cell`s include halo cells in `Bounded` directions, right halo cells in
-`LeftConnected` directions, left halo cells in `RightConnected` directions, and cells
+`LeftConnected` directions, left halo cells in `RightFolded` and `RightConnected` directions, and cells
 within an immersed boundary. Cells that are staggered with respect to tracer cells and
 which lie _on_ the boundary are considered active.
 """
@@ -45,11 +45,11 @@ which lie _on_ the boundary are considered active.
 
 # We use metaprogramming to handle all the permutations between
 # Bounded, LeftConnected, and RightConnected topologies.
-# Note that LeftConnected is equivalent to "RightBounded" and
-# RightConnected is equivalent to "LeftBounded".
-# So LeftConnected and RightConnected are "half Bounded" topologies.
+# Note that LeftConnected is equivalent to "RightBounded" while
+# RightFolded and RightConnected are equivalent to "LeftBounded".
+# So RightFolded, LeftConnected and RightConnected are "half Bounded" topologies.
 
-Topos = (:Bounded, :LeftConnected, :RightConnected)
+Topos = (:Bounded, :RightFolded, :LeftConnected, :RightConnected)
 
 for PrimaryTopo in Topos
 

src/Operators/Operators.jl

@@ -55,7 +55,7 @@ export intrinsic_vector, extrinsic_vector
 export σⁿ, σ⁻, ∂t_σ
 
 using Oceananigans.Grids
-using Oceananigans.Grids: LLGOTF, XRegLLGOTF, YRegLLGOTF
+using Oceananigans.Grids: LLGOTF, XRegLLGOTF, YRegLLGOTF, RightFolded
 
 #####
 ##### Convenient aliases

src/Operators/topology_aware_operators.jl

@@ -2,11 +2,13 @@ using Oceananigans.Grids: AbstractUnderlyingGrid
 
 const AGXB = AbstractUnderlyingGrid{FT, Bounded} where FT
 const AGXP = AbstractUnderlyingGrid{FT, Periodic} where FT
+const AGXF = AbstractUnderlyingGrid{FT, RightFolded} where FT
 const AGXR = AbstractUnderlyingGrid{FT, RightConnected} where FT
 const AGXL = AbstractUnderlyingGrid{FT, LeftConnected} where FT
 
 const AGYB = AbstractUnderlyingGrid{FT, <:Any, Bounded} where FT
 const AGYP = AbstractUnderlyingGrid{FT, <:Any, Periodic} where FT
+const AGYF = AbstractUnderlyingGrid{FT, <:Any, RightFolded} where FT
 const AGYR = AbstractUnderlyingGrid{FT, <:Any, RightConnected} where FT
 const AGYL = AbstractUnderlyingGrid{FT, <:Any, LeftConnected} where FT
 
@@ -43,6 +45,9 @@ const AGYL = AbstractUnderlyingGrid{FT, <:Any, LeftConnected} where FT
 @inline δxTᶠᵃᵃ(i, j, k, grid::AGXB{FT}, f, args...) where FT = ifelse(i == 1, zero(FT), δxᶠᵃᵃ(i, j, k, grid, f, args...))
 @inline δyTᵃᶠᵃ(i, j, k, grid::AGYB{FT}, f, args...) where FT = ifelse(j == 1, zero(FT), δyᵃᶠᵃ(i, j, k, grid, f, args...))
 
+@inline δxTᶠᵃᵃ(i, j, k, grid::AGXF{FT}, f, args...) where FT = ifelse(i == 1, zero(FT), δxᶠᵃᵃ(i, j, k, grid, f, args...))
+@inline δyTᵃᶠᵃ(i, j, k, grid::AGYF{FT}, f, args...) where FT = ifelse(j == 1, zero(FT), δyᵃᶠᵃ(i, j, k, grid, f, args...))
+
 @inline δxTᶠᵃᵃ(i, j, k, grid::AGXR{FT}, f, args...) where FT = ifelse(i == 1, zero(FT), δxᶠᵃᵃ(i, j, k, grid, f, args...))
 @inline δyTᵃᶠᵃ(i, j, k, grid::AGYR{FT}, f, args...) where FT = ifelse(j == 1, zero(FT), δyᵃᶠᵃ(i, j, k, grid, f, args...))
 
@@ -69,6 +74,33 @@ const AGYL = AbstractUnderlyingGrid{FT, <:Any, LeftConnected} where FT
 @inline δxTᶜᵃᵃ(i, j, k, grid::AGXL, u::AbstractArray) = @inbounds ifelse(i == grid.Nx, - u[i, j, k], δxᶜᵃᵃ(i, j, k, grid, u))
 @inline δyTᵃᶜᵃ(i, j, k, grid::AGYL, v::AbstractArray) = @inbounds ifelse(j == grid.Ny, - v[i, j, k], δyᵃᶜᵃ(i, j, k, grid, v))
 
+# Enforce Zipper conditions
+
+@inline δxTᶜᵃᵃ(i, j, k, grid::AGXF, f::Function, args...) =
+    ifelse(i == grid.Nx, - f(i, j, k, grid, args...),
+                         ifelse(i == 1, f(2, j, k, grid, args...),
+                                        δxᶜᵃᵃ(i, j, k, grid, f, args...)))
+
+
+# Changes across the fold, this operator works for `Center` fields in the x-direction.
+@inline δyTᵃᶜᵃ(i, j, k, grid::AGYF, f::Function, args...) = 
+    ifelse(j == grid.Ny, folded_δyᵃᶜᵃ(i, j, k, grid, f, args...),
+                         ifelse(j == 1, f(i, 2, k, grid, args...),
+                                        δyᵃᶜᵃ(i, j, k, grid, f, args...)))
+
+# This operator assume we are computing the difference 
+# of a vector component, which needs to switch direction across the fold
+@inline function folded_δyᵃᶜᵃ(i, j, k, grid, f, args...)
+    # Retrieve the folded index
+    i′ = grid.Nx - i + 1 
+
+    # We switch the sign of the function value at the folded index 
+    f₂ = - f(i′, j, k, grid, args...)
+    f₁ =   f(i,  j, k, grid, args...)
+
+    return f₂ - f₁
+end
+
 @inline δxTᶜᵃᵃ(i, j, k, grid::AGXR, f, args...) = ifelse(i == 1, f(2, j, k, grid, args...), δxᶜᵃᵃ(i, j, k, grid, f, args...))
 @inline δyTᵃᶜᵃ(i, j, k, grid::AGYR, f, args...) = ifelse(j == 1, f(i, 2, k, grid, args...), δyᵃᶜᵃ(i, j, k, grid, f, args...))
 

src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl

@@ -13,7 +13,7 @@ using Oceananigans.ImmersedBoundaries
 using Oceananigans.Utils
 using Oceananigans.Fields: index_binary_search, convert_to_0_360
 using Oceananigans.Grids: RightConnected
-using Oceananigans.Grids: R_Earth,
+using Oceananigans.Grids: R_Earth, RightFolded,
                           halo_size, spherical_area_quadrilateral,
                           lat_lon_to_cartesian, generate_coordinate, topology
 

src/OrthogonalSphericalShellGrids/distributed_tripolar_grid.jl

@@ -95,7 +95,16 @@ 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
+    LY = if workers[2] == 1 
+        RightFolded
+    else
+        if yrank == 0 
+            RightConnected 
+        else
+            FullyConnected
+        end
+    end
+
     LX = workers[1] == 1 ? Periodic : FullyConnected
     ny = nylocal[yrank+1]
     nx = nxlocal[xrank+1]