PolarBoundaryCondition for tracer + views for AbstractOperations

src/AbstractOperations/conditional_operations.jl

@@ -1,9 +1,12 @@
 using Oceananigans.Fields: OneField
 using Oceananigans.Grids: architecture
 
+using OffsetArrays
+
 import Oceananigans.Architectures: on_architecture
 import Oceananigans.Fields: condition_operand, conditional_length, set!, compute_at!, indices
 
+
 # For conditional reductions such as mean(u * v, condition = u .> 0))
 struct ConditionalOperation{LX, LY, LZ, O, F, G, C, M, T} <: AbstractOperation{LX, LY, LZ, G, T}
     operand :: O
@@ -115,6 +118,14 @@ end
     return ConditionalOperation(operand; func, condition, mask)
 end
 
+@inline function condition_operand(data::OffsetArray, grid::AbstractGrid, Loc::Tuple, condition, mask)
+    condition = on_architecture(architecture(grid), condition)    
+    return ConditionalOperation{Loc...}(operand, identity, grid, condition, mask)
+end
+
+condition_operand(data::OffsetArray, grid::AbstractGrid, Loc::Tuple, ::Nothing, mask) = 
+    condition_operand(data, grid, Loc, TrueCondition(), mask)
+
 @inline condition_operand(func::typeof(identity), c::ConditionalOperation, ::Nothing, mask) = ConditionalOperation(c; mask)
 @inline condition_operand(func::Function,         c::ConditionalOperation, ::Nothing, mask) = ConditionalOperation(c; func, mask)
 

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

@@ -178,6 +178,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 +219,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,76 @@
+const PolarBoundaryCondition{V} = BoundaryCondition{<:Value, <:PolarValue}
+
+condition_operand(data, grid, Loc, condition, mask) = data
+
+struct PolarValue{C}
+    c :: C
+end
+
+PolarBoundaryCondition(field) = 
+    ValueBoundaryCondition(PolarValue(field))
+
+@inline getbc(pv::PolarValue, args...) = getbc(pv.c, args...)
+
+# TODO: vectors should have a different treatment since vector components should account for the frame of reference
+# North - South flux boundary conditions are not valid on a Latitude-Longitude grid if the last / first rows represent the poles
+function regularize_north_boundary_condition(bc::DefaultBoundaryCondition, grid::LatitudeLongitudeGrid, loc, dim, args...) where C
+    φmax = φnode(grid.Ny+1, grid, Face()) 
+
+    if φmax == 90 
+        _, LY, LZ = loc
+        field = Field{Nothing, LY, LZ}(grid; indices = (:, grid.Ny, :))
+        return PolarBoundaryCondition(field)
+    else
+        return regularize_boundary_condition(bc, grid, args...)
+    end
+end
+
+# North - South flux boundary conditions are not valid on a Latitude-Longitude grid if the last / first rows represent the poles
+function regularize_south_boundary_condition(bc::DefaultBoundaryCondition, grid::LatitudeLongitudeGrid, loc, dim, args...) 
+    φmin = φnode(1, grid, Face()) 
+
+    if φmin == - 90 
+        _, LY, LZ = loc
+        field = Field{Nothing, LY, LZ}(grid; indices = (:, 1, :))
+        return PolarBoundaryCondition(field)
+    else
+        return regularize_boundary_condition(bc, grid, args...)
+    end
+end
+
+function update_pole_value!(bc::PolarBoundaryCondition, c, grid, loc) 
+    operand = condition_operand(c, grid, loc, nothing, 0)
+    mean!(bc.condition.c, operand)
+    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, 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, 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, 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, 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, c, grid, loc)
+    update_pole_value!(north_bc, 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

src/BoundaryConditions/update_boundary_conditions.jl

@@ -25,3 +25,4 @@ function update_boundary_condition!(fields::Tuple, model)
 
     return nothing
 end
+

src/Fields/field.jl

@@ -9,6 +9,7 @@ using Base: @propagate_inbounds
 import Oceananigans: boundary_conditions
 import Oceananigans.Architectures: on_architecture
 import Oceananigans.BoundaryConditions: fill_halo_regions!, getbc
+import Oceananigans.BoundaryConditions: condition_operand
 import Statistics: norm, mean, mean!
 import Base: ==
 

src/ImmersedBoundaries/immersed_reductions.jl

@@ -28,6 +28,12 @@ const IF = AbstractField{<:Any, <:Any, <:Any, <:ImmersedBoundaryGrid}
     return ConditionalOperation(op; func, condition=ni_condition, mask)
 end
 
+@inline function condition_operand(data::OffsetArray, grid::AbstractGrid, Loc, condition, mask)
+    condition = on_architecture(architecture(grid), condition)    
+    ni_condition = NotImmersed(condition)
+    return ConditionalOperation{Loc...}(operand, identity, grid, condition, mask)
+end
+
 @inline conditional_length(c::IF)       = conditional_length(condition_operand(identity, c, nothing, 0))
 @inline conditional_length(c::IF, dims) = conditional_length(condition_operand(identity, c, nothing, 0), dims)