From 0210e6646d6bab93e2c2d579b7e389e1ccc0c5db Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Fri, 11 Oct 2024 16:30:01 +0200 Subject: [PATCH] Different advection schemes in different directions (#3732) * should be done * add the dicussion * disambiguation * remove repeat code * precopmiles * keep these changes for later * correct the outside buffer * fix tests * change the show method * back to previous buffer scheme * allow for changes in the advection * switch position of functions * better comment * better info * bugfix * better info * make sure N==1 -> nothing * Update src/Advection/momentum_advection_operators.jl Co-authored-by: Gregory L. Wagner * add topologies * add topology argument * better adaptation * formatting * whoops * remove topology and less metaprogramming * add advection * dont adapt nothing * retain same behavior as before * fix multi region advection * bugfix scalar diffusivities parameter types * bugfix * enzyme error * make sure constructors are type-stable... * force required_halo_size to be an integer * add better comment * add a couple of unit tests * test biharmonic diffusivity * fix enzyme tests * correct adaptbetter explanation * add a test for adapt_advection * whoops wrong test * revert changes * add couple more tests * fix tests * without the extra space * fix enzyme test * fix enzyme tests * fix nonhydrostatic tests * rmove one newline * simplify test * Update src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl Co-authored-by: Gregory L. Wagner * remove extra dispatches * relaunch the tests * change comment * switch to B * new syntax for KernelParameters * bugfix * another bugfix * make sure the range is correct * Update src/Advection/adapt_advection_order.jl Co-authored-by: Gregory L. Wagner * Update src/Advection/adapt_advection_order.jl Co-authored-by: Gregory L. Wagner * Update src/Advection/adapt_advection_order.jl Co-authored-by: Gregory L. Wagner * remove comments * this should be type stable for the moment * bugfix * bugfix * better comment * fix the docs * whoops --------- Co-authored-by: Gregory L. Wagner --- src/Advection/Advection.jl | 11 ++- src/Advection/adapt_advection_order.jl | 98 +++++++++++++++++++ src/Advection/flux_form_advection.jl | 58 +++++++++++ src/Advection/momentum_advection_operators.jl | 31 +++--- ...topologically_conditional_interpolation.jl | 27 +++-- src/Advection/tracer_advection_operators.jl | 72 +------------- src/Advection/vector_invariant_advection.jl | 22 ++++- src/Advection/weno_reconstruction.jl | 2 +- src/Grids/automatic_halo_sizing.jl | 68 ++++++++++--- src/ImmersedBoundaries/conditional_fluxes.jl | 86 ++++++++++------ ...static_free_surface_boundary_tendencies.jl | 22 ++--- .../hydrostatic_free_surface_model.jl | 6 +- ...pute_nonhydrostatic_boundary_tendencies.jl | 66 ++++++------- .../nonhydrostatic_model.jl | 7 +- src/MultiRegion/multi_region_models.jl | 7 +- src/Oceananigans.jl | 2 +- src/TurbulenceClosures/TurbulenceClosures.jl | 6 +- src/TurbulenceClosures/closure_tuples.jl | 4 +- .../isopycnal_skew_symmetric_diffusivity.jl | 8 +- .../scalar_biharmonic_diffusivity.jl | 23 +++-- .../scalar_diffusivity.jl | 27 +++-- test/test_immersed_advection.jl | 4 +- test/test_nonhydrostatic_models.jl | 24 +++++ test/test_turbulence_closures.jl | 25 ++++- 24 files changed, 487 insertions(+), 219 deletions(-) create mode 100644 src/Advection/adapt_advection_order.jl create mode 100644 src/Advection/flux_form_advection.jl diff --git a/src/Advection/Advection.jl b/src/Advection/Advection.jl index c54b310475..27b5ea6e27 100644 --- a/src/Advection/Advection.jl +++ b/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,6 +75,7 @@ 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") @@ -79,5 +83,6 @@ 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 diff --git a/src/Advection/adapt_advection_order.jl b/src/Advection/adapt_advection_order.jl new file mode 100644 index 0000000000..41f3f60c9a --- /dev/null +++ b/src/Advection/adapt_advection_order.jl @@ -0,0 +1,98 @@ +using Oceananigans.Grids: topology + +""" + 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 first order +upwind or 2nd order centered scheme, depending on the original user-specified advection scheme. 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. + +If the order of advection is changed in at least one direction, the adapted advection scheme with adjusted advection order returned +by this function is a `FluxFormAdvection`. +""" +function adapt_advection_order(advection, grid::AbstractGrid) + advection_x = x_advection(advection) + advection_y = y_advection(advection) + advection_z = z_advection(advection) + + new_advection_x = adapt_advection_order(advection_x, size(grid, 1), grid) + new_advection_y = adapt_advection_order(advection_y, size(grid, 2), grid) + new_advection_z = adapt_advection_order(advection_z, size(grid, 3), grid) + + # Check that we indeed changed the advection operator + changed_x = new_advection_x != advection_x + changed_y = new_advection_y != advection_y + changed_z = new_advection_z != advection_z + + new_advection = FluxFormAdvection(new_advection_x, new_advection_y, new_advection_z) + changed_advection = any((changed_x, changed_y, changed_z)) + + if changed_x + @info "Using the advection scheme $(summary(new_advection.x)) in the x-direction because size(grid, 1) = $(size(grid, 1))" + end + if changed_y + @info "Using the advection scheme $(summary(new_advection.y)) in the y-direction because size(grid, 2) = $(size(grid, 2))" + end + if changed_z + @info "Using the advection scheme $(summary(new_advection.z)) in the x-direction because size(grid, 3) = $(size(grid, 3))" + end + + return ifelse(changed_advection, new_advection, advection) +end + + +x_advection(flux_form::FluxFormAdvection) = flux_form.x +y_advection(flux_form::FluxFormAdvection) = flux_form.y +z_advection(flux_form::FluxFormAdvection) = flux_form.z + +x_advection(advection) = advection +y_advection(advection) = advection +z_advection(advection) = advection + +# For the moment, we do not adapt the advection order for the VectorInvariant advection scheme +adapt_advection_order(advection::VectorInvariant, grid::AbstractGrid) = advection +adapt_advection_order(advection::Nothing, grid::AbstractGrid) = nothing +adapt_advection_order(advection::Nothing, N::Int, grid::AbstractGrid) = nothing + +##### +##### Directional adapt advection order +##### + +function adapt_advection_order(advection::Centered{B}, N::Int, grid::AbstractGrid) where B + if N >= B + return advection + else + return Centered(; order = 2N) + end +end + +function adapt_advection_order(advection::UpwindBiased{B}, N::Int, grid::AbstractGrid) where B + if N >= B + return advection + else + return UpwindBiased(; order = 2N - 1) + end +end + +""" + new_weno_scheme(grid, order, bounds, XT, YT, ZT) + +Constructs a new WENO scheme based on the given parameters. `XT`, `YT`, and `ZT` is the type of the precomputed weno coefficients in the +x-direction, y-direction and z-direction. A _non-stretched_ WENO scheme has `T` equal to `Nothing` everywhere. 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(::WENO, grid, order, bounds, ::Type{Nothing}, ::Type{Nothing}, ::Type{Nothing},) = WENO(; order, bounds) +new_weno_scheme(::WENO, grid, order, bounds, XT, YT, ZT) = WENO(grid; order, bounds) + +function adapt_advection_order(advection::WENO{B, FT, XT, YT, ZT}, N::Int, grid::AbstractGrid) where {B, FT, XT, YT, ZT} + if N >= B + return advection + else + return new_weno_scheme(advection, grid, 2N - 1, advection.bounds, XT, YT, ZT) + end +end \ No newline at end of file diff --git a/src/Advection/flux_form_advection.jl b/src/Advection/flux_form_advection.jl new file mode 100644 index 0000000000..faa35fb668 --- /dev/null +++ b/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...) diff --git a/src/Advection/momentum_advection_operators.jl b/src/Advection/momentum_advection_operators.jl index a2aa5c29ed..eda8f78423 100644 --- a/src/Advection/momentum_advection_operators.jl +++ b/src/Advection/momentum_advection_operators.jl @@ -28,17 +28,9 @@ const ZeroU = NamedTuple{(:u, :v, :w), Tuple{ZeroField, ZeroField, ZeroField}} @inline div_𝐯v(i, j, k, grid, advection, U, ::ZeroField) = zero(grid) @inline div_𝐯w(i, j, k, grid, advection, U, ::ZeroField) = zero(grid) -@inline div_𝐯u(i, j, k, grid, ::Nothing, U, u) = zero(grid) -@inline div_𝐯v(i, j, k, grid, ::Nothing, U, v) = zero(grid) -@inline div_𝐯w(i, j, k, grid, ::Nothing, U, w) = zero(grid) - -@inline div_𝐯u(i, j, k, grid, ::Nothing, ::ZeroU, u) = zero(grid) -@inline div_𝐯v(i, j, k, grid, ::Nothing, ::ZeroU, v) = zero(grid) -@inline div_𝐯w(i, j, k, grid, ::Nothing, ::ZeroU, w) = zero(grid) - -@inline div_𝐯u(i, j, k, grid, ::Nothing, U, ::ZeroField) = zero(grid) -@inline div_𝐯v(i, j, k, grid, ::Nothing, U, ::ZeroField) = zero(grid) -@inline div_𝐯w(i, j, k, grid, ::Nothing, U, ::ZeroField) = zero(grid) +@inline div_𝐯u(i, j, k, grid, advection, ::ZeroU, ::ZeroField) = zero(grid) +@inline div_𝐯v(i, j, k, grid, advection, ::ZeroU, ::ZeroField) = zero(grid) +@inline div_𝐯w(i, j, k, grid, advection, ::ZeroU, ::ZeroField) = zero(grid) """ div_𝐯u(i, j, k, grid, advection, U, u) @@ -89,3 +81,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! +##### + +# Fallback for `nothing` advection +@inline _advective_momentum_flux_Uu(i, j, k, grid, ::Nothing, args...) = zero(grid) +@inline _advective_momentum_flux_Uv(i, j, k, grid, ::Nothing, args...) = zero(grid) +@inline _advective_momentum_flux_Uw(i, j, k, grid, ::Nothing, args...) = zero(grid) + +@inline _advective_momentum_flux_Vu(i, j, k, grid, ::Nothing, args...) = zero(grid) +@inline _advective_momentum_flux_Vv(i, j, k, grid, ::Nothing, args...) = zero(grid) +@inline _advective_momentum_flux_Vw(i, j, k, grid, ::Nothing, args...) = zero(grid) + +@inline _advective_momentum_flux_Wu(i, j, k, grid, ::Nothing, args...) = zero(grid) +@inline _advective_momentum_flux_Wv(i, j, k, grid, ::Nothing, args...) = zero(grid) +@inline _advective_momentum_flux_Ww(i, j, k, grid, ::Nothing, args...) = zero(grid) \ No newline at end of file diff --git a/src/Advection/topologically_conditional_interpolation.jl b/src/Advection/topologically_conditional_interpolation.jl index 094f5a7253..9c623c6b57 100644 --- a/src/Advection/topologically_conditional_interpolation.jl +++ b/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 diff --git a/src/Advection/tracer_advection_operators.jl b/src/Advection/tracer_advection_operators.jl index fcae1923ce..c22dee4cd9 100644 --- a/src/Advection/tracer_advection_operators.jl +++ b/src/Advection/tracer_advection_operators.jl @@ -1,81 +1,17 @@ -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 tracer fluxes! +##### # 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 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) - ##### ##### Tracer advection operator ##### @@ -100,7 +36,9 @@ end # Fallbacks for zero velocities, zero tracer and `nothing` advection @inline div_Uc(i, j, k, grid, advection, ::ZeroU, c) = zero(grid) @inline div_Uc(i, j, k, grid, advection, U, ::ZeroField) = zero(grid) +@inline div_Uc(i, j, k, grid, advection, ::ZeroU, ::ZeroField) = zero(grid) @inline div_Uc(i, j, k, grid, ::Nothing, U, c) = zero(grid) @inline div_Uc(i, j, k, grid, ::Nothing, ::ZeroU, c) = zero(grid) @inline div_Uc(i, j, k, grid, ::Nothing, U, ::ZeroField) = zero(grid) +@inline div_Uc(i, j, k, grid, ::Nothing, ::ZeroU, ::ZeroField) = zero(grid) diff --git a/src/Advection/vector_invariant_advection.jl b/src/Advection/vector_invariant_advection.jl index a9a89b2623..e16cd7438a 100644 --- a/src/Advection/vector_invariant_advection.jl +++ b/src/Advection/vector_invariant_advection.jl @@ -116,7 +116,14 @@ function VectorInvariant(; vorticity_scheme = EnstrophyConserving(), upwinding = OnlySelfUpwinding(; cross_scheme = divergence_scheme), multi_dimensional_stencil = false) - N = required_halo_size(vorticity_scheme) + N = max(required_halo_size_x(vorticity_scheme), + required_halo_size_y(vorticity_scheme), + required_halo_size_x(divergence_scheme), + required_halo_size_y(divergence_scheme), + required_halo_size_x(kinetic_energy_gradient_scheme), + required_halo_size_y(kinetic_energy_gradient_scheme), + required_halo_size_z(vertical_scheme)) + FT = eltype(vorticity_scheme) return VectorInvariant{N, FT, multi_dimensional_stencil}(vorticity_scheme, @@ -216,11 +223,20 @@ function WENOVectorInvariant(FT::DataType = Float64; upwinding) end - # Since vorticity itself requires one halo, if we use an upwinding scheme (N > 1) we require one additional # halo for vector invariant advection -required_halo_size(scheme::VectorInvariant{N}) where N = N == 1 ? N : N + 1 +@inline function required_halo_size_x(scheme::VectorInvariant) + Hx₁ = required_halo_size_x(scheme.vorticity_scheme) + Hx₂ = required_halo_size_x(scheme.divergence_scheme) + Hx₃ = required_halo_size_x(scheme.kinetic_energy_gradient_scheme) + + Hx = max(Hx₁, Hx₂, Hx₃) + return Hx == 1 ? Hx : Hx + 1 +end +@inline required_halo_size_y(scheme::VectorInvariant) = required_halo_size_x(scheme) +@inline required_halo_size_z(scheme::VectorInvariant) = required_halo_size_z(scheme.vertical_scheme) + Adapt.adapt_structure(to, scheme::VectorInvariant{N, FT, M}) where {N, FT, M} = VectorInvariant{N, FT, M}(Adapt.adapt(to, scheme.vorticity_scheme), Adapt.adapt(to, scheme.vorticity_stencil), diff --git a/src/Advection/weno_reconstruction.jl b/src/Advection/weno_reconstruction.jl index 5ed5a16adc..f91c633557 100644 --- a/src/Advection/weno_reconstruction.jl +++ b/src/Advection/weno_reconstruction.jl @@ -111,8 +111,8 @@ function WENO(FT::DataType=Float64; N = Int((order + 1) ÷ 2) weno_coefficients = compute_reconstruction_coefficients(grid, FT, :WENO; order = N) - buffer_scheme = WENO(FT; grid, order = order - 2, bounds) advecting_velocity_scheme = Centered(FT; grid, order = order - 1) + buffer_scheme = WENO(FT; grid, order = order - 2, bounds) end return WENO{N, FT}(weno_coefficients..., bounds, buffer_scheme, advecting_velocity_scheme) diff --git a/src/Grids/automatic_halo_sizing.jl b/src/Grids/automatic_halo_sizing.jl index a4fd19b526..262728efa1 100644 --- a/src/Grids/automatic_halo_sizing.jl +++ b/src/Grids/automatic_halo_sizing.jl @@ -1,7 +1,7 @@ """ - required_halo_size(tendency_term) + required_halo_size_x(tendency_term) -Return the required size of halos for a term appearing +Return the required size of halos in the x direction for a term appearing in the tendency for a velocity field or tracer field. Example @@ -9,17 +9,61 @@ Example ```jldoctest using Oceananigans.Advection: CenteredFourthOrder -using Oceananigans.Grids: required_halo_size +using Oceananigans.Grids: required_halo_size_x -required_halo_size(CenteredFourthOrder()) +required_halo_size_x(CenteredFourthOrder()) # output 2 """ -function required_halo_size end +function required_halo_size_x end -required_halo_size(tendency_term) = 1 -required_halo_size(::Nothing) = 0 +""" + required_halo_size_y(tendency_term) + +Return the required size of halos in the y direction for a term appearing +in the tendency for a velocity field or tracer field. + +Example +======= + +```jldoctest +using Oceananigans.Advection: CenteredFourthOrder +using Oceananigans.Grids: required_halo_size_y + +required_halo_size_y(CenteredFourthOrder()) + +# output +2 +""" +function required_halo_size_y end + +""" + required_halo_size_z(tendency_term) + +Return the required size of halos in the y direction for a term appearing +in the tendency for a velocity field or tracer field. + +Example +======= + +```jldoctest +using Oceananigans.Advection: CenteredFourthOrder +using Oceananigans.Grids: required_halo_size_z + +required_halo_size_z(CenteredFourthOrder()) + +# output +2 +""" +function required_halo_size_z end + +required_halo_size_x(tendency_term) = 1 +required_halo_size_x(::Nothing) = 0 +required_halo_size_y(tendency_term) = 1 +required_halo_size_y(::Nothing) = 0 +required_halo_size_z(tendency_term) = 1 +required_halo_size_z(::Nothing) = 0 inflate_halo_size_one_dimension(req_H, old_H, _, grid) = max(req_H, old_H) inflate_halo_size_one_dimension(req_H, old_H, ::Type{Flat}, grid) = 0 @@ -27,10 +71,12 @@ inflate_halo_size_one_dimension(req_H, old_H, ::Type{Flat}, grid) = 0 function inflate_halo_size(Hx, Hy, Hz, grid, tendency_terms...) topo = topology(grid) for term in tendency_terms - H = required_halo_size(term) - Hx = inflate_halo_size_one_dimension(H, Hx, topo[1], grid) - Hy = inflate_halo_size_one_dimension(H, Hy, topo[2], grid) - Hz = inflate_halo_size_one_dimension(H, Hz, topo[3], grid) + Hx_required = required_halo_size_x(term) + Hy_required = required_halo_size_y(term) + Hz_required = required_halo_size_z(term) + Hx = inflate_halo_size_one_dimension(Hx_required, Hx, topo[1], grid) + Hy = inflate_halo_size_one_dimension(Hy_required, Hy, topo[2], grid) + Hz = inflate_halo_size_one_dimension(Hz_required, Hz, topo[3], grid) end return Hx, Hy, Hz diff --git a/src/ImmersedBoundaries/conditional_fluxes.jl b/src/ImmersedBoundaries/conditional_fluxes.jl index 260e397c75..fd2d45fb8b 100644 --- a/src/ImmersedBoundaries/conditional_fluxes.jl +++ b/src/ImmersedBoundaries/conditional_fluxes.jl @@ -1,7 +1,7 @@ using Oceananigans.Advection: AbstractAdvectionScheme, advection_buffers using Oceananigans.Operators: ℑxᶠᵃᵃ, ℑxᶜᵃᵃ, ℑyᵃᶠᵃ, ℑyᵃᶜᵃ, ℑzᵃᵃᶠ, ℑzᵃᵃᶜ using Oceananigans.TurbulenceClosures: AbstractTurbulenceClosure, AbstractTimeDiscretization -using Oceananigans.Advection: LOADV, HOADV, WENO, TracerAdvection +using Oceananigans.Advection: LOADV, HOADV, WENO, FluxFormAdvection using Oceananigans.Fields: ZeroField const ATC = AbstractTurbulenceClosure @@ -83,47 +83,69 @@ end @inline _advective_tracer_flux_z(i, j, k, ibg::IBG, args...) = conditional_flux_ccf(i, j, k, ibg, zero(ibg), advective_tracer_flux_z(i, j, k, ibg, args...)) # Disambiguation for tracer fluxes.... -@inline _advective_tracer_flux_x(i, j, k, ibg::IBG, advection::TracerAdvection, args...) = +@inline _advective_tracer_flux_x(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = _advective_tracer_flux_x(i, j, k, ibg, advection.x, args...) -@inline _advective_tracer_flux_y(i, j, k, ibg::IBG, advection::TracerAdvection, args...) = +@inline _advective_tracer_flux_y(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = _advective_tracer_flux_y(i, j, k, ibg, advection.y, args...) -@inline _advective_tracer_flux_z(i, j, k, ibg::IBG, advection::TracerAdvection, args...) = +@inline _advective_tracer_flux_z(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = _advective_tracer_flux_z(i, j, k, ibg, advection.z, args...) -# Fallback for `nothing` advection +# Disambiguation for momentum fluxes in x.... +@inline _advective_momentum_flux_Uu(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = + _advective_momentum_flux_Uu(i, j, k, ibg, advection.x, args...) + +@inline _advective_momentum_flux_Uv(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = + _advective_momentum_flux_Uv(i, j, k, ibg, advection.x, args...) + +@inline _advective_momentum_flux_Uw(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = + _advective_momentum_flux_Uw(i, j, k, ibg, advection.x, args...) + +# Disambiguation for momentum fluxes in y.... +@inline _advective_momentum_flux_Vu(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = + _advective_momentum_flux_Vu(i, j, k, ibg, advection.y, args...) + +@inline _advective_momentum_flux_Vv(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = + _advective_momentum_flux_Vv(i, j, k, ibg, advection.y, args...) + +@inline _advective_momentum_flux_Vw(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = + _advective_momentum_flux_Vw(i, j, k, ibg, advection.y, args...) + +# Disambiguation for momentum fluxes in z.... +@inline _advective_momentum_flux_Wu(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = + _advective_momentum_flux_Wu(i, j, k, ibg, advection.z, args...) + +@inline _advective_momentum_flux_Wv(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = + _advective_momentum_flux_Wv(i, j, k, ibg, advection.z, args...) + +@inline _advective_momentum_flux_Ww(i, j, k, ibg::IBG, advection::FluxFormAdvection, args...) = + _advective_momentum_flux_Ww(i, j, k, ibg, advection.z, args...) + +# Fallback for `nothing` Advection + +# dx(uu), dy(vu), dz(wu) +# ccc, ffc, fcf +@inline _advective_momentum_flux_Uu(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) +@inline _advective_momentum_flux_Vu(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) +@inline _advective_momentum_flux_Wu(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) + +# dx(uv), dy(vv), dz(wv) +# ffc, ccc, cff +@inline _advective_momentum_flux_Uv(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) +@inline _advective_momentum_flux_Vv(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) +@inline _advective_momentum_flux_Wv(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) + +# dx(uw), dy(vw), dz(ww) +# fcf, cff, ccc +@inline _advective_momentum_flux_Uw(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) +@inline _advective_momentum_flux_Vw(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) +@inline _advective_momentum_flux_Ww(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) + @inline _advective_tracer_flux_x(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) @inline _advective_tracer_flux_y(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) @inline _advective_tracer_flux_z(i, j, k, ibg::IBG, ::Nothing, args...) = zero(ibg) -# Fallback for `nothing` advection and `ZeroField` tracers and velocities -@inline _advective_tracer_flux_x(i, j, k, ibg::IBG, ::Nothing, ::ZeroField, ::ZeroField) = zero(ibg) -@inline _advective_tracer_flux_y(i, j, k, ibg::IBG, ::Nothing, ::ZeroField, ::ZeroField) = zero(ibg) -@inline _advective_tracer_flux_z(i, j, k, ibg::IBG, ::Nothing, ::ZeroField, ::ZeroField) = zero(ibg) - -@inline _advective_tracer_flux_x(i, j, k, ibg::IBG, ::Nothing, U, ::ZeroField) = zero(ibg) -@inline _advective_tracer_flux_y(i, j, k, ibg::IBG, ::Nothing, V, ::ZeroField) = zero(ibg) -@inline _advective_tracer_flux_z(i, j, k, ibg::IBG, ::Nothing, W, ::ZeroField) = zero(ibg) -@inline _advective_tracer_flux_x(i, j, k, ibg::IBG, ::Nothing, ::ZeroField, c) = zero(ibg) -@inline _advective_tracer_flux_y(i, j, k, ibg::IBG, ::Nothing, ::ZeroField, c) = zero(ibg) -@inline _advective_tracer_flux_z(i, j, k, ibg::IBG, ::Nothing, ::ZeroField, c) = zero(ibg) - -# Fallback for `ZeroField` tracers and velocities -@inline _advective_tracer_flux_x(i, j, k, ibg::IBG, scheme, ::ZeroField, ::ZeroField) = zero(ibg) -@inline _advective_tracer_flux_y(i, j, k, ibg::IBG, scheme, ::ZeroField, ::ZeroField) = zero(ibg) -@inline _advective_tracer_flux_z(i, j, k, ibg::IBG, scheme, ::ZeroField, ::ZeroField) = zero(ibg) - -# Fallback for `ZeroField` tracers -@inline _advective_tracer_flux_x(i, j, k, ibg::IBG, scheme, U, ::ZeroField) = zero(ibg) -@inline _advective_tracer_flux_y(i, j, k, ibg::IBG, scheme, V, ::ZeroField) = zero(ibg) -@inline _advective_tracer_flux_z(i, j, k, ibg::IBG, scheme, W, ::ZeroField) = zero(ibg) - -# Fallback for `ZeroField` velocities -@inline _advective_tracer_flux_x(i, j, k, ibg::IBG, scheme, ::ZeroField, c) = zero(ibg) -@inline _advective_tracer_flux_y(i, j, k, ibg::IBG, scheme, ::ZeroField, c) = zero(ibg) -@inline _advective_tracer_flux_z(i, j, k, ibg::IBG, scheme, ::ZeroField, c) = zero(ibg) - ##### ##### "Boundary-aware" reconstruct ##### diff --git a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_boundary_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_boundary_tendencies.jl index 375ab00883..bbc13a43d1 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_boundary_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_boundary_tendencies.jl @@ -8,8 +8,6 @@ using Oceananigans.Models.NonhydrostaticModels: boundary_tendency_kernel_paramet boundary_κ_kernel_parameters, boundary_parameters -using Oceananigans.TurbulenceClosures: required_halo_size - # We assume here that top/bottom BC are always synchronized (no partitioning in z) function compute_boundary_tendencies!(model::HydrostaticFreeSurfaceModel) grid = model.grid @@ -19,7 +17,7 @@ function compute_boundary_tendencies!(model::HydrostaticFreeSurfaceModel) p_parameters = boundary_p_kernel_parameters(grid, arch) κ_parameters = boundary_κ_kernel_parameters(grid, model.closure, arch) - # We need new values for `w`, `p` and `κ` + # Compute new values for `w`, `p` and `κ` on the perifery compute_auxiliaries!(model; w_parameters, p_parameters, κ_parameters) # parameters for communicating North / South / East / West side @@ -58,19 +56,15 @@ function boundary_w_kernel_parameters(grid, arch) Nx, Ny, _ = size(grid) Hx, Hy, _ = halo_size(grid) - Sx = (Hx, Ny+2) - Sy = (Nx+2, Hy) - # Offsets in tangential direction are == -1 to # cover the required corners - Oxᴸ = (-Hx+1, -1) - Oyᴸ = (-1, -Hy+1) - Oxᴿ = (Nx-1, -1) - Oyᴿ = (-1, Ny-1) + param_west = (-Hx+2:1, 0:Ny+1) + param_east = (Nx:Nx+Hx-1, 0:Ny+1) + param_south = (0:Nx+1, -Hy+2:1) + param_north = (0:Nx+1, Ny:Ny+Hy-1) - sizes = (Sx, Sy, Sx, Sy) - offs = (Oxᴸ, Oyᴸ, Oxᴿ, Oyᴿ) - - return boundary_parameters(sizes, offs, grid, arch) + params = (param_west, param_east, param_south, param_north) + + return boundary_parameters(params, grid, arch) end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index 7bd3ee6337..6b5a03fd00 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -3,7 +3,7 @@ using OrderedCollections: OrderedDict using Oceananigans.DistributedComputations using Oceananigans.Architectures: AbstractArchitecture -using Oceananigans.Advection: AbstractAdvectionScheme, CenteredSecondOrder, VectorInvariant +using Oceananigans.Advection: AbstractAdvectionScheme, CenteredSecondOrder, VectorInvariant, adapt_advection_order using Oceananigans.BuoyancyModels: validate_buoyancy, regularize_buoyancy, SeawaterBuoyancy, g_Earth using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions using Oceananigans.Biogeochemistry: validate_biogeochemistry, AbstractBiogeochemistry, biogeochemical_auxiliary_fields @@ -130,6 +130,10 @@ function HydrostaticFreeSurfaceModel(; grid, tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example) + # Reduce the advection order in directions that do not have enough grid points + momentum_advection = adapt_advection_order(momentum_advection, grid) + tracer_advection = adapt_advection_order(tracer_advection, grid) + tracers, auxiliary_fields = validate_biogeochemistry(tracers, merge(auxiliary_fields, biogeochemical_auxiliary_fields(biogeochemistry)), biogeochemistry, grid, clock) validate_buoyancy(buoyancy, tracernames(tracers)) buoyancy = regularize_buoyancy(buoyancy) diff --git a/src/Models/NonhydrostaticModels/compute_nonhydrostatic_boundary_tendencies.jl b/src/Models/NonhydrostaticModels/compute_nonhydrostatic_boundary_tendencies.jl index 933f732d8d..cfe3691075 100644 --- a/src/Models/NonhydrostaticModels/compute_nonhydrostatic_boundary_tendencies.jl +++ b/src/Models/NonhydrostaticModels/compute_nonhydrostatic_boundary_tendencies.jl @@ -1,6 +1,6 @@ import Oceananigans.Models: compute_boundary_tendencies! -using Oceananigans.TurbulenceClosures: required_halo_size +using Oceananigans.TurbulenceClosures: required_halo_size_x, required_halo_size_y using Oceananigans.Grids: XFlatGrid, YFlatGrid # TODO: the code in this file is difficult to understand. @@ -28,70 +28,60 @@ end function boundary_tendency_kernel_parameters(grid, arch) Nx, Ny, Nz = size(grid) Hx, Hy, _ = halo_size(grid) - - Sx = (Hx, Ny, Nz) - Sy = (Nx, Hy, Nz) - Oᴸ = (0, 0, 0) - Oxᴿ = (Nx-Hx, 0, 0) - Oyᴿ = (0, Ny-Hy, 0) + param_west = (1:Hx, 1:Ny, 1:Nz) + param_east = (Nx-Hx+1:Nx, 1:Ny, 1:Nz) + param_south = (1:Nx, 1:Hy, 1:Nz) + param_north = (1:Nx, Ny-Hy+1:Ny, 1:Nz) - sizes = (Sx, Sy, Sx, Sy) - offs = (Oᴸ, Oᴸ, Oxᴿ, Oyᴿ) + params = (param_west, param_east, param_south, param_north) - return boundary_parameters(sizes, offs, grid, arch) + return boundary_parameters(params, grid, arch) end # p needs computing in the range 0 : 0 and N + 1 : N + 1 function boundary_p_kernel_parameters(grid, arch) Nx, Ny, _ = size(grid) - Sx = (1, Ny) - Sy = (Nx, 1) + param_west = (0:0, 1:Ny) + param_east = (Nx+1:Nx+1, 1:Ny) + param_south = (1:Nx, 0:0) + param_north = (1:Nx, Ny+1:Ny+1) - Oxᴸ = (-1, 0) - Oyᴸ = (0, -1) - Oxᴿ = (Nx, 0) - Oyᴿ = (0, Ny) + params = (param_west, param_east, param_south, param_north) - sizes = (Sx, Sy, Sx, Sy) - offs = (Oxᴸ, Oyᴸ, Oxᴿ, Oyᴿ) - - return boundary_parameters(sizes, offs, grid, arch) + return boundary_parameters(params, grid, arch) end # diffusivities need recomputing in the range 0 : B and N - B + 1 : N + 1 function boundary_κ_kernel_parameters(grid, closure, arch) Nx, Ny, Nz = size(grid) - B = required_halo_size(closure) - - Sx = (B+1, Ny, Nz) - Sy = (Nx, B+1, Nz) + Bx = required_halo_size_x(closure) + By = required_halo_size_y(closure) - Oxᴸ = (-1, 0, 0) - Oyᴸ = (0, -1, 0) - Oxᴿ = (Nx-B, 0, 0) - Oyᴿ = (0, Ny-B, 0) + param_west = (0:Bx, 1:Ny, 1:Nz) + param_east = (Nx-Bx+1:Nx+1, 1:Ny, 1:Nz) + param_south = (1:Nx, 0:By, 1:Nz) + param_north = (1:Nx, Ny-By+1:Ny+1, 1:Nz) - sizes = (Sx, Sy, Sx, Sy) - offs = (Oxᴸ, Oyᴸ, Oxᴿ, Oyᴿ) + params = (param_west, param_east, param_south, param_north) - return boundary_parameters(sizes, offs, grid, arch) + return boundary_parameters(params, grid, arch) end # Recompute only on communicating sides -function boundary_parameters(S, O, grid, arch) +function boundary_parameters(parameters, grid, arch) Rx, Ry, _ = arch.ranks Tx, Ty, _ = topology(grid) - include_xᴸ = !isa(grid, XFlatGrid) && (Rx != 1) && !(Tx == RightConnected) - include_yᴸ = !isa(grid, YFlatGrid) && (Ry != 1) && !(Ty == RightConnected) - include_xᴿ = !isa(grid, XFlatGrid) && (Rx != 1) && !(Tx == LeftConnected) - include_yᴿ = !isa(grid, YFlatGrid) && (Ry != 1) && !(Ty == LeftConnected) + include_west = !isa(grid, XFlatGrid) && (Rx != 1) && !(Tx == RightConnected) + include_east = !isa(grid, XFlatGrid) && (Rx != 1) && !(Tx == LeftConnected) + include_south = !isa(grid, YFlatGrid) && (Ry != 1) && !(Ty == RightConnected) + include_north = !isa(grid, YFlatGrid) && (Ry != 1) && !(Ty == LeftConnected) - include_side = (include_xᴸ, include_yᴸ, include_xᴿ, include_yᴿ) + include_side = (include_west, include_east, include_south, include_north) - return Tuple(KernelParameters(S[i], O[i]) for i in findall(include_side)) + return Tuple(KernelParameters(parameters[i]...) for i in findall(include_side)) end diff --git a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl index 6769891656..43b41d5715 100644 --- a/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl +++ b/src/Models/NonhydrostaticModels/nonhydrostatic_model.jl @@ -3,7 +3,7 @@ using OrderedCollections: OrderedDict using Oceananigans.Architectures: AbstractArchitecture using Oceananigans.DistributedComputations: Distributed -using Oceananigans.Advection: CenteredSecondOrder +using Oceananigans.Advection: CenteredSecondOrder, adapt_advection_order using Oceananigans.BuoyancyModels: validate_buoyancy, regularize_buoyancy, SeawaterBuoyancy using Oceananigans.Biogeochemistry: validate_biogeochemistry, AbstractBiogeochemistry, biogeochemical_auxiliary_fields using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions @@ -172,6 +172,11 @@ function NonhydrostaticModel(; grid, validate_buoyancy(buoyancy, tracernames(tracers)) buoyancy = regularize_buoyancy(buoyancy) + # Adjust advection scheme to be valid on a particular grid size. i.e. if the grid size + # is smaller than the advection order, reduce the order of the advection in that particular + # direction + advection = adapt_advection_order(advection, grid) + # Adjust halos when the advection scheme or turbulence closure requires it. # Note that halos are isotropic by default; however we respect user-input here # by adjusting each (x, y, z) halo individually. diff --git a/src/MultiRegion/multi_region_models.jl b/src/MultiRegion/multi_region_models.jl index 9d835342f7..95f677e434 100644 --- a/src/MultiRegion/multi_region_models.jl +++ b/src/MultiRegion/multi_region_models.jl @@ -9,12 +9,17 @@ using Oceananigans.Advection: OnlySelfUpwinding, CrossAndSelfUpwinding using Oceananigans.ImmersedBoundaries: GridFittedBottom, PartialCellBottom, GridFittedBoundary using Oceananigans.Solvers: ConjugateGradientSolver -import Oceananigans.Advection: WENO, cell_advection_timescale +import Oceananigans.Advection: WENO, cell_advection_timescale, adapt_advection_order import Oceananigans.Models.HydrostaticFreeSurfaceModels: build_implicit_step_solver, validate_tracer_advection import Oceananigans.TurbulenceClosures: implicit_diffusion_solver const MultiRegionModel = HydrostaticFreeSurfaceModel{<:Any, <:Any, <:AbstractArchitecture, <:Any, <:MultiRegionGrids} +function adapt_advection_order(advection::MultiRegionObject, grid::MultiRegionGrids) + @apply_regionally new_advection = adapt_advection_order(advection, grid) + return new_advection +end + # Utility to generate the inputs to complex `getregion`s function getregionalproperties(T, inner=true) type = getglobal(@__MODULE__, T) diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index 2abd97e7a0..52935f1a0e 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -34,7 +34,7 @@ export UpwindBiased, UpwindBiasedFirstOrder, UpwindBiasedThirdOrder, UpwindBiasedFifthOrder, WENO, WENOThirdOrder, WENOFifthOrder, VectorInvariant, WENOVectorInvariant, EnergyConserving, EnstrophyConserving, - TracerAdvection, + FluxFormAdvection, # Boundary conditions BoundaryCondition, diff --git a/src/TurbulenceClosures/TurbulenceClosures.jl b/src/TurbulenceClosures/TurbulenceClosures.jl index f810ce5e2f..a99d3481f7 100644 --- a/src/TurbulenceClosures/TurbulenceClosures.jl +++ b/src/TurbulenceClosures/TurbulenceClosures.jl @@ -50,7 +50,7 @@ using Oceananigans.Utils using Oceananigans.Architectures: AbstractArchitecture, device using Oceananigans.Fields: FunctionField -import Oceananigans.Advection: required_halo_size +import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z import Oceananigans.Architectures: on_architecture const VerticallyBoundedGrid{FT} = AbstractGrid{FT, <:Any, <:Any, <:Bounded} @@ -77,7 +77,9 @@ compute_diffusivities!(K, closure::AbstractTurbulenceClosure, args...; kwargs... # point at each side to calculate viscous fluxes at the edge of the domain. # If diffusivity itself requires one halo to be computed (e.g. κ = ℑxᶠᵃᵃ(i, j, k, grid, ℑxᶜᵃᵃ, T), # or `AnisotropicMinimumDissipation` and `SmagorinskyLilly`) then B = 2 -@inline required_halo_size(::AbstractTurbulenceClosure{TD, B}) where {TD, B} = B +@inline required_halo_size_x(::AbstractTurbulenceClosure{TD, B}) where {TD, B} = B +@inline required_halo_size_y(::AbstractTurbulenceClosure{TD, B}) where {TD, B} = B +@inline required_halo_size_z(::AbstractTurbulenceClosure{TD, B}) where {TD, B} = B const ClosureKinda = Union{Nothing, AbstractTurbulenceClosure, AbstractArray{<:AbstractTurbulenceClosure}} add_closure_specific_boundary_conditions(closure::ClosureKinda, bcs, args...) = bcs diff --git a/src/TurbulenceClosures/closure_tuples.jl b/src/TurbulenceClosures/closure_tuples.jl index 6321800b14..b552331bb0 100644 --- a/src/TurbulenceClosures/closure_tuples.jl +++ b/src/TurbulenceClosures/closure_tuples.jl @@ -89,7 +89,9 @@ function add_closure_specific_boundary_conditions(closure_tuple::Tuple, bcs, arg return bcs end -required_halo_size(closure_tuple::Tuple) = maximum(map(required_halo_size, closure_tuple)) +required_halo_size_x(closure_tuple::Tuple) = maximum(map(required_halo_size_x, closure_tuple)) +required_halo_size_y(closure_tuple::Tuple) = maximum(map(required_halo_size_y, closure_tuple)) +required_halo_size_z(closure_tuple::Tuple) = maximum(map(required_halo_size_z, closure_tuple)) ##### ##### Compiler-inferrable time_discretization for tuples diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl index e79a1a76d1..87f6ceff5f 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl @@ -37,15 +37,15 @@ function IsopycnalSkewSymmetricDiffusivity(time_disc::TD = VerticallyImplicitTim κ_symmetric = 0, isopycnal_tensor = SmallSlopeIsopycnalTensor(), slope_limiter = FluxTapering(1e-2), - required_halo_size = 1) where TD + required_halo_size::Int = 1) where TD isopycnal_tensor isa SmallSlopeIsopycnalTensor || error("Only isopycnal_tensor=SmallSlopeIsopycnalTensor() is currently supported.") return IsopycnalSkewSymmetricDiffusivity{TD, required_halo_size}(convert_diffusivity(FT, κ_skew), - convert_diffusivity(FT, κ_symmetric), - isopycnal_tensor, - slope_limiter) + convert_diffusivity(FT, κ_symmetric), + isopycnal_tensor, + slope_limiter) end IsopycnalSkewSymmetricDiffusivity(FT::DataType; kw...) = diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl index 5d5056461c..003fabbeea 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl @@ -1,4 +1,4 @@ -import Oceananigans.Grids: required_halo_size +import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z using Oceananigans.Utils: prettysummary """ @@ -6,10 +6,10 @@ using Oceananigans.Utils: prettysummary Holds viscosity and diffusivities for models with prescribed isotropic diffusivities. """ -struct ScalarBiharmonicDiffusivity{F, V, K, N} <: AbstractScalarBiharmonicDiffusivity{F, N} +struct ScalarBiharmonicDiffusivity{F, N, V, K} <: AbstractScalarBiharmonicDiffusivity{F, N} ν :: V κ :: K - ScalarBiharmonicDiffusivity{F, N}(ν::V, κ::K) where {F, V, K, N} = new{F, V, K, N}(ν, κ) + ScalarBiharmonicDiffusivity{F, N}(ν::V, κ::K) where {F, V, K, N} = new{F, N, V, K}(ν, κ) end # Aliases that allow specify the floating type, assuming that the discretization is Explicit in time @@ -50,6 +50,9 @@ Keyword arguments * `discrete_form`: `Boolean`; default: `false`. +* `required_halo_size = 2`: the required halo size for the closure. This value should be an integer. + change only if using a function for `ν` or `κ` that requires a halo size larger than 1 to compute. + When prescribing the viscosities or diffusivities as functions, depending on the value of keyword argument `discrete_form`, the constructor expects: @@ -75,14 +78,22 @@ function ScalarBiharmonicDiffusivity(formulation = ThreeDimensionalFormulation() discrete_form = false, loc = (nothing, nothing, nothing), parameters = nothing, - required_halo_size = 2) + required_halo_size::Int = 2) ν = convert_diffusivity(FT, ν; discrete_form, loc, parameters) κ = convert_diffusivity(FT, κ; discrete_form, loc, parameters) + + # Force a type-stable constructor if ν and κ are numbers + # This particular short-circuiting of the required_halo_size kwargs is necessary to perform parameter + # estimation of the diffusivity coefficients using autodiff. + if ν isa Number && κ isa Number + return ScalarBiharmonicDiffusivity{typeof(formulation), 2}(ν, κ) + end + return ScalarBiharmonicDiffusivity{typeof(formulation), required_halo_size}(ν, κ) end -function with_tracers(tracers, closure::ScalarBiharmonicDiffusivity{F, N}) where {F, N} +function with_tracers(tracers, closure::ScalarBiharmonicDiffusivity{F, N, V, K}) where {F, N, V, K} κ = tracer_diffusivities(tracers, closure.κ) return ScalarBiharmonicDiffusivity{F, N}(closure.ν, κ) end @@ -116,4 +127,4 @@ function on_architecture(to, closure::ScalarBiharmonicDiffusivity{F, <:Any, <:An ν = on_architecture(to, closure.ν) κ = on_architecture(to, closure.κ) return ScalarBiharmonicDiffusivity{F, N}(ν, κ) -end +end \ No newline at end of file diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/scalar_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/scalar_diffusivity.jl index d17f1eb72e..a4f85814b8 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/scalar_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/scalar_diffusivity.jl @@ -1,12 +1,12 @@ using Oceananigans.Utils: prettysummary import Adapt -import Oceananigans.Grids: required_halo_size +import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z -struct ScalarDiffusivity{TD, F, V, K, N} <: AbstractScalarDiffusivity{TD, F, N} +struct ScalarDiffusivity{TD, F, N, V, K} <: AbstractScalarDiffusivity{TD, F, N} ν :: V κ :: K - ScalarDiffusivity{TD, F, N}(ν::V, κ::K) where {TD, F, V, K, N} = new{TD, F, V, K, N}(ν, κ) + ScalarDiffusivity{TD, F, N}(ν::V, κ::K) where {TD, F, N, V, K} = new{TD, F, N, V, K}(ν, κ) end """ @@ -63,6 +63,9 @@ value of keyword argument `discrete_form`, the constructor expects: - with `loc = (ℓx, ℓy, ℓz)` and specified `parameters`: functions of `(i, j, k, grid, clock, fields, parameters)`. +* `required_halo_size = 1`: the required halo size for the closure. This value should be an integer. + change only if using a function for `ν` or `κ` that requires a halo size larger than 1 to compute. + * `parameters`: `NamedTuple` with parameters used by the functions that compute viscosity and/or diffusivity; default: `nothing`. @@ -111,21 +114,29 @@ ScalarDiffusivity{ExplicitTimeDiscretization}(ν=0.0, κ=Oceananigans.Turbulence ``` """ function ScalarDiffusivity(time_discretization=ExplicitTimeDiscretization(), - formulation=ThreeDimensionalFormulation(), FT=Float64; + formulation=ThreeDimensionalFormulation(), + FT=Float64; ν=0, κ=0, discrete_form = false, loc = (nothing, nothing, nothing), parameters = nothing, - required_halo_size = 1) + required_halo_size::Int = 1) if formulation == HorizontalFormulation() && time_discretization == VerticallyImplicitTimeDiscretization() - throw(ArgumentError("VerticallyImplicitTimeDiscretization is only supported for \ + throw(ArgumentError("VerticallyImplicitTimeDiscretization is only supported for \ `VerticalFormulation` or `ThreeDimensionalFormulation`")) end κ = convert_diffusivity(FT, κ; discrete_form, loc, parameters) ν = convert_diffusivity(FT, ν; discrete_form, loc, parameters) + # Force a type-stable constructor if ν and κ are numbers + # This particular short-circuiting of the required_halo_size kwargs is necessary to perform parameter + # estimation of the diffusivity coefficients using autodiff. + if ν isa Number && κ isa Number + return ScalarDiffusivity{typeof(time_discretization), typeof(formulation), 1}(ν, κ) + end + return ScalarDiffusivity{typeof(time_discretization), typeof(formulation), required_halo_size}(ν, κ) end @@ -173,8 +184,6 @@ Shorthand for a `ScalarDiffusivity` with `HorizontalDivergenceFormulation()`. Se HorizontalScalarDiffusivity(FT::DataType; kwargs...) = ScalarDiffusivity(ExplicitTimeDiscretization(), HorizontalFormulation(), FT; kwargs...) HorizontalDivergenceScalarDiffusivity(FT::DataType; kwargs...) = ScalarDiffusivity(ExplicitTimeDiscretization(), HorizontalDivergenceFormulation(), FT; kwargs...) -required_halo_size(closure::ScalarDiffusivity) = 1 - @inline function with_tracers(tracers, closure::ScalarDiffusivity{TD, F, N}) where {TD, F, N} κ = tracer_diffusivities(tracers, closure.κ) return ScalarDiffusivity{TD, F, N}(closure.ν, κ) @@ -218,4 +227,4 @@ function on_architecture(to, closure::ScalarDiffusivity{TD, F, <:Any, <:Any, N}) ν = on_architecture(to, closure.ν) κ = on_architecture(to, closure.κ) return ScalarDiffusivity{TD, F, N}(ν, κ) -end +end \ No newline at end of file diff --git a/test/test_immersed_advection.jl b/test/test_immersed_advection.jl index aaae16a899..dc36776b6a 100644 --- a/test/test_immersed_advection.jl +++ b/test/test_immersed_advection.jl @@ -10,7 +10,7 @@ using Oceananigans.Advection: _biased_interpolate_xᶠᵃᵃ, _biased_interpolate_yᵃᶜᵃ, _biased_interpolate_yᵃᶠᵃ, - TracerAdvection + FluxFormAdvection advection_schemes = [Centered, UpwindBiased, WENO] @@ -128,7 +128,7 @@ for arch in archs for adv in advection_schemes, buffer in [1, 2, 3, 4, 5] directional_scheme = adv(order = advective_order(buffer, adv)) - scheme = TracerAdvection(directional_scheme, directional_scheme, directional_scheme) + scheme = FluxFormAdvection(directional_scheme, directional_scheme, directional_scheme) for g in [grid, ibg] @info " Testing immersed tracer conservation [$(typeof(arch)), $(summary(scheme)), $(typeof(g).name.wrapper)]" run_tracer_conservation_test(g, scheme) diff --git a/test/test_nonhydrostatic_models.jl b/test/test_nonhydrostatic_models.jl index dedc95c9ce..bd10653bfd 100644 --- a/test/test_nonhydrostatic_models.jl +++ b/test/test_nonhydrostatic_models.jl @@ -1,5 +1,7 @@ include("dependencies_for_runtests.jl") +using Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z + @testset "Models" begin @info "Testing models..." @@ -66,6 +68,28 @@ include("dependencies_for_runtests.jl") model = NonhydrostaticModel(closure=ScalarBiharmonicDiffusivity(), grid=funny_grid) @test model.grid.Hx == 2 && model.grid.Hy == 3 && model.grid.Hz == 4 + + @info " Testing adjustment of advection schemes in NonhydrostaticModel constructor..." + small_grid = RectilinearGrid(size=(4, 2, 4), extent=(1, 2, 3), halo=(1, 1, 1)) + + # Model ensures that halos are at least of size 1 + model = NonhydrostaticModel(grid=small_grid, advection=WENO()) + @test model.advection isa FluxFormAdvection + @test required_halo_size_x(model.advection) == 3 + @test required_halo_size_y(model.advection) == 2 + @test required_halo_size_z(model.advection) == 3 + + model = NonhydrostaticModel(grid=small_grid, advection=UpwindBiased(; order = 9)) + @test model.advection isa FluxFormAdvection + @test required_halo_size_x(model.advection) == 4 + @test required_halo_size_y(model.advection) == 2 + @test required_halo_size_z(model.advection) == 4 + + model = NonhydrostaticModel(grid=small_grid, advection=Centered(; order = 10)) + @test model.advection isa FluxFormAdvection + @test required_halo_size_x(model.advection) == 4 + @test required_halo_size_y(model.advection) == 2 + @test required_halo_size_z(model.advection) == 4 end @testset "Model construction with single tracer and nothing tracer" begin diff --git a/test/test_turbulence_closures.jl b/test/test_turbulence_closures.jl index d2e5cafb95..00b8e34bef 100644 --- a/test/test_turbulence_closures.jl +++ b/test/test_turbulence_closures.jl @@ -1,8 +1,9 @@ include("dependencies_for_runtests.jl") -using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity, RiBasedVerticalDiffusivity +using Oceananigans.TurbulenceClosures: CATKEVerticalDiffusivity, RiBasedVerticalDiffusivity, DiscreteDiffusionFunction -using Oceananigans.TurbulenceClosures: viscosity_location, diffusivity_location +using Oceananigans.TurbulenceClosures: viscosity_location, diffusivity_location, + required_halo_size_x, required_halo_size_y, required_halo_size_z using Oceananigans.TurbulenceClosures: diffusive_flux_x, diffusive_flux_y, diffusive_flux_z, viscous_flux_ux, viscous_flux_uy, viscous_flux_uz, @@ -246,6 +247,26 @@ end @test closure.κ.T == T(κ) run_constant_isotropic_diffusivity_fluxdiv_tests(T) end + + @info " Testing ScalarDiffusivity with different halo requirements..." + closure = ScalarDiffusivity(ν=0.3) + @test required_halo_size_x(closure) == 1 + @test required_halo_size_y(closure) == 1 + @test required_halo_size_z(closure) == 1 + + closure = ScalarBiharmonicDiffusivity(ν=0.3) + @test required_halo_size_x(closure) == 2 + @test required_halo_size_y(closure) == 2 + @test required_halo_size_z(closure) == 2 + + @inline ν(i, j, k, grid, ℓx, ℓy, ℓz, clock, fields) = ℑxᶠᵃᵃ(i, j, k, grid, ℑxᶜᵃᵃ, fields.u) + closure = ScalarDiffusivity(; ν, discrete_form=true, required_halo_size=2) + + @test closure.ν isa DiscreteDiffusionFunction + @test required_halo_size_x(closure) == 2 + @test required_halo_size_y(closure) == 2 + @test required_halo_size_z(closure) == 2 + end @testset "HorizontalScalarDiffusivity" begin