Tripolar grid in Oceananigans (only serial)

.buildkite/pipeline.yml

@@ -211,6 +211,25 @@ steps:
           limit: 1
     depends_on: "init_cpu"
 
+#####
+##### Tripolar Grid tests
+#####
+
+  - label: "🦃 cpu tripolar grid tests"
+    env:
+      JULIA_DEPOT_PATH: "$TARTARUS_HOME/.julia-$BUILDKITE_BUILD_NUMBER"
+      TEST_GROUP: "tripolar_grid"
+    commands:
+      - "$TARTARUS_HOME/julia-$JULIA_VERSION/bin/julia -O0 --color=yes --project -e 'using Pkg; Pkg.test()'"
+    agents:
+      queue: Oceananigans
+      architecture: CPU
+    retry:
+      automatic:
+        - exit_status: 1      
+          limit: 1
+    depends_on: "init_cpu"
+
 #####
 ##### NonhydrostaticModel and time stepping (part 1)
 #####

docs/src/appendix/library.md

@@ -86,6 +86,13 @@ Modules = [Oceananigans.Grids]
 Private = false
 ```
 
+## OrthogonalSphericalShellGrids
+
+```@autodocs
+Modules = [Oceananigans.OrthogonalSphericalShellGrids]
+Private = false
+```
+
 ## Immersed boundaries
 
 ```@autodocs

docs/src/grids.md

@@ -88,7 +88,7 @@ The shape of the physical domain determines what grid type should be used:
 3. [`OrthogonalSphericalShellGrid`](@ref Oceananigans.Grids.OrthogonalSphericalShellGrid) represents sectors of thin spherical shells divided with mesh lines that intersect at right angles (thus, orthogonal) but are otherwise arbitrary.
 
 !!! note "OrthogonalSphericalShellGrids.jl"
-    See the auxiliary package [`OrthogonalSphericalShellGrids.jl`](https://github.com/CliMA/OrthogonalSphericalShellGrids.jl)
+    See the auxiliary module [`OrthogonalSphericalShellGrids.jl`](@ref Oceananigans.OrthogonalSphericalShellGrids)
     for recipes that implement some useful `OrthogonalSphericalShellGrid`, including the
     ["tripolar" grid](https://www.sciencedirect.com/science/article/abs/pii/S0021999196901369).
 

src/BoundaryConditions/BoundaryConditions.jl

@@ -33,6 +33,7 @@ include("fill_halo_regions_open.jl")
 include("fill_halo_regions_periodic.jl")
 include("fill_halo_regions_flux.jl")
 include("fill_halo_regions_nothing.jl")
+include("fill_halo_regions_zipper.jl")
 
 include("apply_flux_bcs.jl")
 

src/BoundaryConditions/apply_flux_bcs.jl

@@ -90,6 +90,9 @@ end
 @inline    apply_z_top_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
 @inline apply_z_bottom_bc!(Gc, loc, ::NotFluxBC, args...) = nothing
 
+# shortcut for the zipper BC
+@inline apply_y_north_bc!(Gc, loc, ::ZBC, args...) = nothing
+
 @inline flip(::Center) = Face()
 @inline flip(::Face) = Center()
 

src/BoundaryConditions/boundary_condition.jl

@@ -82,20 +82,23 @@ const GBC  = BoundaryCondition{<:Gradient}
 const ZFBC = BoundaryCondition{Flux, Nothing} # "zero" flux
 const MCBC = BoundaryCondition{<:MultiRegionCommunication}
 const DCBC = BoundaryCondition{<:DistributedCommunication}
+const ZBC  = BoundaryCondition{<:Zipper}
 
 # More readable BC constructors for the public API.
                 PeriodicBoundaryCondition() = BoundaryCondition(Periodic(),                 nothing)
                   NoFluxBoundaryCondition() = BoundaryCondition(Flux(),                     nothing)
             ImpenetrableBoundaryCondition() = BoundaryCondition(Open(), nothing)
 MultiRegionCommunicationBoundaryCondition() = BoundaryCondition(MultiRegionCommunication(), nothing)
 DistributedCommunicationBoundaryCondition() = BoundaryCondition(DistributedCommunication(), nothing)
+            ZipperBoundaryCondition(sign=1) = BoundaryCondition(Zipper(), sign)
 
                     FluxBoundaryCondition(val; kwargs...) = BoundaryCondition(Flux(), val; kwargs...)
                    ValueBoundaryCondition(val; kwargs...) = BoundaryCondition(Value(), val; kwargs...)
                 GradientBoundaryCondition(val; kwargs...) = BoundaryCondition(Gradient(), val; kwargs...)
                     OpenBoundaryCondition(val; kwargs...) = BoundaryCondition(Open(nothing), val; kwargs...)
 MultiRegionCommunicationBoundaryCondition(val; kwargs...) = BoundaryCondition(MultiRegionCommunication(), val; kwargs...)
 DistributedCommunicationBoundaryCondition(val; kwargs...) = BoundaryCondition(DistributedCommunication(), val; kwargs...)
+                  ZipperBoundaryCondition(val; kwargs...) = BoundaryCondition(Zipper(), val; kwargs...)
 
 # Support for various types of boundary conditions.
 #

src/BoundaryConditions/boundary_condition_classifications.jl

@@ -82,3 +82,11 @@ struct MultiRegionCommunication <: AbstractBoundaryConditionClassification end
 A classification specifying a distributed memory communicating boundary condition 
 """
 struct DistributedCommunication <: AbstractBoundaryConditionClassification end
+
+"""
+    struct Zipper <: AbstractBoundaryConditionClassification
+
+A classification specifying a Zipper boundary condition where one boundary is folded onto itself.
+Used only for a tripolar grid.
+"""
+struct Zipper <: AbstractBoundaryConditionClassification end

src/BoundaryConditions/fill_halo_regions_zipper.jl

@@ -0,0 +1,134 @@
+using Oceananigans.Grids: Center, Face
+using Oceananigans.Utils: KernelParameters, launch!
+
+"""
+    ZipperBoundaryCondition(sign = 1)
+
+Create a zipper boundary condition specific to the `TripolarGrid`.
+A Zipper boundary condition is similar to a periodic boundary condition, but, 
+instead of retrieving the value from the opposite boundary, it splits the boundary 
+in two and retrieves the value from the opposite side of the boundary. 
+It is possible to think of it as a periodic boundary over a folded domain.
+
+When copying in halos, folded velocities need to switch sign, while tracers or similar fields do not.
+
+Note: the Tripolar boundary condition is particular because the last grid point 
+at the north edge is repeated for tracers. For this reason we do not start copying the halo 
+from the last grid point but from the second to last grid point.
+
+Example
+=======
+
+Consider the northern edge of a tripolar grid where P indicates the j - location of the poles
+
+```
+                  P                         P
+                  |            |            |            |            
+ Ny (center) ->   u₁    c₁     u₂    c₂     u₃    c₃     u₄    c₄     
+ Ny (face)   ->   |---- v₁ ----|---- v₂ ----|---- v₃ ----|---- v₄ ----
+                                                      Nx    Nx
+                                                    (face) (center)
+```
+The grid ends at `Ny` because it is periodic in nature, but, given the fold, 
+```
+c₁ == c₄
+```
+
+and
+```
+c₂ == c₃
+```
+
+This is not the case for the v-velocity (or any field on the j-faces) where the last grid point is not repeated.
+"""
+
+#####
+##### Outer functions for filling halo regions for Zipper boundary conditions.
+#####
+
+@inline function fold_north_face_face!(i, k, grid, sign, c)
+    Nx, Ny, _ = size(grid)
+
+    i′ = Nx - i + 2 # Remember! element Nx + 1 does not exist!
+    sign = ifelse(i′ > Nx , abs(sign), sign) # for periodic elements we change the sign
+    i′ = ifelse(i′ > Nx, i′ - Nx, i′) # Periodicity is hardcoded in the x-direction!!
+    Hy = grid.Hy
+
+    for j = 1 : Hy
+        @inbounds begin
+            c[i, Ny + j, k] = sign * c[i′, Ny - j + 1, k] 
+        end
+    end
+
+    return nothing
+end
+
+@inline function fold_north_face_center!(i, k, grid, sign, c)
+    Nx, Ny, _ = size(grid)
+
+    i′ = Nx - i + 2 # Remember! element Nx + 1 does not exist!
+    sign  = ifelse(i′ > Nx , abs(sign), sign) # for periodic elements we change the sign
+    i′ = ifelse(i′ > Nx, i′ - Nx, i′) # Periodicity is hardcoded in the x-direction!!
+    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
+    end
+
+    # We substitute the redundant part of the last row to ensure consistency
+    @inbounds c[i, Ny, k] = ifelse(i > Nx ÷ 2, sign * c[i′, Ny, k], c[i, Ny, k])
+
+    return nothing
+end
+
+@inline function fold_north_center_face!(i, k, grid, sign, c)
+    Nx, Ny, _ = size(grid)
+
+    i′ = Nx - i + 1
+    Hy = grid.Hy
+
+    for j = 1 : Hy
+        @inbounds begin
+            c[i, Ny + j, k] = sign * c[i′, Ny - j + 1, k] 
+        end
+    end
+
+    return nothing
+end
+
+@inline function fold_north_center_center!(i, k, grid, sign, c)
+    Nx, Ny, _ = size(grid)
+
+    i′ = Nx - i + 1
+    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
+    end
+
+    # We substitute the redundant part of the last row to ensure consistency
+    @inbounds c[i, Ny, k] = ifelse(i > Nx ÷ 2, sign * c[i′, Ny, k], c[i, Ny, k])
+
+    return nothing
+end
+
+const CCLocation = Tuple{<:Center, <:Center, <:Any} 
+const FCLocation = Tuple{<:Face,   <:Center, <:Any} 
+const CFLocation = Tuple{<:Center, <:Face,   <:Any} 
+const FFLocation = Tuple{<:Face,   <:Face,   <:Any} 
+
+# tracers or similar fields
+@inline _fill_north_halo!(i, k, grid, c, bc::ZBC, ::CCLocation, args...) = fold_north_center_center!(i, k, grid, bc.condition, c)
+
+# u-velocity or similar fields
+@inline _fill_north_halo!(i, k, grid, u, bc::ZBC, ::FCLocation, args...) = fold_north_face_center!(i, k, grid, bc.condition, u)
+
+# v-velocity or similar fields
+@inline _fill_north_halo!(i, k, grid, v, bc::ZBC, ::CFLocation, args...) = fold_north_center_face!(i, k, grid, bc.condition, v)
+
+# vorticity or similar fields
+@inline _fill_north_halo!(i, k, grid, ζ, bc::ZBC, ::FFLocation, args...) = fold_north_face_face!(i, k, grid, bc.condition, ζ)

src/BoundaryConditions/show_boundary_conditions.jl

@@ -14,6 +14,7 @@ bc_str(::IBC)                    = "Impenetrable"
 bc_str(::DFBC)                   = "Default"
 bc_str(::MCBC)                   = "MultiRegionCommunication"
 bc_str(::DCBC)                   = "DistributedCommunication"
+bc_str(::ZBC)                    = "Zipper"
 bc_str(::Nothing)                = "Nothing"
 
 #####
@@ -28,6 +29,7 @@ Base.summary(bc::VBC)                     = string("ValueBoundaryCondition: ", p
 Base.summary(bc::GBC)                     = string("GradientBoundaryCondition: ", prettysummary(bc.condition))
 Base.summary(::PBC)                       = string("PeriodicBoundaryCondition")
 Base.summary(bc::DCBC)                    = string("DistributedBoundaryCondition: ", prettysummary(bc.condition))
+Base.summary(bc::ZBC)                     = string("ZipperBoundaryCondition: ", prettysummary(bc.condition))
 
 show(io::IO, bc::BoundaryCondition) = print(io, summary(bc))
 

src/Fields/field.jl

@@ -1,4 +1,4 @@
-using Oceananigans.BoundaryConditions: OBC, MCBC, BoundaryCondition
+using Oceananigans.BoundaryConditions: OBC, MCBC, BoundaryCondition, Zipper
 using Oceananigans.Grids: parent_index_range, index_range_offset, default_indices, all_indices, validate_indices
 using Oceananigans.Grids: index_range_contains
 
@@ -80,6 +80,14 @@ function validate_boundary_conditions(loc, grid, bcs)
     return nothing
 end
 
+# Some special validation for a zipper boundary condition
+validate_boundary_condition_location(bc::Zipper, loc::Center, side) = 
+    side == :north ? nothing : throw(ArgumentError("Cannot specify $side boundary condition $bc on a field at $(loc) (north only)!"))
+
+validate_boundary_condition_location(bc::Zipper, loc::Face, side) = 
+    side == :north ? nothing : throw(ArgumentError("Cannot specify $side boundary condition $bc on a field at $(loc) (north only)!"))
+
+
 #####
 ##### Some basic constructors
 #####

src/Oceananigans.jl

@@ -11,7 +11,7 @@ export
     # Grids
     Center, Face,
     Periodic, Bounded, Flat,
-    RectilinearGrid, LatitudeLongitudeGrid, OrthogonalSphericalShellGrid,
+    RectilinearGrid, LatitudeLongitudeGrid, OrthogonalSphericalShellGrid, TripolarGrid,
     nodes, xnodes, ynodes, rnodes, znodes, λnodes, φnodes,
     xspacings, yspacings, rspacings, zspacings, λspacings, φspacings,
     minimum_xspacing, minimum_yspacing, minimum_zspacing,
@@ -208,12 +208,13 @@ include("Operators/Operators.jl")
 include("BoundaryConditions/BoundaryConditions.jl")
 include("Fields/Fields.jl")
 include("AbstractOperations/AbstractOperations.jl")
-include("TimeSteppers/TimeSteppers.jl")
 include("ImmersedBoundaries/ImmersedBoundaries.jl")
+include("TimeSteppers/TimeSteppers.jl")
 include("Advection/Advection.jl")
 include("Solvers/Solvers.jl")
 include("OutputReaders/OutputReaders.jl")
 include("DistributedComputations/DistributedComputations.jl")
+include("OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl")
 
 # TODO: move here
 #include("MultiRegion/MultiRegion.jl")
@@ -257,6 +258,7 @@ using .OutputReaders
 using .Forcings
 using .ImmersedBoundaries
 using .DistributedComputations
+using .OrthogonalSphericalShellGrids
 using .Models
 using .TimeSteppers
 using .Diagnostics

src/OrthogonalSphericalShellGrids/OrthogonalSphericalShellGrids.jl

@@ -0,0 +1,29 @@
+module OrthogonalSphericalShellGrids
+
+# The only thing we need!
+export TripolarGrid
+
+using Oceananigans
+using Oceananigans.Grids
+using Oceananigans.Architectures: device, on_architecture
+using Oceananigans.BoundaryConditions
+using Oceananigans.Fields: index_binary_search, convert_to_0_360
+using Oceananigans.Grids: RightConnected
+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 Distances
+using Adapt 
+using KernelAbstractions: @kernel, @index
+using KernelAbstractions.Extras.LoopInfo: @unroll
+using OffsetArrays
+
+include("generate_tripolar_coordinates.jl")
+include("tripolar_grid.jl")
+include("tripolar_field_extensions.jl")
+
+end