(0.94.4) PolarBoundaryCondition for tracer + views for AbstractOperations

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.94.3"
+version = "0.94.4"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/AbstractOperations/AbstractOperations.jl

@@ -21,6 +21,7 @@ using Oceananigans.Operators: interpolation_operator
 using Oceananigans.Architectures: device
 
 import Adapt
+import Base
 
 import Oceananigans.Architectures: architecture, on_architecture
 import Oceananigans.BoundaryConditions: fill_halo_regions!
@@ -42,6 +43,10 @@ architecture(a::AbstractOperation) = architecture(a.grid)
 # AbstractOperation macros add their associated functions to this list
 const operators = Set()
 
+# To allow view(o::AbstractOperation, i, j, k) to work
+@inline get_field_view(a::Union{AbstractOperation, AbstractField}, i, j, k) = view(a, i, j, k)
+@inline get_field_view(a, i, j, k) = a
+
 """
     at(loc, abstract_operation)
 

src/AbstractOperations/at.jl

@@ -77,6 +77,7 @@ compute_index_intersection(::Colon, to_loc; kw...) = Colon()
 
 compute_index_intersection(to_idx, to_loc, op; dim) =
     _compute_index_intersection(to_idx, indices(op)[dim],
+                                topology(op.grid)[dim],
                                 to_loc, location(op, dim))
 
 """Compute index intersection recursively for `dim`ension ∈ (1, 2, 3)."""
@@ -85,23 +86,25 @@ function compute_index_intersection(to_idx, to_loc, op1, op2, more_ops...; dim)
     return compute_index_intersection(new_to_idx, to_loc, op2, more_ops...; dim)
 end
 
+const Range = Union{UnitRange, Base.OneTo, Base.Slice{<:Base.OneTo}}
+
 # Life is pretty simple in this case.
-_compute_index_intersection(to_idx::Colon, from_idx::Colon, args...) = Colon()
+_compute_index_intersection(::Colon, ::Colon, args...) = Colon()
 
 # Because `from_idx` imposes no restrictions, we just return `to_idx`.
-_compute_index_intersection(to_idx::UnitRange, from_idx::Colon, args...) = to_idx
+_compute_index_intersection(to_idx::Range, ::Colon, args...) = to_idx
 
 # This time we account for the possible range-reducing effect of interpolation on `from_idx`.
-function _compute_index_intersection(to_idx::Colon, from_idx::UnitRange, to_loc, from_loc)
+function _compute_index_intersection(::Colon, from_idx::Range, to_loc, from_loc)
     shifted_idx = restrict_index_for_interpolation(from_idx, from_loc, to_loc)
-    validate_shifted_index(shifted_idx)
+    validate_shifted_index(shifted_idx, from_idx, from_loc, to_loc)
     return shifted_idx
 end
 
 # Compute the intersection of two index ranges
-function _compute_index_intersection(to_idx::UnitRange, from_idx::UnitRange, to_loc, from_loc)
+function _compute_index_intersection(to_idx::Range, from_idx::Range, to_loc, from_loc)
     shifted_idx = restrict_index_for_interpolation(from_idx, from_loc, to_loc)
-    validate_shifted_index(shifted_idx)
+    validate_shifted_index(shifted_idx, from_idx, from_loc, to_loc)
 
     range_intersection = UnitRange(max(first(shifted_idx), first(to_idx)), min(last(shifted_idx), last(to_idx)))
 
@@ -112,7 +115,7 @@ function _compute_index_intersection(to_idx::UnitRange, from_idx::UnitRange, to_
     return range_intersection
 end
 
-validate_shifted_index(shifted_idx) = first(shifted_idx) > last(shifted_idx) &&
+validate_shifted_index(shifted_idx, from_idx, from_loc, to_loc) = first(shifted_idx) > last(shifted_idx) &&
     throw(ArgumentError("Cannot compute index intersection for indices $(from_idx) interpolating from $(from_loc) to $(to_loc)!"))
 
 """
@@ -128,3 +131,10 @@ restrict_index_for_interpolation(from_idx, ::Type{Face},   ::Type{Face})   = Uni
 restrict_index_for_interpolation(from_idx, ::Type{Center}, ::Type{Center}) = UnitRange(first(from_idx),   last(from_idx))
 restrict_index_for_interpolation(from_idx, ::Type{Face},   ::Type{Center}) = UnitRange(first(from_idx),   last(from_idx)-1)
 restrict_index_for_interpolation(from_idx, ::Type{Center}, ::Type{Face})   = UnitRange(first(from_idx)+1, last(from_idx))
+
+# No restrictions for interpolating from `Nothing` or to `Nothing`
+restrict_index_for_interpolation(from_idx, ::Type{Nothing}, ::Type{Face})    = from_idx
+restrict_index_for_interpolation(from_idx, ::Type{Nothing}, ::Type{Center})  = from_idx
+restrict_index_for_interpolation(from_idx, ::Type{Face},    ::Type{Nothing}) = from_idx
+restrict_index_for_interpolation(from_idx, ::Type{Center},  ::Type{Nothing}) = from_idx
+restrict_index_for_interpolation(from_idx, ::Type{Nothing}, ::Type{Nothing}) = from_idx

src/AbstractOperations/binary_operations.jl

@@ -225,3 +225,14 @@ on_architecture(to, binary::BinaryOperation{LX, LY, LZ}) where {LX, LY, LZ} =
                                 on_architecture(to, binary.▶a),
                                 on_architecture(to, binary.▶b),
                                 on_architecture(to, binary.grid))
+
+# Extending views for `BinaryOperation`s
+function Base.view(binary::BinaryOperation{LX, LY, LZ}, i, j, k) where {LX, LY, LZ}
+    # Propagate view over all arguments
+    view_a = get_field_view(binary.a, i, j, k) 
+    view_b = get_field_view(binary.b, i, j, k) 
+    view_▶a = get_field_view(binary.▶a, i, j, k) 
+    view_▶b = get_field_view(binary.▶b, i, j, k) 
+
+    return BinaryOperation{LX, LY, LZ}(binary.op, view_a, view_b, view_▶a, view_▶b, binary.grid)
+end

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
@@ -159,3 +162,10 @@ Base.show(io::IO, operation::ConditionalOperation) =
           "├── func: ", summary(operation.func), "\n",
           "├── condition: ", summary(operation.condition), "\n",
           "└── mask: ", operation.mask)
+
+# Extending views for `ConditionalyOperation`s
+function Base.view(co::ConditionalOperation{LX, LY, LZ}, i, j, k) where {LX, LY, LZ}
+    # Propagate view over all arguments
+    view_operand = get_field_view(co.operand, i, j, k) 
+    return ConditionalOperation{LX, LY, LZ}(view_operand, co.func, co.grid, co.condition, co.mask)
+end

src/AbstractOperations/derivatives.jl

@@ -132,4 +132,12 @@ on_architecture(to, deriv::Derivative{LX, LY, LZ}) where {LX, LY, LZ} =
                            on_architecture(to, deriv.▶),
                            deriv.abstract_∂,
                            on_architecture(to, deriv.grid))
-
+
+# Extending views for `Derivative`s
+function Base.view(deriv::Derivative{LX, LY, LZ}, i, j, k) where {LX, LY, LZ}
+    # Propagate view over all arguments
+    view_arg = get_field_view(deriv.arg, i, j, k)
+    view_▶ =   get_field_view(deriv.▶, i, j, k)
+
+    return Derivative{LX, LY, LZ}(deriv.∂, view_arg, view_▶, deriv.grid)
+end

src/AbstractOperations/grid_metrics.jl

@@ -145,6 +145,9 @@ indices(::GridMetricOperation) = default_indices(3)
 # Special constructor for BinaryOperation
 GridMetricOperation(L, metric, grid) = GridMetricOperation{L[1], L[2], L[3]}(metric_function(L, metric), grid)
 
+# Extend view on a GridMetricOperation (we return the same object)
+Base.view(gm::GridMetricOperation, i, j, k) = gm
+
 #####
 ##### Spacings
 #####

src/AbstractOperations/kernel_function_operation.jl

@@ -97,3 +97,10 @@ Base.show(io::IO, kfo::KernelFunctionOperation) =
                              prettysummary(kfo.arguments[end])
                          end
 )
+
+# Extending views for KernelFunctionOperation
+function Base.view(ko::KernelFunctionOperation{LX, LY, LZ}, i, j, k) where {LX, LY, LZ}
+    # Propagate view over all arguments
+    view_args = [get_field_view(a, i, j, k) for a in ko.arguments]
+    return KernelFunctionOperation{LX, LY, LZ}(ko.kernel_function, ko.grid, view_args...)
+end

src/AbstractOperations/multiary_operations.jl

@@ -156,4 +156,12 @@ on_architecture(to, multiary::MultiaryOperation{LX, LY, LZ}) where {LX, LY, LZ}
                                   on_architecture(to, multiary.args),
                                   on_architecture(to, multiary.▶),
                                   on_architecture(to, multiary.grid))
-
+
+# Extending views for `MultiaryOperation`s
+function Base.view(multiary::MultiaryOperation{LX, LY, LZ}, i, j, k) where {LX, LY, LZ}
+    # Propagate view over all arguments
+    view_args = Tuple(get_field_view(arg, i, j, k) for args in multiary.args)
+    view_▶ =    Tuple(get_field_view(▶, i, j, k) for ▶ in multiary.▶)
+
+    return MultiaryOperation{LX, LY, LZ}(multiary.op, view_args, view_▶, multiary.grid)
+end

src/AbstractOperations/unary_operations.jl

@@ -136,3 +136,12 @@ on_architecture(to, unary::UnaryOperation{LX, LY, LZ}) where {LX, LY, LZ} =
                                on_architecture(to, unary.arg),
                                on_architecture(to, unary.▶),
                                on_architecture(to, unary.grid))
+
+# Extending views for `UnaryOperation`s
+function Base.view(unary::UnaryOperation{LX, LY, LZ}, i, j, k) where {LX, LY, LZ}
+    # Propagate view over all arguments
+    view_arg = get_field_view(unary.arg, i, j, k)
+    view_▶ =   get_field_view(unary.▶, i, j, k)
+
+    return UnaryOperation{LX, LY, LZ}(unary.op, view_arg, view_▶, unary.grid)
+end

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

@@ -62,8 +62,10 @@ FieldBoundaryConditions(indices::Tuple, bcs::FieldBoundaryConditions) =
 
 FieldBoundaryConditions(indices::Tuple, ::Nothing) = nothing
 
-window_boundary_conditions(::Colon,     left, right) = left, right
-window_boundary_conditions(::UnitRange, left, right) = nothing, nothing
+const Range = Union{UnitRange, Base.OneTo, Base.Slice{<:Base.OneTo}}
+
+window_boundary_conditions(::Colon, left, right) = left, right
+window_boundary_conditions(::Range, left, right) = nothing, nothing
 
 on_architecture(arch, fbcs::FieldBoundaryConditions) =
     FieldBoundaryConditions(on_architecture(arch, fbcs.west),
@@ -178,6 +180,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 +221,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,121 @@
+using Oceananigans.Architectures: device
+using Oceananigans.Grids: inactive_node
+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)
+    FT = eltype(grid)
+    arch = architecture(grid)
+    data = on_architecture(arch, zeros(FT, grid.Nz))
+    return ValueBoundaryCondition(PolarValue(data, side))
+end
+
+# Just a column
+@inline getbc(pv::BC{<:Value, <:PolarValue}, i, k, args...) = @inbounds pv.condition.data[k]
+
+# 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 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()) 
+
+    # No problem!
+    if φmax < 90 || loc[1] == Nothing
+        return default_bc
+    end
+
+    return PolarBoundaryCondition(grid, :north)
+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()) 
+
+    # No problem!
+    if φmin > -90 || loc[1] == Nothing
+        return default_bc
+    end
+
+    return PolarBoundaryCondition(grid, :south)
+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)
+    k = @index(Global, Linear)
+    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[k] = ifelse(n == 0,  0,  c̄ / n)
+end
+
+function update_pole_value!(bc::PolarValue, c, grid, loc) 
+    j = bc.side == :north ? grid.Ny : 1
+    dev = device(architecture(grid))
+    _average_pole_value!(dev, 16, grid.Nz)(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

src/Fields/abstract_field.jl

@@ -58,10 +58,12 @@ Base.parent(f::AbstractField) = f
 const Abstract3DField = AbstractField{<:Any, <:Any, <:Any, <:Any, <:Any, 3}
 const Abstract4DField = AbstractField{<:Any, <:Any, <:Any, <:Any, <:Any, 4}
 
+const Range = Union{UnitRange, Base.OneTo, Base.Slice{<:Base.OneTo}}
+
 # TODO: to omit boundaries on Face fields, we have to return 2:N
 # when topo=Bounded, and loc=Face
 @inline axis(::Colon, N) = Base.OneTo(N)
-@inline axis(index::UnitRange, N) = index
+@inline axis(index::Range, N) = index
 
 @inline function Base.axes(f::Abstract3DField)
     Nx, Ny, Nz = size(f)

src/Fields/constant_field.jl

@@ -26,3 +26,7 @@ const CF = Union{ConstantField, ZeroField, OneField}
 fill_halo_regions!(::ZeroField, args...; kw...) = nothing
 fill_halo_regions!(::ConstantField, args...; kw...) = nothing
 
+# Views of OneField, ZeroField, ConstantField are themselves
+Base.view(f::OneField,      args...) = f
+Base.view(f::ZeroField,     args...) = f
+Base.view(f::ConstantField, args...) = f