Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix Biharmonic Viscosity on Cubed-Sphere Grid #3850

Draft
wants to merge 13 commits into
base: main
Choose a base branch
from
Draft

Fix Biharmonic Viscosity on Cubed-Sphere Grid #3850

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
0c09c1f
allow vector invariant form
simone-silvestri Oct 18, 2024
1a72658
Rename Laplacian operators and correct ∇²v formula
siddharthabishnu Oct 18, 2024
116aed2
Define vector-invariant ζ★ᶠᶠᶜ
siddharthabishnu Oct 18, 2024
8a1f70e
Define VectorInvariantForm struct and fix bugs
siddharthabishnu Oct 18, 2024
5cf8017
Remove spurious parameters
siddharthabishnu Oct 19, 2024
dedd0da
Merge branch 'main' into ss-sb-jmc/biharmonic-vector-invariant
navidcy Oct 19, 2024
12117f4
Define Δy_∇²v and Δx_∇²u, and fix bugs
siddharthabishnu Oct 22, 2024
eac550c
Union{AHBD, VectorInvariantASBD} --> AHBD
siddharthabishnu Oct 22, 2024
8138819
Fix adapt_structure and on_architecture for νhb
siddharthabishnu Oct 23, 2024
31cb273
Fix bug
siddharthabishnu Oct 23, 2024
fe6701a
Merge branch 'main' into ss-sb-jmc/biharmonic-vector-invariant
navidcy Oct 25, 2024
1c04419
Update src/TurbulenceClosures/turbulence_closure_implementations/scal…
siddharthabishnu Oct 30, 2024
cc0773d
Update src/TurbulenceClosures/turbulence_closure_implementations/scal…
siddharthabishnu Oct 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 31 additions & 18 deletions src/TurbulenceClosures/abstract_scalar_biharmonic_diffusivity_closure.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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)

Expand All @@ -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)
Comment on lines +84 to +85
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Are ∇²v and ∇²u still appropriate names for these operators? It's a bit confusing.


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

#####
Expand Down
7 changes: 7 additions & 0 deletions src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
27 changes: 15 additions & 12 deletions src/TurbulenceClosures/turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.ν
Expand All @@ -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