Polar boundary conditions

src/BoundaryConditions/BoundaryConditions.jl

@@ -37,6 +37,7 @@ include("fill_halo_regions_nothing.jl")
 include("apply_flux_bcs.jl")
 
 include("update_boundary_conditions.jl")
+include("polar_boundary_condition.jl")
 
 include("flat_extrapolation_open_boundary_matching_scheme.jl")
 end # module

src/BoundaryConditions/field_boundary_conditions.jl

@@ -59,7 +59,6 @@ end
 FieldBoundaryConditions(indices::Tuple, bcs::FieldBoundaryConditions) =
     FieldBoundaryConditions(indices, (getproperty(bcs, side) for side in propertynames(bcs))...)
 
-
 FieldBoundaryConditions(indices::Tuple, ::Nothing) = nothing
 
 window_boundary_conditions(::Colon,     left, right) = left, right
@@ -178,6 +177,13 @@ function regularize_immersed_boundary_condition(ibc, grid, loc, field_name, args
     return NoFluxBoundaryCondition()
 end
 
+  regularize_west_boundary_condition(bc, args...) = regularize_boundary_condition(bc, args...)    
+  regularize_east_boundary_condition(bc, args...) = regularize_boundary_condition(bc, args...)    
+ regularize_south_boundary_condition(bc, args...) = regularize_boundary_condition(bc, args...)   
+ regularize_north_boundary_condition(bc, args...) = regularize_boundary_condition(bc, args...)  
+regularize_bottom_boundary_condition(bc, args...) = regularize_boundary_condition(bc, args...)
+   regularize_top_boundary_condition(bc, args...) = regularize_boundary_condition(bc, args...)
+
 # regularize default boundary conditions
 function regularize_boundary_condition(default::DefaultBoundaryCondition, grid, loc, dim, args...)
     default_bc = default_prognostic_bc(topology(grid, dim)(), loc[dim](), default)
@@ -212,12 +218,12 @@ function regularize_field_boundary_conditions(bcs::FieldBoundaryConditions,
 
     loc = assumed_field_location(field_name)
 
-    west   = regularize_boundary_condition(bcs.west,   grid, loc, 1, LeftBoundary,  prognostic_names)
-    east   = regularize_boundary_condition(bcs.east,   grid, loc, 1, RightBoundary, prognostic_names)
-    south  = regularize_boundary_condition(bcs.south,  grid, loc, 2, LeftBoundary,  prognostic_names)
-    north  = regularize_boundary_condition(bcs.north,  grid, loc, 2, RightBoundary, prognostic_names)
-    bottom = regularize_boundary_condition(bcs.bottom, grid, loc, 3, LeftBoundary,  prognostic_names)
-    top    = regularize_boundary_condition(bcs.top,    grid, loc, 3, RightBoundary, prognostic_names)
+    west   = regularize_west_boundary_condition(bcs.west,     grid, loc, 1, LeftBoundary,  prognostic_names)
+    east   = regularize_east_boundary_condition(bcs.east,     grid, loc, 1, RightBoundary, prognostic_names)
+    south  = regularize_south_boundary_condition(bcs.south,   grid, loc, 2, LeftBoundary,  prognostic_names)
+    north  = regularize_north_boundary_condition(bcs.north,   grid, loc, 2, RightBoundary, prognostic_names)
+    bottom = regularize_bottom_boundary_condition(bcs.bottom, grid, loc, 3, LeftBoundary,  prognostic_names)
+    top    = regularize_top_boundary_condition(bcs.top,       grid, loc, 3, RightBoundary, prognostic_names)
 
     immersed = regularize_immersed_boundary_condition(bcs.immersed, grid, loc, field_name, prognostic_names)
 

src/BoundaryConditions/polar_boundary_condition.jl

@@ -0,0 +1,131 @@
+using Oceananigans.Grids: inactive_node, new_data
+using CUDA: @allowscalar
+
+struct PolarValue{D, S}
+    data :: D
+    side :: S
+end
+
+Adapt.adapt_structure(to, pv::PolarValue) = PolarValue(Adapt.adapt(to, pv.data), nothing)
+
+const PolarBoundaryCondition{V} = BoundaryCondition{<:Value, <:PolarValue}
+
+function PolarBoundaryCondition(grid, side, zloc)
+    FT   = eltype(grid)
+    loc  = (Nothing, Nothing, zloc)
+    data = new_data(FT, grid, loc)
+    return ValueBoundaryCondition(PolarValue(data, side))
+end
+
+# Just a column
+@inline getbc(pv::BC{<:Value, <:PolarValue}, i, k, args...) = @inbounds pv.condition.data[1, 1, k]
+
+# TODO: vectors should have a different treatment since vector components should account for the frame of reference.
+# For the moment, the `PolarBoundaryConditions` is implemented only for fields that have `loc[1] == loc[2] == Center()`, which
+# we assume are not components of horizontal vectors that would require rotation. (The `w` velocity if not a tracer, but it does
+# not require rotation since it is a scalar field.)
+# North - South flux boundary conditions are not valid on a Latitude-Longitude grid if the last / first rows represent the poles
+function latitude_north_auxiliary_bc(grid, loc, default_bc=DefaultBoundaryCondition()) 
+    # Check if the halo lies beyond the north pole
+    φmax = @allowscalar φnode(grid.Ny+1, grid, Center()) 
+
+    # Assumption: fields at `Center`s in x and y are not vector components
+    rotated_field = loc[1] != Center || loc[2] != Center
+
+    # No problem!
+    if φmax < 90 || rotated_field
+        return default_bc
+    end
+
+    return PolarBoundaryCondition(grid, :north, loc[3])
+end
+
+# North - South flux boundary conditions are not valid on a Latitude-Longitude grid if the last / first rows represent the poles
+function latitude_south_auxiliary_bc(grid, loc, default_bc=DefaultBoundaryCondition()) 
+    # Check if the halo lies beyond the south pole
+    φmin = @allowscalar φnode(0, grid, Face()) 
+
+    # Assumption: fields at `Center`s in x and y are not vector components
+    rotated_field = loc[1] != Center || loc[2] != Center
+
+    # No problem!
+    if φmin > -90 || rotated_field
+        return default_bc
+    end
+
+    return PolarBoundaryCondition(grid, :south, loc[3])
+end
+
+regularize_north_boundary_condition(bc::DefaultBoundaryCondition, grid::LatitudeLongitudeGrid, loc, args...) = 
+    regularize_boundary_condition(latitude_north_auxiliary_bc(grid, loc, bc), grid, loc, args...)
+
+regularize_south_boundary_condition(bc::DefaultBoundaryCondition, grid::LatitudeLongitudeGrid, loc, args...) = 
+    regularize_boundary_condition(latitude_south_auxiliary_bc(grid, loc, bc), grid, loc, args...)
+
+@kernel function _average_pole_value!(data, c, j, grid, loc)
+    i′, j′, k = @index(Global, NTuple)
+    c̄ = zero(grid)
+    n = 0
+    @inbounds for i in 1:grid.Nx
+        inactive = inactive_node(i, j, k, grid, loc...)
+        c̄ += ifelse(inactive, 0, c[i, j, k])
+        n += ifelse(inactive, 0, 1)
+    end
+    @inbounds data[i′, j′, k] = ifelse(n == 0,  0,  c̄ / n)
+end
+
+function update_pole_value!(bc::PolarValue, c, grid, loc) 
+    j = bc.side == :north ? grid.Ny : 1
+    Nz = size(c, 3)
+    Oz = c.offsets[3]
+    params = KernelParameters(1:1, 1:1, 1+Oz:Nz+Oz)
+    launch!(architecture(grid), grid, params, _average_pole_value!, bc.data, c, j, grid, loc)
+    return nothing
+end
+
+function fill_south_halo!(c, bc::PolarBoundaryCondition, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) 
+    update_pole_value!(bc.condition, c, grid, loc)
+    return launch!(arch, grid, KernelParameters(size, offset),
+                   _fill_only_south_halo!, c, bc, loc, grid, Tuple(args); kwargs...)
+end
+
+function fill_north_halo!(c, bc::PolarBoundaryCondition, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...) 
+    update_pole_value!(bc.condition, c, grid, loc)
+    return launch!(arch, grid, KernelParameters(size, offset),
+                   _fill_only_north_halo!, c, bc, loc, grid, Tuple(args); kwargs...)
+end
+
+function fill_south_and_north_halo!(c, south_bc::PolarBoundaryCondition, north_bc, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...)
+    update_pole_value!(south_bc.condition, c, grid, loc)
+    return launch!(arch, grid, KernelParameters(size, offset),
+                   _fill_south_and_north_halo!, c, south_bc, north_bc, loc, grid, Tuple(args); kwargs...)
+end
+
+function fill_south_and_north_halo!(c, south_bc, north_bc::PolarBoundaryCondition, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...)
+    update_pole_value!(north_bc.condition, c, grid, loc)
+    return launch!(arch, grid, KernelParameters(size, offset),
+                   _fill_south_and_north_halo!, c, south_bc, north_bc, loc, grid, Tuple(args); kwargs...)
+end
+
+function fill_south_and_north_halo!(c, south_bc::PolarBoundaryCondition, north_bc::PolarBoundaryCondition, size, offset, loc, arch, grid, args...; only_local_halos = false, kwargs...)
+    update_pole_value!(south_bc.condition, c, grid, loc)
+    update_pole_value!(north_bc.condition, c, grid, loc)
+    return launch!(arch, grid, KernelParameters(size, offset),
+                   _fill_south_and_north_halo!, c, south_bc, north_bc, loc, grid, Tuple(args); kwargs...)
+end
+
+# If it is a LatitudeLongitudeGrid, we include the PolarBoundaryConditions
+function FieldBoundaryConditions(grid::LatitudeLongitudeGrid, location, indices=(:, :, :);
+                                 west     = default_auxiliary_bc(topology(grid, 1)(), location[1]()),
+                                 east     = default_auxiliary_bc(topology(grid, 1)(), location[1]()),
+                                 south    = default_auxiliary_bc(topology(grid, 2)(), location[2]()),
+                                 north    = default_auxiliary_bc(topology(grid, 2)(), location[2]()),
+                                 bottom   = default_auxiliary_bc(topology(grid, 3)(), location[3]()),
+                                 top      = default_auxiliary_bc(topology(grid, 3)(), location[3]()),
+                                 immersed = NoFluxBoundaryCondition())
+
+    north = latitude_north_auxiliary_bc(grid, location, north)
+    south = latitude_south_auxiliary_bc(grid, location, south)
+
+    return FieldBoundaryConditions(indices, west, east, south, north, bottom, top, immersed)
+end

test/test_boundary_conditions.jl

@@ -1,6 +1,6 @@
 include("dependencies_for_runtests.jl")
 
-using Oceananigans.BoundaryConditions: PBC, ZFBC, OBC, ContinuousBoundaryFunction, DiscreteBoundaryFunction, regularize_field_boundary_conditions
+using Oceananigans.BoundaryConditions: PBC, ZFBC, VBC, OBC, ContinuousBoundaryFunction, DiscreteBoundaryFunction, regularize_field_boundary_conditions
 using Oceananigans.Fields: Face, Center
 
 simple_bc(ξ, η, t) = exp(ξ) * cos(η) * sin(t)
@@ -252,5 +252,17 @@ end
         @test T_bcs.south  === one_bc
         @test T_bcs.top    === one_bc
         @test T_bcs.bottom === one_bc
+
+        grid = LatitudeLongitudeGrid(size=(10, 10, 10), latitude=(-90, 90), longitude=(0, 360), z = (0, 1))
+        f = CenterField(grid)
+
+        @test f.boundary_conditions.north isa VBC
+        @test f.boundary_conditions.south isa VBC
+
+        set!(f, (x, y, z) -> x)
+        fill_halo_regions!(f)
+
+        @test all(f.data[1:10, 0,  1:10] .== 2 * mean(f.data[1:10, 1,  1:10]) .- f.data[1:10, 1,  1:10])
+        @test all(f.data[1:10, 11, 1:10] .== 2 * mean(f.data[1:10, 10, 1:10]) .- f.data[1:10, 10, 1:10])
     end
 end