Adds a pertubation advection open boundary matching scheme

src/BoundaryConditions/BoundaryConditions.jl

@@ -40,4 +40,5 @@ include("update_boundary_conditions.jl")
 include("polar_boundary_condition.jl")
 
 include("flat_extrapolation_open_boundary_matching_scheme.jl")
+include("perturbation_advection_open_boundary_matching_scheme.jl")
 end # module

src/BoundaryConditions/perturbation_advection_open_boundary_matching_scheme.jl

@@ -0,0 +1,222 @@
+using Oceananigans.Operators: Δxᶠᶜᶜ, Δyᶜᶠᶜ, Δzᶜᶜᶠ, Ax_qᶠᶜᶜ, Ay_qᶜᶠᶜ, Az_qᶜᶜᶠ
+
+"""
+    PerturbationAdvection
+
+For cases where we assume that the internal flow is a small perturbation from 
+an external prescribed or coarser flow, we can split the velocity into background
+and perturbation components:
+...
+see latex document for now
+
+TODO: check what the coriolis is doing, and check what happens if U is the mean velocity
+"""
+struct PerturbationAdvection{VT, FT}
+       backward_step :: VT
+    inflow_timescale :: FT
+   outflow_timescale :: FT
+end
+
+Adapt.adapt_structure(to, pe::PerturbationAdvection) = 
+    PerturbationAdvection(adapt(to, pe.backward_step),
+                          adapt(to, pe.outflow_timescale),
+                          adapt(to, pe.inflow_timescale))
+
+function PerturbationAdvectionOpenBoundaryCondition(val, FT = Float64; 
+                                                    backward_step = true,
+                                                    outflow_timescale = Inf, 
+                                                    inflow_timescale = 300.0, kwargs...)
+
+    classification = Open(PerturbationAdvection(Val(backward_step), inflow_timescale, outflow_timescale))
+
+    @warn "`PerturbationAdvection` open boundaries matching scheme is experimental and un-tested/validated"
+
+    return BoundaryCondition(classification, val; kwargs...)
+end
+
+const PAOBC = BoundaryCondition{<:Open{<:PerturbationAdvection}}
+
+const BPAOBC = BoundaryCondition{<:Open{<:PerturbationAdvection{Val{true}}}}
+const FPAOBC = BoundaryCondition{<:Open{<:PerturbationAdvection{Val{false}}}}
+
+@inline function step_right_boundary!(bc::BPAOBC, l, m, boundary_indices, boundary_adjacent_indices, 
+                                      grid, u, clock, model_fields, ΔX)
+    Δt = clock.last_stage_Δt
+
+    Δt = ifelse(isinf(Δt), 0, Δt)
+
+    ūⁿ⁺¹ = getbc(bc, l, m, grid, clock, model_fields)
+
+    uᵢⁿ     = @inbounds getindex(u, boundary_indices...)
+    uᵢ₋₁ⁿ⁺¹ = @inbounds getindex(u, boundary_adjacent_indices...)
+
+    U = max(0, min(1, Δt / ΔX * ūⁿ⁺¹))
+
+    pa = bc.classification.matching_scheme
+
+    τ = ifelse(ūⁿ⁺¹ >= 0, pa.outflow_timescale, pa.inflow_timescale)
+
+    τ̃ = Δt / τ
+
+    uᵢⁿ⁺¹ = (uᵢⁿ + U * uᵢ₋₁ⁿ⁺¹ + ūⁿ⁺¹ * τ̃) / (1 + τ̃ + U)
+
+    @inbounds setindex!(u, uᵢⁿ⁺¹, boundary_indices...)
+
+    return nothing
+end
+
+@inline function step_left_boundary!(bc::BPAOBC, l, m, boundary_indices, boundary_adjacent_indices, boundary_secret_storage_indices, 
+                                     grid, u, clock, model_fields, ΔX)
+    Δt = clock.last_stage_Δt
+
+    Δt = ifelse(isinf(Δt), 0, Δt)
+
+    ūⁿ⁺¹ = getbc(bc, l, m, grid, clock, model_fields)
+
+    uᵢⁿ     = @inbounds getindex(u, boundary_secret_storage_indices...)
+    uᵢ₋₁ⁿ⁺¹ = @inbounds getindex(u, boundary_adjacent_indices...)
+
+    U = min(0, max(-1, Δt / ΔX * ūⁿ⁺¹))
+
+    pa = bc.classification.matching_scheme
+
+    τ = ifelse(ūⁿ⁺¹ <= 0, pa.outflow_timescale, pa.inflow_timescale)
+
+    τ̃ = Δt / τ
+
+    u₁ⁿ⁺¹ = (uᵢⁿ - U * uᵢ₋₁ⁿ⁺¹ + ūⁿ⁺¹ * τ̃) / (1 + τ̃ - U)
+
+    @inbounds setindex!(u, u₁ⁿ⁺¹, boundary_indices...)
+    @inbounds setindex!(u, u₁ⁿ⁺¹, boundary_secret_storage_indices...)
+
+    return nothing
+end
+
+
+@inline function step_right_boundary!(bc::FPAOBC, l, m, boundary_indices, boundary_adjacent_indices, 
+                                      grid, u, clock, model_fields, ΔX)
+    Δt = clock.last_stage_Δt
+
+    Δt = ifelse(isinf(Δt), 0, Δt)
+
+    ūⁿ⁺¹ = getbc(bc, l, m, grid, clock, model_fields)
+
+    uᵢⁿ     = @inbounds getindex(u, boundary_indices...)
+    uᵢ₋₁ⁿ⁺¹ = @inbounds getindex(u, boundary_adjacent_indices...)
+
+    U = max(0, min(1, Δt / ΔX * ūⁿ⁺¹))
+
+    pa = bc.classification.matching_scheme
+
+    τ = ifelse(ūⁿ⁺¹ >= 0, pa.outflow_timescale, pa.inflow_timescale)
+
+    τ̃ = Δt / τ
+
+    uᵢⁿ⁺¹ = uᵢⁿ + U * (uᵢ₋₁ⁿ⁺¹ - ūⁿ⁺¹) + (ūⁿ⁺¹ - uᵢⁿ) * τ̃
+
+    @inbounds setindex!(u, uᵢⁿ⁺¹, boundary_indices...)
+
+    return nothing
+end
+
+@inline function step_left_boundary!(bc::FPAOBC, l, m, boundary_indices, boundary_adjacent_indices, boundary_secret_storage_indices, 
+                                     grid, u, clock, model_fields, ΔX)
+    Δt = clock.last_stage_Δt
+
+    Δt = ifelse(isinf(Δt), 0, Δt)
+
+    ūⁿ⁺¹ = getbc(bc, l, m, grid, clock, model_fields)
+
+    uᵢⁿ     = @inbounds getindex(u, boundary_secret_storage_indices...)
+    uᵢ₋₁ⁿ⁺¹ = @inbounds getindex(u, boundary_adjacent_indices...)
+
+    U = min(0, max(-1, Δt / ΔX * ūⁿ⁺¹))
+
+    pa = bc.classification.matching_scheme
+
+    τ = ifelse(ūⁿ⁺¹ <= 0, pa.outflow_timescale, pa.inflow_timescale)
+
+    τ̃ = Δt / τ
+
+    u₁ⁿ⁺¹ = uᵢⁿ - U * (uᵢ₋₁ⁿ⁺¹ - ūⁿ⁺¹) + (ūⁿ⁺¹ - uᵢⁿ) * τ̃
+
+    @inbounds setindex!(u, u₁ⁿ⁺¹, boundary_indices...)
+    @inbounds setindex!(u, u₁ⁿ⁺¹, boundary_secret_storage_indices...)
+
+    return nothing
+end
+
+@inline function _fill_east_halo!(j, k, grid, u, bc::PAOBC, ::Tuple{Face, Any, Any}, clock, model_fields)
+    i = grid.Nx + 1
+
+    boundary_indices = (i, j, k)
+    boundary_adjacent_indices = (i-1, j, k)
+
+    Δx = Δxᶠᶜᶜ(i, j, k, grid)
+
+    step_right_boundary!(bc, j, k, boundary_indices, boundary_adjacent_indices, grid, u, clock, model_fields, Δx)
+
+    return nothing
+end
+
+@inline function _fill_west_halo!(j, k, grid, u, bc::PAOBC, ::Tuple{Face, Any, Any}, clock, model_fields)
+    boundary_indices = (1, j, k)
+    boundary_adjacent_indices = (2, j, k)
+    boundary_secret_storage_indices = (0, j, k)
+
+    Δx = Δxᶠᶜᶜ(1, j, k, grid)
+
+    step_left_boundary!(bc, j, k, boundary_indices, boundary_adjacent_indices, boundary_secret_storage_indices, grid, u, clock, model_fields, Δx)
+
+    return nothing
+end
+
+@inline function _fill_north_halo!(i, k, grid, u, bc::PAOBC, ::Tuple{Any, Face, Any}, clock, model_fields)
+    j = grid.Ny + 1
+
+    boundary_indices = (i, j, k)
+    boundary_adjacent_indices = (i, j-1, k)
+
+    Δy = Δyᶜᶠᶜ(i, j, k, grid)
+
+    step_right_boundary!(bc, i, k, boundary_indices, boundary_adjacent_indices, grid, u, clock, model_fields, Δy)
+
+    return nothing
+end
+
+@inline function _fill_south_halo!(i, k, grid, u, bc::PAOBC, ::Tuple{Any, Face, Any}, clock, model_fields)
+    boundary_indices = (i, 1, k)
+    boundary_adjacent_indices = (i, 1, k)
+    boundary_secret_storage_indices = (i, 0, k)
+
+    Δy = Δyᶜᶠᶜ(i, 1, k, grid)
+
+    step_left_boundary!(bc, i, k, boundary_indices, boundary_adjacent_indices, boundary_secret_storage_indices, grid, u, clock, model_fields, Δy)
+
+    return nothing
+end
+
+@inline function _fill_top_halo!(i, j, grid, u, bc::PAOBC, ::Tuple{Any, Any, Face}, clock, model_fields)
+    k = grid.Nz + 1
+
+    boundary_indices = (i, j, k)
+    boundary_adjacent_indices = (i, j, k-1)
+
+    Δz = Δzᶜᶜᶠ(i, j, k, grid)
+
+    step_right_boundary!(bc, i, j, boundary_indices, boundary_adjacent_indices, grid, u, clock, model_fields, Δz)
+
+    return nothing
+end
+
+@inline function _fill_bottom_halo!(i, j, grid, u, bc::PAOBC, ::Tuple{Any, Any, Face}, clock, model_fields)
+    boundary_indices = (i, j, 1)
+    boundary_adjacent_indices = (i, j, 1)
+    boundary_secret_storage_indices = (i, j, 0)
+
+    Δz = Δzᶜᶜᶠ(i, j, 1, grid)
+
+    step_left_boundary!(bc, i, j, boundary_indices, boundary_adjacent_indices, boundary_secret_storage_indices, grid, u, clock, model_fields, Δz)
+
+    return nothing
+end

src/Oceananigans.jl

@@ -221,6 +221,8 @@ include("Simulations/Simulations.jl")
 # Abstractions for distributed and multi-region models
 include("MultiRegion/MultiRegion.jl")
 
+include("boundary_mean.jl")
+
 #####
 ##### Needed so we can export names from sub-modules at the top-level
 #####

src/boundary_mean.jl

@@ -0,0 +1,60 @@
+using Adapt
+using Oceananigans: instantiated_location
+using Oceananigans.Fields: Center, Face
+using Oceananigans.AbstractOperations: GridMetricOperation, Ax, Ay, Az
+using Oceananigans.BoundaryConditions: BoundaryCondition, Open, PerturbationAdvection
+
+import Adapt: adapt_structure
+import Oceananigans.BoundaryConditions: update_boundary_condition!
+
+struct BoundaryAdjacentMean{V}
+   value :: V
+
+   BoundaryAdjacentMean(; value::BV = Ref(0.0)) where BV = new{BV}(value)
+end
+
+@inline (bam::BoundaryAdjacentMean)(args...) = bam.value[]
+
+Adapt.adapt_structure(to, mo::BoundaryAdjacentMean) = 
+    BoundaryAdjacentMean(; value = adapt(to, mo.value[]))
+
+const MOPABC = BoundaryCondition{<:Open{<:PerturbationAdvection}, <:BoundaryAdjacentMean}
+
+@inline boundary_normal_area(::Union{Val{:west}, Val{:east}}, grid)   = GridMetricOperation((Face, Center, Center), Ax, grid)
+@inline boundary_normal_area(::Union{Val{:south}, Val{:north}}, grid) = GridMetricOperation((Center, Face, Center), Ay, grid)
+@inline boundary_normal_area(::Union{Val{:bottom}, Val{:top}}, grid)  = GridMetricOperation((Center, Center, Face), Az, grid)
+
+@inline boundary_adjacent_indices(::Val{:east}, grid, loc) = (size(grid, 1), 1, 1), (2, 3)
+@inline boundary_adjacent_indices(val_side::Val{:west}, grid, loc) = (first_interior_index(val_side, loc), 1, 1), (2, 3)
+
+@inline boundary_adjacent_indices(::Val{:north}, grid, loc) = (1, size(grid, 2), 1), (1, 3)
+@inline boundary_adjacent_indices(val_side::Val{:south}, grid, loc) = (1, first_interior_index(val_side, loc), 1), (1, 3)
+
+@inline boundary_adjacent_indices(::Val{:top}, grid, loc) = (1, 1, size(grid, 3)), (2, 3)
+@inline boundary_adjacent_indices(val_side::Val{:bottom}, grid, loc) = (1, 1, first_interior_index(val_side, loc)), (2, 3)
+
+@inline first_interior_index(::Union{Val{:west}, Val{:east}}, ::Tuple{Center, <:Any, <:Any}) = 1
+@inline first_interior_index(::Union{Val{:west}, Val{:east}}, ::Tuple{Face, <:Any, <:Any}) = 2
+
+@inline first_interior_index(::Union{Val{:south}, Val{:north}}, ::Tuple{<:Any, Center, <:Any}) = 1
+@inline first_interior_index(::Union{Val{:south}, Val{:north}}, ::Tuple{<:Any, Face, <:Any}) = 2
+
+@inline first_interior_index(::Union{Val{:bottom}, Val{:top}}, ::Tuple{<:Any, <:Any, Center}) = 1
+@inline first_interior_index(::Union{Val{:bottom}, Val{:top}}, ::Tuple{<:Any, <:Any, Face}) = 2
+
+function update_boundary_condition!(bc::MOPABC, val_side, u, model)
+    grid = model.grid
+    loc = instantiated_location(u)
+
+    An = boundary_normal_area(val_side, grid)
+
+    (i, j, k), dims = boundary_adjacent_indices(val_side, grid, loc)
+
+    total_area = CUDA.@allowscalar sum(An; dims)[i, j, k]
+
+    Ū = sum(u * An; dims) / total_area
+
+    bc.condition.value[] = CUDA.@allowscalar Ū[i, j, k]
+
+    return nothing
+end

test/test_boundary_conditions_integration.jl

@@ -1,6 +1,10 @@
 include("dependencies_for_runtests.jl")
 
-using Oceananigans.BoundaryConditions: ContinuousBoundaryFunction, FlatExtrapolationOpenBoundaryCondition, fill_halo_regions!
+using Oceananigans.BoundaryConditions: ContinuousBoundaryFunction, 
+                                       FlatExtrapolationOpenBoundaryCondition, 
+                                       PerturbationAdvectionOpenBoundaryCondition
+                                       fill_halo_regions!
+
 using Oceananigans: prognostic_fields
 
 function test_boundary_condition(arch, FT, topo, side, field_name, boundary_condition)
@@ -147,7 +151,7 @@ function test_flat_extrapolation_open_boundary_conditions(arch, FT)
         fill_halo_regions!(u1, clock, (); fill_boundary_normal_velocities = false)
         fill_halo_regions!(u2, clock, (); fill_boundary_normal_velocities = false)
 
-        # we can stop the wall normal halos being filled after the pressure solve
+        # we can stop the wall normal halos being filled after the pressure solve - this serves more as a test of the general OBC stuff
         @test interior(u1, 1, 1, 1) .== 2
         @test interior(u2, end_position(Val(orientation), grid)...) .== 2
 
@@ -160,6 +164,67 @@ function test_flat_extrapolation_open_boundary_conditions(arch, FT)
     end
 end
 
+wall_normal_boundary_condition(::Val{1}, obc) = (; u = FieldBoundaryConditions(east = obc, west = obc))
+wall_normal_boundary_condition(::Val{2}, obc) = (; v = FieldBoundaryConditions(south = obc, north = obc))
+wall_normal_boundary_condition(::Val{3}, obc) = (; w = FieldBoundaryConditions(bottom = obc, top = obc))
+
+normal_velocity(::Val{1}, model) = model.velocities.u
+normal_velocity(::Val{2}, model) = model.velocities.v
+normal_velocity(::Val{3}, model) = model.velocities.w
+
+velocity_forcing(::Val{1}, forcing) = (; u = forcing)
+velocity_forcing(::Val{2}, forcing) = (; v = forcing)
+velocity_forcing(::Val{3}, forcing) = (; w = forcing)
+
+function test_pertubation_advection_open_boundary_conditions(arch, FT)
+    for orientation in 1:3
+        topology = tuple(map(n -> ifelse(n == orientation, Bounded, Flat), 1:3)...)
+
+        grid = RectilinearGrid(arch, FT; topology, size = (4, ), x = (0, 4), y = (0, 4), z = (0, 4), halo = (1, ))
+
+        obc = PerturbationAdvectionOpenBoundaryCondition(1, inflow_timescale = 0.1)
+
+        boundary_conditions = wall_normal_boundary_condition(Val(orientation), obc)
+
+        model = NonhydrostaticModel(; grid, boundary_conditions, timestepper = :QuasiAdamsBashforth2)
+
+        u = normal_velocity(Val(orientation), model)
+
+        CUDA.@allowscalar u.data .= 1
+
+        time_step!(model, 1)
+
+        @test all(u .== 1)
+
+        end_index = tuple(map(n -> ifelse(n == orientation, 5, 1), 1:3)...)
+
+        CUDA.@allowscalar u[end_index...] = 2
+
+        time_step!(model, 1)
+
+        CUDA.@allowscalar @test all(interior(u, end_index...) .== 1.5)
+
+        obc = PerturbationAdvectionOpenBoundaryCondition((t) -> 0.1*t, inflow_timescale = 0.01, outflow_timescale = 0.5)
+
+        forcing = velocity_forcing(Val(orientation), Forcing((x, t) -> 0.1))
+
+        boundary_conditions = wall_normal_boundary_condition(Val(orientation), obc)
+
+        model = NonhydrostaticModel(; grid, 
+                              boundary_conditions, 
+                              timestepper = :QuasiAdamsBashforth2,
+                              forcing)
+
+        u = normal_velocity(Val(orientation), model)
+
+        for _ in 1:100
+            time_step!(model, 0.1)
+        end
+
+        @test all(isapprox.(interior(u), 1, atol = 0.1))
+    end
+end
+
 test_boundary_conditions(C, FT, ArrayType) = (integer_bc(C, FT, ArrayType),
                                               float_bc(C, FT, ArrayType),
                                               irrational_bc(C, FT, ArrayType),
@@ -336,6 +401,7 @@ test_boundary_conditions(C, FT, ArrayType) = (integer_bc(C, FT, ArrayType),
             A = typeof(arch)
             @info "  Testing open boundary conditions [$A, $FT]..."
             test_flat_extrapolation_open_boundary_conditions(arch, FT)
+            test_pertubation_advection_open_boundary_conditions(arch, FT)
         end
     end
 end