diff --git a/src/TurbulenceClosures/abstract_scalar_biharmonic_diffusivity_closure.jl b/src/TurbulenceClosures/abstract_scalar_biharmonic_diffusivity_closure.jl index 357970226a..a41752adba 100644 --- a/src/TurbulenceClosures/abstract_scalar_biharmonic_diffusivity_closure.jl +++ b/src/TurbulenceClosures/abstract_scalar_biharmonic_diffusivity_closure.jl @@ -5,11 +5,12 @@ using Oceananigans.Grids: peripheral_node Abstract type for closures with scalar biharmonic diffusivities. """ -abstract type AbstractScalarBiharmonicDiffusivity{F, N} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, N} end +abstract type AbstractScalarBiharmonicDiffusivity{F, N, VI} <: AbstractTurbulenceClosure{ExplicitTimeDiscretization, N} end @inline formulation(::AbstractScalarBiharmonicDiffusivity{F}) where {F} = F() const ASBD = AbstractScalarBiharmonicDiffusivity +const VectorInvariantASBD = AbstractScalarBiharmonicDiffusivity{<:HorizontalFormulation, <:Nothing, <:VectorInvariantForm} ##### ##### Coefficient extractors @@ -34,6 +35,11 @@ const AHBD = AbstractScalarBiharmonicDiffusivity{<:HorizontalFormulation} const ADBD = AbstractScalarBiharmonicDiffusivity{<:HorizontalDivergenceFormulation} const AVBD = AbstractScalarBiharmonicDiffusivity{<:VerticalFormulation} +@inline ν_σᶜᶜᶜ(i, j, k, grid, closure::AHBD, K, clock, fields, σᶜᶜᶜ, args...) = νᶜᶜᶜ(i, j, k, grid, closure, K, clock, fields) * σᶜᶜᶜ(i, j, k, grid, closure, args...) +@inline ν_σᶠᶠᶜ(i, j, k, grid, closure::AHBD, K, clock, fields, σᶠᶠᶜ, args...) = νᶠᶠᶜ(i, j, k, grid, closure, K, clock, fields) * σᶠᶠᶜ(i, j, k, grid, closure, args...) +@inline ν_σᶠᶜᶠ(i, j, k, grid, closure::AHBD, K, clock, fields, σᶠᶜᶠ, args...) = νᶠᶜᶠ(i, j, k, grid, closure, K, clock, fields) * σᶠᶜᶠ(i, j, k, grid, closure, args...) +@inline ν_σᶜᶠᶠ(i, j, k, grid, closure::AHBD, K, clock, fields, σᶜᶠᶠ, args...) = νᶜᶠᶠ(i, j, k, grid, closure, K, clock, fields) * σᶜᶠᶠ(i, j, k, grid, closure, args...) + @inline viscous_flux_ux(i, j, k, grid, closure::AIBD, K, clk, fields, b) = + ν_σᶜᶜᶜ(i, j, k, grid, closure, K, clk, fields, ∂xᶜᶜᶜ, biharmonic_mask_x, ∇²ᶠᶜᶜ, fields.u) @inline viscous_flux_vx(i, j, k, grid, closure::AIBD, K, clk, fields, b) = + ν_σᶠᶠᶜ(i, j, k, grid, closure, K, clk, fields, biharmonic_mask_x, ∂xᶠᶠᶜ, ∇²ᶜᶠᶜ, fields.v) @inline viscous_flux_wx(i, j, k, grid, closure::AIBD, K, clk, fields, b) = + ν_σᶠᶜᶠ(i, j, k, grid, closure, K, clk, fields, biharmonic_mask_x, ∂xᶠᶜᶠ, ∇²ᶜᶜᶠ, fields.w) @@ -52,7 +58,6 @@ const AVBD = AbstractScalarBiharmonicDiffusivity{<:VerticalFormulation} @inline viscous_flux_uz(i, j, k, grid, closure::AVBD, K, clk, fields, b) = + ν_σᶠᶜᶠ(i, j, k, grid, closure, K, clk, fields, biharmonic_mask_z, ∂zᶠᶜᶠ, ∂²zᶠᶜᶜ, fields.u) @inline viscous_flux_vz(i, j, k, grid, closure::AVBD, K, clk, fields, b) = + ν_σᶜᶠᶠ(i, j, k, grid, closure, K, clk, fields, biharmonic_mask_z, ∂zᶜᶠᶠ, ∂²zᶜᶠᶜ, fields.v) @inline viscous_flux_wz(i, j, k, grid, closure::AVBD, K, clk, fields, b) = + ν_σᶜᶜᶜ(i, j, k, grid, closure, K, clk, fields, ∂zᶜᶜᶜ, biharmonic_mask_z, ∂²zᶜᶜᶠ, fields.w) - @inline viscous_flux_ux(i, j, k, grid, closure::ADBD, K, clk, fields, b) = + ν_σᶜᶜᶜ(i, j, k, grid, closure, K, clk, fields, δ★ᶜᶜᶜ, fields.u, fields.v) @inline viscous_flux_vy(i, j, k, grid, closure::ADBD, K, clk, fields, b) = + ν_σᶜᶜᶜ(i, j, k, grid, closure, K, clk, fields, δ★ᶜᶜᶜ, fields.u, fields.v) @@ -71,27 +76,35 @@ const AVBD = AbstractScalarBiharmonicDiffusivity{<:VerticalFormulation} ##### Biharmonic-specific viscous operators ##### -# See https://mitgcm.readthedocs.io/en/latest/algorithm/algorithm.html#horizontal-dissipation -@inline function δ★ᶜᶜᶜ(i, j, k, grid, u, v) +@inline ∇²u_vector_invariantᶠᶜᶜ(i, j, k, grid, u, v) = δxᶠᶜᶜ(i, j, k, grid, div_xyᶜᶜᶜ, u, v) - δyᶠᶜᶜ(i, j, k, grid, ζ₃ᶠᶠᶜ, u, v) +@inline ∇²v_vector_invariantᶜᶠᶜ(i, j, k, grid, u, v) = δxᶜᶠᶜ(i, j, k, grid, ζ₃ᶠᶠᶜ, u, v) + δyᶜᶠᶜ(i, j, k, grid, div_xyᶜᶜᶜ, u, v) - # These closures seem to be needed to help the compiler infer types - # (either of u and v or of the function arguments) - @inline Δy_∇²u(i, j, k, grid, u) = Δy_qᶠᶜᶜ(i, j, k, grid, biharmonic_mask_x, ∇²hᶠᶜᶜ, u) - @inline Δx_∇²v(i, j, k, grid, v) = Δx_qᶜᶠᶜ(i, j, k, grid, biharmonic_mask_y, ∇²hᶜᶠᶜ, v) +# These closures seem to be needed to help the compiler infer types +# (either of u and v or of the function arguments) +@inline Δy_∇²u(i, j, k, grid, closure, u, v) = Δy_qᶠᶜᶜ(i, j, k, grid, biharmonic_mask_x, ∇²hᶠᶜᶜ, u) +@inline Δx_∇²v(i, j, k, grid, closure, u, v) = Δx_qᶜᶠᶜ(i, j, k, grid, biharmonic_mask_y, ∇²hᶜᶠᶜ, v) - return 1 / Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Δy_∇²u, u) + - δyᵃᶜᵃ(i, j, k, grid, Δx_∇²v, v)) -end +@inline Δy_∇²v(i, j, k, grid, closure, u, v) = Δy_qᶜᶠᶜ(i, j, k, grid, biharmonic_mask_y, ∇²hᶜᶠᶜ, v) +@inline Δx_∇²u(i, j, k, grid, closure, u, v) = Δx_qᶠᶜᶜ(i, j, k, grid, biharmonic_mask_x, ∇²hᶠᶜᶜ, u) -@inline function ζ★ᶠᶠᶜ(i, j, k, grid, u, v) +@inline Δy_∇²u(i, j, k, grid, ::VectorInvariantASBD, u, v) = Δy_qᶠᶜᶜ(i, j, k, grid, biharmonic_mask_x, ∇²u_vector_invariantᶠᶜᶜ, u, v) +@inline Δx_∇²v(i, j, k, grid, ::VectorInvariantASBD, u, v) = Δx_qᶜᶠᶜ(i, j, k, grid, biharmonic_mask_y, ∇²v_vector_invariantᶜᶠᶜ, u, v) - # These closures seem to be needed to help the compiler infer types - # (either of u and v or of the function arguments) - @inline Δy_∇²v(i, j, k, grid, v) = Δy_qᶜᶠᶜ(i, j, k, grid, biharmonic_mask_y, ∇²hᶜᶠᶜ, v) - @inline Δx_∇²u(i, j, k, grid, u) = Δx_qᶠᶜᶜ(i, j, k, grid, biharmonic_mask_x, ∇²hᶠᶜᶜ, u) +# See https://mitgcm.readthedocs.io/en/latest/algorithm/algorithm.html#horizontal-dissipation +@inline function δ★ᶜᶜᶜ(i, j, k, grid, closure, u, v) + return 1 / Azᶜᶜᶜ(i, j, k, grid) * (δxᶜᵃᵃ(i, j, k, grid, Δy_∇²u, closure, u, v) + + δyᵃᶜᵃ(i, j, k, grid, Δx_∇²v, closure, u, v)) +end + +@inline function ζ★ᶠᶠᶜ(i, j, k, grid, closure, u, v) + return 1 / Azᶠᶠᶜ(i, j, k, grid) * (δxᶠᵃᵃ(i, j, k, grid, Δy_∇²v, closure, u, v) - + δyᵃᶠᵃ(i, j, k, grid, Δx_∇²u, closure, u, v)) +end - return 1 / Azᶠᶠᶜ(i, j, k, grid) * (δxᶠᵃᵃ(i, j, k, grid, Δy_∇²v, v) - - δyᵃᶠᵃ(i, j, k, grid, Δx_∇²u, u)) +@inline function ζ★ᶠᶠᶜ(i, j, k, grid, ::VectorInvariantASBD, u, v) + ∇²u = ∇²u_vector_invariantᶠᶜᶜ(i, j, k, grid, u, v) + ∇²v = ∇²v_vector_invariantᶜᶠᶜ(i, j, k, grid, u, v) + return ζ₃ᶠᶠᶜ(i, j, k, grid, ∇²u, ∇²v) end ##### diff --git a/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl b/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl index 01e322645d..20b7e66679 100644 --- a/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl +++ b/src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl @@ -27,6 +27,13 @@ Specifies a horizontally-isotropic, `VectorInvariant`, `ScalarDiffusivity`. """ struct HorizontalFormulation <: AbstractDiffusivityFormulation end +""" + struct VectorInvariantForm end + +Specifies a `VectorInvariant` `ScalarDiffusivity`. +""" +struct VectorInvariantForm <: AbstractDiffusivityFormulation end + """ struct HorizontalDivergenceFormulation end diff --git a/src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl b/src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl index 003fabbeea..0169e9217b 100644 --- a/src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl +++ b/src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl @@ -2,14 +2,14 @@ import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_ using Oceananigans.Utils: prettysummary """ - struct ScalarBiharmonicDiffusivity{F, N, K} <: AbstractScalarBiharmonicDiffusivity{F} + struct ScalarBiharmonicDiffusivity{F, N, VI, V, K} <: AbstractScalarBiharmonicDiffusivity{F} Holds viscosity and diffusivities for models with prescribed isotropic diffusivities. """ -struct ScalarBiharmonicDiffusivity{F, N, V, K} <: AbstractScalarBiharmonicDiffusivity{F, N} +struct ScalarBiharmonicDiffusivity{F, N, VI, V, K} <: AbstractScalarBiharmonicDiffusivity{F, N, VI} ν :: V κ :: K - ScalarBiharmonicDiffusivity{F, N}(ν::V, κ::K) where {F, V, K, N} = new{F, N, V, K}(ν, κ) + ScalarBiharmonicDiffusivity{F, N, VI}(ν::V, κ::K) where {F, N, VI, V, K} = new{F, N, VI, V, K}(ν, κ) end # Aliases that allow specify the floating type, assuming that the discretization is Explicit in time @@ -76,26 +76,29 @@ function ScalarBiharmonicDiffusivity(formulation = ThreeDimensionalFormulation() ν = 0, κ = 0, discrete_form = false, + vector_invariant_form = true, loc = (nothing, nothing, nothing), parameters = nothing, required_halo_size::Int = 2) ν = convert_diffusivity(FT, ν; discrete_form, loc, parameters) κ = convert_diffusivity(FT, κ; discrete_form, loc, parameters) - + + VI = vector_invariant_form ? VectorInvariantForm : nothing + # 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}(ν, κ) + return ScalarBiharmonicDiffusivity{typeof(formulation), 2, VI}(ν, κ) end - return ScalarBiharmonicDiffusivity{typeof(formulation), required_halo_size}(ν, κ) + return ScalarBiharmonicDiffusivity{typeof(formulation), required_halo_size, VI}(ν, κ) end -function with_tracers(tracers, closure::ScalarBiharmonicDiffusivity{F, N, V, K}) where {F, N, V, K} +function with_tracers(tracers, closure::ScalarBiharmonicDiffusivity{F, N, VI}) where {F, N, VI} κ = tracer_diffusivities(tracers, closure.κ) - return ScalarBiharmonicDiffusivity{F, N}(closure.ν, κ) + return ScalarBiharmonicDiffusivity{F, N, VI}(closure.ν, κ) end @inline viscosity(closure::ScalarBiharmonicDiffusivity, K) = closure.ν @@ -117,14 +120,14 @@ end Base.show(io::IO, closure::ScalarBiharmonicDiffusivity) = print(io, summary(closure)) -function Adapt.adapt_structure(to, closure::ScalarBiharmonicDiffusivity{F, <:Any, <:Any, N}) where {F, N} +function Adapt.adapt_structure(to, closure::ScalarBiharmonicDiffusivity{F, N, VI, <:Any, <:Any}) where {F, N, VI} ν = Adapt.adapt(to, closure.ν) κ = Adapt.adapt(to, closure.κ) - return ScalarBiharmonicDiffusivity{F, N}(ν, κ) + return ScalarBiharmonicDiffusivity{F, N, VI}(ν, κ) end -function on_architecture(to, closure::ScalarBiharmonicDiffusivity{F, <:Any, <:Any, N}) where {F, N} +function on_architecture(to, closure::ScalarBiharmonicDiffusivity{F, N, VI}) where {F, N, VI} ν = on_architecture(to, closure.ν) κ = on_architecture(to, closure.κ) - return ScalarBiharmonicDiffusivity{F, N}(ν, κ) + return ScalarBiharmonicDiffusivity{F, N, VI}(ν, κ) end \ No newline at end of file