Different advection schemes in different directions

src/Advection/Advection.jl

@@ -21,7 +21,7 @@ export
     UpwindBiased, UpwindBiasedFirstOrder, UpwindBiasedThirdOrder, UpwindBiasedFifthOrder,
     WENO, WENOThirdOrder, WENOFifthOrder,
     VectorInvariant, WENOVectorInvariant,
-    TracerAdvection,
+    FluxFormAdvection,
     EnergyConserving,
     EnstrophyConserving
 
@@ -38,7 +38,7 @@ using Oceananigans.Architectures: architecture, CPU
 using Oceananigans.Operators
 
 import Base: show, summary
-import Oceananigans.Grids: required_halo_size
+import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z
 import Oceananigans.Architectures: on_architecture
 
 abstract type AbstractAdvectionScheme{B, FT} end
@@ -55,9 +55,12 @@ abstract type AbstractUpwindBiasedAdvectionScheme{B, FT} <: AbstractAdvectionSch
 # Note that it is not possible to compile schemes for `advection_buffer = 41` or higher.
 const advection_buffers = [1, 2, 3, 4, 5, 6]
 
-@inline required_halo_size(::AbstractAdvectionScheme{B}) where B = B
 @inline Base.eltype(::AbstractAdvectionScheme{<:Any, FT}) where FT = FT
 
+@inline required_halo_size_x(::AbstractAdvectionScheme{B}) where B = B
+@inline required_halo_size_y(::AbstractAdvectionScheme{B}) where B = B
+@inline required_halo_size_z(::AbstractAdvectionScheme{B}) where B = B
+
 include("centered_advective_fluxes.jl")
 include("upwind_biased_advective_fluxes.jl")
 
@@ -72,12 +75,14 @@ include("vector_invariant_upwinding.jl")
 include("vector_invariant_advection.jl")
 include("vector_invariant_self_upwinding.jl")
 include("vector_invariant_cross_upwinding.jl")
+include("flux_form_advection.jl")
 
 include("flat_advective_fluxes.jl")
 include("topologically_conditional_interpolation.jl")
 include("momentum_advection_operators.jl")
 include("tracer_advection_operators.jl")
 include("positivity_preserving_tracer_advection_operators.jl")
 include("cell_advection_timescale.jl")
+include("adapt_advection_order.jl")
 
 end # module

src/Advection/adapt_advection_order.jl

@@ -0,0 +1,120 @@
+
+
+function adapt_advection_order_x(advection::Centered{H}, grid::AbstractGrid, N::Int) where H
+    if N == 1
+        return nothing
+    elseif N >= H
+        return advection
+    else
+        return Centered(; order = N)
+    end
+end
+
+function adapt_advection_order_x(advection::UpwindBiased{H}, grid::AbstractGrid, N::Int) where H
+    if N == 1
+        return nothing
+    elseif N >= H
+        return advection
+    else
+        return UpwindBiased(; order = N)
+    end
+end
+
+adapt_advection_order_y(advection, grid::AbstractGrid, N::Int) = adapt_advection_order_x(advection, grid, N)
+adapt_advection_order_z(advection, grid::AbstractGrid, N::Int) = adapt_advection_order_x(advection, grid, N)
+
+"""
+    new_weno_scheme(grid, order, bounds, T)
+
+Constructs a new WENO scheme based on the given parameters. `T` is the type of the weno coefficients. 
+A _non-stretched_ WENO scheme has `T` equal to `Nothing`. In case of a non-stretched WENO scheme, 
+we rebuild the advection without passing the grid information, otherwise we use the grid to account for stretched directions.
+"""
+new_weno_scheme(grid, order, bounds, T) = ifelse(T == Nothing, WENO(; order, bounds), WENO(grid; order, bounds))
+
+function adapt_advection_order_x(advection::WENO{H, FT, XT, YT, ZT}, grid::AbstractGrid) where {H, FT, XT, YT, ZT}
+
+    if N == 1
+        return nothing
+    elseif N >= H
+        return advection
+    else
+        return new_weno_scheme(grid, N, advection.bounds, XT)
+    end
+end
+
+function adapt_advection_order_y(advection::WENO{H, FT, XT, YT, ZT}, grid::AbstractGrid, N::Int) where {H, FT, XT, YT, ZT}
+
+    if N > H
+        return advection
+    else
+        return new_weno_scheme(grid, N, advection.bounds, YT)
+    end
+end
+
+function adapt_advection_order_z(advection::WENO{H, FT, XT, YT, ZT}, grid::AbstractGrid, N::Int) where {H, FT, XT, YT, ZT}
+
+    if N == 1
+        return nothing
+    elseif N >= H
+        return advection
+    else
+        return new_weno_scheme(grid, N, advection.bounds, XT)
+    end
+end
+
+"""
+    adapt_advection_order(advection, grid::AbstractGrid)
+
+Adapts the advection operator `advection` based on the grid `grid` by adjusting the order of advection in each direction.
+For example, if the grid has only one point in the x-direction, the advection operator in the x-direction is set to `nothing`.
+A high order advection sheme is reduced to a lower order advection scheme if the grid has fewer points in that direction.
+
+# Arguments
+- `advection`: The original advection scheme.
+- `grid::AbstractGrid`: The grid on which the advection scheme is applied.
+
+The adapted advection scheme with adjusted advection order returned by this function is a `FluxFormAdvection`.
+"""
+function adapt_advection_order(advection, grid::AbstractGrid)
+    advection_x = adapt_advection_order_x(advection, grid, grid.Nx)
+    advection_y = adapt_advection_order_y(advection, grid, grid.Ny)
+    advection_z = adapt_advection_order_z(advection, grid, grid.Nz)
+
+    # Check that we indeed changed the advection operator
+    changed_x = advection_x != advection
+    changed_y = advection_y != advection
+    changed_z = advection_z != advection
+
+    new_advection = FluxFormAdvection(advection_x, advection_y, advection_z)
+    changed_advection = any((changed_x, changed_y, changed_z))
+
+    if changed_advection
+        @info "User-defined advection scheme $(advection) reduced to $(new_advection) to comply with grid-size limitations."
+    end
+
+    return ifelse(changed_advection, new_advection, advection)
+end
+
+function adapt_advection_order(advection::FluxFormAdvection, grid::AbstractGrid)
+    advection_x = adapt_advection_order_x(advection.x, grid, grid.Nx)
+    advection_y = adapt_advection_order_y(advection.y, grid, grid.Ny)
+    advection_z = adapt_advection_order_z(advection.z, grid, grid.Nz)
+
+    # Check that we indeed changed the advection operator
+    changed_x = advection_x != advection.x
+    changed_y = advection_y != advection.y
+    changed_z = advection_z != advection.z
+
+    new_advection = FluxFormAdvection(advection_x, advection_y, advection_z)
+    changed_advection = any((changed_x, changed_y, changed_z))
+
+    if changed_advection
+        @info "User-defined advection scheme $(advection) reduced to $(new_advection) to comply with grid-size limitations."
+    end
+
+    return ifelse(changed_advection, new_advection, advection)
+end
+
+# For the moment, we do not adapt the advection order for the VectorInvariant advection scheme
+adapt_advection_order(advection::VectorInvariant, grid::AbstractGrid) = advection

src/Advection/flux_form_advection.jl

@@ -0,0 +1,58 @@
+using Oceananigans.Operators: Vᶜᶜᶜ
+using Oceananigans.Fields: ZeroField
+
+struct FluxFormAdvection{N, FT, A, B, C} <: AbstractAdvectionScheme{N, FT}
+    x :: A
+    y :: B
+    z :: C
+
+    FluxFormAdvection{N, FT}(x::A, y::B, z::C) where {N, FT, A, B, C} = new{N, FT, A, B, C}(x, y, z)
+end
+
+"""
+    function FluxFormAdvection(x, y, z)
+
+Builds a `FluxFormAdvection` type with reconstructions schemes `x`, `y`, and `z` to be applied in
+the x, y, and z direction, respectively.
+"""
+function FluxFormAdvection(x_advection, y_advection, z_advection)
+    Hx = required_halo_size_x(x_advection)
+    Hy = required_halo_size_y(y_advection)
+    Hz = required_halo_size_z(z_advection)
+
+    FT = eltype(x_advection)
+    H = max(Hx, Hy, Hz)
+
+    return FluxFormAdvection{H, FT}(x_advection, y_advection, z_advection)
+end
+
+Base.show(io::IO, scheme::FluxFormAdvection) = 
+    print(io, "FluxFormAdvection with reconstructions: ", " \n",
+          "    ├── x: ", summary(scheme.x), "\n",
+          "    ├── y: ", summary(scheme.y), "\n",
+          "    └── z: ", summary(scheme.z))
+
+@inline required_halo_size_x(scheme::FluxFormAdvection) = required_halo_size_x(scheme.x)
+@inline required_halo_size_y(scheme::FluxFormAdvection) = required_halo_size_y(scheme.y)
+@inline required_halo_size_z(scheme::FluxFormAdvection) = required_halo_size_z(scheme.z)
+
+Adapt.adapt_structure(to, scheme::FluxFormAdvection{N, FT}) where {N, FT} = 
+    FluxFormAdvection{N, FT}(Adapt.adapt(to, scheme.x),
+                             Adapt.adapt(to, scheme.y),
+                             Adapt.adapt(to, scheme.z))
+
+@inline _advective_tracer_flux_x(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_tracer_flux_x(i, j, k, grid, advection.x, args...)
+@inline _advective_tracer_flux_y(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_tracer_flux_y(i, j, k, grid, advection.y, args...)
+@inline _advective_tracer_flux_z(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_tracer_flux_z(i, j, k, grid, advection.z, args...)
+
+@inline _advective_momentum_flux_Uu(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_momentum_flux_Uu(i, j, k, grid, advection.x, args...)
+@inline _advective_momentum_flux_Vu(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_momentum_flux_Vu(i, j, k, grid, advection.y, args...)
+@inline _advective_momentum_flux_Wu(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_momentum_flux_Wu(i, j, k, grid, advection.z, args...)
+
+@inline _advective_momentum_flux_Uv(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_momentum_flux_Uv(i, j, k, grid, advection.x, args...)
+@inline _advective_momentum_flux_Vv(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_momentum_flux_Vv(i, j, k, grid, advection.y, args...)
+@inline _advective_momentum_flux_Wv(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_momentum_flux_Wv(i, j, k, grid, advection.z, args...)
+
+@inline _advective_momentum_flux_Uw(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_momentum_flux_Uw(i, j, k, grid, advection.x, args...)
+@inline _advective_momentum_flux_Vw(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_momentum_flux_Vw(i, j, k, grid, advection.y, args...)
+@inline _advective_momentum_flux_Ww(i, j, k, grid, advection::FluxFormAdvection, args...) = _advective_momentum_flux_Ww(i, j, k, grid, advection.z, args...)

src/Advection/momentum_advection_operators.jl

@@ -89,3 +89,20 @@ which ends up at the location `ccf`.
                                     δyᵃᶜᵃ(i, j, k, grid, _advective_momentum_flux_Vw, advection, U[2], w) +
                                     δzᵃᵃᶠ(i, j, k, grid, _advective_momentum_flux_Ww, advection, U[3], w))
 end
+
+#####
+##### Fallback advection fluxes!
+#####
+
+for flux_dir in (:Uu, :Vu, :Wu, :Uv, :Vv, :Wv, :Uw, :Vw, :Ww)
+    advective_momentum_flux = Symbol(:_advective_momentum_flux_, flux_dir)
+
+    @eval begin
+        @inline $advective_momentum_flux(i, j, k, grid, ::Nothing, args...)                  = zero(grid)
+        @inline $advective_momentum_flux(i, j, k, grid, ::Nothing, ::ZeroField, ::ZeroField) = zero(grid)
+        @inline $advective_momentum_flux(i, j, k, grid, ::Nothing, U, ::ZeroField)           = zero(grid)
+        @inline $advective_momentum_flux(i, j, k, grid, scheme, ::ZeroField, ::ZeroField)    = zero(grid)
+        @inline $advective_momentum_flux(i, j, k, grid, scheme, U, ::ZeroField)              = zero(grid)
+        @inline $advective_momentum_flux(i, j, k, grid, scheme, ::ZeroField, u)              = zero(grid)
+    end
+end

src/Advection/topologically_conditional_interpolation.jl

@@ -26,14 +26,23 @@ const AUGXYZ = AUG{<:Any, <:Bounded, <:Bounded, <:Bounded}
 # Left-biased buffers are smaller by one grid point on the right side; vice versa for right-biased buffers
 # Center interpolation stencil look at i + 1 (i.e., require one less point on the left)
 
-@inline outside_symmetric_haloᶠ(i, N, adv) = (i >= required_halo_size(adv) + 1) & (i <= N + 1 - required_halo_size(adv))
-@inline outside_symmetric_haloᶜ(i, N, adv) = (i >= required_halo_size(adv))     & (i <= N + 1 - required_halo_size(adv))
-
-@inline outside_biased_haloᶠ(i, N, adv) = (i >= required_halo_size(adv) + 1) & (i <= N + 1 - (required_halo_size(adv) - 1)) &  # Left bias
-                                          (i >= required_halo_size(adv))     & (i <= N + 1 - required_halo_size(adv))          # Right bias
-@inline outside_biased_haloᶜ(i, N, adv) = (i >= required_halo_size(adv))     & (i <= N + 1 - (required_halo_size(adv) - 1)) &  # Left bias
-                                          (i >= required_halo_size(adv) - 1) & (i <= N + 1 - required_halo_size(adv))          # Right bias
-
+for dir in (:x, :y, :z)
+    outside_symmetric_haloᶠ = Symbol(:outside_symmetric_halo_, dir, :ᶠ)
+    outside_symmetric_haloᶜ = Symbol(:outside_symmetric_halo_, dir, :ᶜ)
+    outside_biased_haloᶠ = Symbol(:outside_biased_halo_, dir, :ᶠ)
+    outside_biased_haloᶜ = Symbol(:outside_biased_halo_, dir, :ᶜ)
+    required_halo_size = Symbol(:required_halo_size_, dir)
+
+    @eval begin
+     @inline $outside_symmetric_haloᶠ(i, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - $required_halo_size(adv))
+     @inline $outside_symmetric_haloᶜ(i, N, adv) = (i >= $required_halo_size(adv))     & (i <= N + 1 - $required_halo_size(adv))
+
+     @inline $outside_biased_haloᶠ(i, N, adv) = (i >= $required_halo_size(adv) + 1) & (i <= N + 1 - ($required_halo_size(adv) - 1)) &  # Left bias
+                                                (i >= $required_halo_size(adv))     & (i <= N + 1 - $required_halo_size(adv))          # Right bias
+     @inline $outside_biased_haloᶜ(i, N, adv) = (i >= $required_halo_size(adv))     & (i <= N + 1 - ($required_halo_size(adv) - 1)) &  # Left bias
+                                                (i >= $required_halo_size(adv) - 1) & (i <= N + 1 - $required_halo_size(adv))          # Right bias
+    end
+end
 # Separate High order advection from low order advection
 const HOADV = Union{WENO, 
                     Tuple(Centered{N} for N in advection_buffers[2:end])...,
@@ -61,7 +70,7 @@ for bias in (:symmetric, :biased)
                 @eval @inline $alt1_interp(i, j, k, grid::$GridType, scheme::LOADV, args...) = $interp(i, j, k, grid, scheme, args...)
             end
 
-            outside_buffer = Symbol(:outside_, bias, :_halo, loc)
+            outside_buffer = Symbol(:outside_, bias, :_halo_, ξ, loc)
 
             # Conditional high-order interpolation in Bounded directions
             if ξ == :x

src/Advection/tracer_advection_operators.jl

@@ -1,80 +1,24 @@
-using Oceananigans.Operators: Vᶜᶜᶜ
-using Oceananigans.Fields: ZeroField
-
-struct TracerAdvection{N, FT, A, B, C} <: AbstractAdvectionScheme{N, FT}
-    x :: A
-    y :: B
-    z :: C
-
-    TracerAdvection{N, FT}(x::A, y::B, z::C) where {N, FT, A, B, C} = new{N, FT, A, B, C}(x, y, z)
-end
-
-"""
-    function TracerAdvection(x, y, z)
-
-Builds a `TracerAdvection` type with reconstructions schemes `x`, `y`, and `z` to be applied in
-the x, y, and z direction, respectively.
-"""
-function TracerAdvection(x_advection, y_advection, z_advection)
-    Hx = required_halo_size(x_advection)
-    Hy = required_halo_size(y_advection)
-    Hz = required_halo_size(z_advection)
-
-    FT = eltype(x_advection)
-    H = max(Hx, Hy, Hz)
-
-    return TracerAdvection{H, FT}(x_advection, y_advection, z_advection)
-end
-
-Adapt.adapt_structure(to, scheme::TracerAdvection{N, FT}) where {N, FT} = 
-    TracerAdvection{N, FT}(Adapt.adapt(to, scheme.x),
-                           Adapt.adapt(to, scheme.y),
-                           Adapt.adapt(to, scheme.z))
 
 @inline _advective_tracer_flux_x(args...) = advective_tracer_flux_x(args...)
 @inline _advective_tracer_flux_y(args...) = advective_tracer_flux_y(args...)
 @inline _advective_tracer_flux_z(args...) = advective_tracer_flux_z(args...)
 
-@inline _advective_tracer_flux_x(i, j, k, grid, advection::TracerAdvection, args...) =
-        _advective_tracer_flux_x(i, j, k, grid, advection.x, args...)
-
-@inline _advective_tracer_flux_y(i, j, k, grid, advection::TracerAdvection, args...) =
-        _advective_tracer_flux_y(i, j, k, grid, advection.y, args...)
-
-@inline _advective_tracer_flux_z(i, j, k, grid, advection::TracerAdvection, args...) =
-        _advective_tracer_flux_z(i, j, k, grid, advection.z, args...)
-
-# Fallback for `nothing` advection
-@inline _advective_tracer_flux_x(i, j, k, grid, ::Nothing, args...) = zero(grid)
-@inline _advective_tracer_flux_y(i, j, k, grid, ::Nothing, args...) = zero(grid)
-@inline _advective_tracer_flux_z(i, j, k, grid, ::Nothing, args...) = zero(grid)
-
-# Fallback for `nothing` advection and `ZeroField` tracers and velocities
-@inline _advective_tracer_flux_x(i, j, k, grid, ::Nothing, ::ZeroField, ::ZeroField) = zero(grid)
-@inline _advective_tracer_flux_y(i, j, k, grid, ::Nothing, ::ZeroField, ::ZeroField) = zero(grid)
-@inline _advective_tracer_flux_z(i, j, k, grid, ::Nothing, ::ZeroField, ::ZeroField) = zero(grid)
-
-@inline _advective_tracer_flux_x(i, j, k, grid, ::Nothing, U, ::ZeroField) = zero(grid)
-@inline _advective_tracer_flux_y(i, j, k, grid, ::Nothing, V, ::ZeroField) = zero(grid)
-@inline _advective_tracer_flux_z(i, j, k, grid, ::Nothing, W, ::ZeroField) = zero(grid)
-@inline _advective_tracer_flux_x(i, j, k, grid, ::Nothing, ::ZeroField, c) = zero(grid)
-@inline _advective_tracer_flux_y(i, j, k, grid, ::Nothing, ::ZeroField, c) = zero(grid)
-@inline _advective_tracer_flux_z(i, j, k, grid, ::Nothing, ::ZeroField, c) = zero(grid)
-
-# Fallback for `ZeroField` tracers and velocities
-@inline _advective_tracer_flux_x(i, j, k, grid, scheme, ::ZeroField, ::ZeroField) = zero(grid)
-@inline _advective_tracer_flux_y(i, j, k, grid, scheme, ::ZeroField, ::ZeroField) = zero(grid)
-@inline _advective_tracer_flux_z(i, j, k, grid, scheme, ::ZeroField, ::ZeroField) = zero(grid)
-
-# Fallback for `ZeroField` tracers
-@inline _advective_tracer_flux_x(i, j, k, grid, scheme, U, ::ZeroField) = zero(grid)
-@inline _advective_tracer_flux_y(i, j, k, grid, scheme, V, ::ZeroField) = zero(grid)
-@inline _advective_tracer_flux_z(i, j, k, grid, scheme, W, ::ZeroField) = zero(grid)
+#####
+##### Fallback tracer fluxes!
+#####
 
-# Fallback for `ZeroField` velocities
-@inline _advective_tracer_flux_x(i, j, k, grid, scheme, ::ZeroField, c) = zero(grid)
-@inline _advective_tracer_flux_y(i, j, k, grid, scheme, ::ZeroField, c) = zero(grid)
-@inline _advective_tracer_flux_z(i, j, k, grid, scheme, ::ZeroField, c) = zero(grid)
+for flux_dir in (:x, :y, :z)
+    advective_tracer_flux = Symbol(:_advective_tracer_flux_, flux_dir)
+
+    @eval begin
+        @inline $advective_tracer_flux(i, j, k, grid, ::Nothing, args...)                  = zero(grid)
+        @inline $advective_tracer_flux(i, j, k, grid, ::Nothing, ::ZeroField, ::ZeroField) = zero(grid)
+        @inline $advective_tracer_flux(i, j, k, grid, ::Nothing, U, ::ZeroField)           = zero(grid)
+        @inline $advective_tracer_flux(i, j, k, grid, scheme, ::ZeroField, ::ZeroField)    = zero(grid)
+        @inline $advective_tracer_flux(i, j, k, grid, scheme, U, ::ZeroField)              = zero(grid)
+        @inline $advective_tracer_flux(i, j, k, grid, scheme, ::ZeroField, c)              = zero(grid)
+    end
+end
 
 #####
 ##### Tracer advection operator