diff --git a/src/Advection/adapt_advection_order.jl b/src/Advection/adapt_advection_order.jl index 41f3f60c9a..038653be2d 100644 --- a/src/Advection/adapt_advection_order.jl +++ b/src/Advection/adapt_advection_order.jl @@ -39,7 +39,7 @@ function adapt_advection_order(advection, grid::AbstractGrid) @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))" + @info "Using the advection scheme $(summary(new_advection.z)) in the z-direction because size(grid, 3) = $(size(grid, 3))" end return ifelse(changed_advection, new_advection, advection) diff --git a/src/Models/Models.jl b/src/Models/Models.jl index 9399a5b503..78aab8b8bf 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -188,4 +188,9 @@ default_nan_checker(::OnlyParticleTrackingModel) = nothing # applicable to both `NonhydrostaticModel` and `HydrostaticFreeSurfaceModel` include("seawater_density.jl") +# Implementation of the diagnostic for computing the dissipation rate +include("VarianceDissipationComputation/VarianceDissipationComputation.jl") + +using .VarianceDissipationComputation + end # module diff --git a/src/Models/VarianceDissipationComputation/VarianceDissipationComputation.jl b/src/Models/VarianceDissipationComputation/VarianceDissipationComputation.jl new file mode 100644 index 0000000000..7138682e47 --- /dev/null +++ b/src/Models/VarianceDissipationComputation/VarianceDissipationComputation.jl @@ -0,0 +1,118 @@ +module VarianceDissipationComputation + +export VarianceDissipation, get_dissipation_fields + +using Oceananigans.Grids: architecture +using Oceananigans.Utils +using Oceananigans.TimeSteppers +using Oceananigans.Fields +using Oceananigans.Fields: Field, VelocityFields +using Oceananigans.Operators +using Oceananigans.BoundaryConditions +using Oceananigans.TurbulenceClosures: viscosity, + diffusivity, + ScalarDiffusivity, + ScalarBiharmonicDiffusivity, + AbstractTurbulenceClosure, + HorizontalFormulation + +using Oceananigans.Advection: _advective_tracer_flux_x, + _advective_tracer_flux_y, + _advective_tracer_flux_z, + horizontal_advection_U, + horizontal_advection_V + +using Oceananigans.Operators: volume +using KernelAbstractions: @kernel, @index + +struct VarianceDissipation{P, K, A, D, S, G} + advective_production :: P + diffusive_production :: K + advective_fluxes :: A + diffusive_fluxes :: D + previous_state :: S + gradient_squared :: G +end + +include("dissipation_utils.jl") + +function VarianceDissipation(model; + tracers = propertynames(model.tracers), + include_vorticity = true) + + if !(model.timestepper isa QuasiAdamsBashforth2TimeStepper) + throw(ArgumentError("DissipationComputation requires a QuasiAdamsBashforth2TimeStepper")) + end + + tracers = tupleit(tracers) + diffusivities = model.diffusivity_fields + closure = model.closure + grid = model.grid + + P = NamedTuple{tracers}(tracer_fluxes(grid) for tracer in tracers) + + K = NamedTuple{tracers}(tracer_closure_dissipation(grid, diffusivities, closure, id) for id in eachindex(tracers)) + Vⁿ = NamedTuple{tracers}(tracer_closure_dissipation(grid, diffusivities, closure, id) for id in eachindex(tracers)) + Vⁿ⁻¹ = NamedTuple{tracers}(tracer_closure_dissipation(grid, diffusivities, closure, id) for id in eachindex(tracers)) + + K = NamedTuple{tracers}(tracer_fluxes(grid) for tracer in tracers) + Fⁿ = NamedTuple{tracers}(tracer_fluxes(grid) for tracer in tracers) + Fⁿ⁻¹ = NamedTuple{tracers}(tracer_fluxes(grid) for tracer in tracers) + + Uⁿ⁻¹ = VelocityFields(grid) + Uⁿ = VelocityFields(grid) + + cⁿ⁻¹ = NamedTuple{tracers}(CenterField(grid) for tracer in tracers) + + if include_vorticity + Fζⁿ = vorticity_fluxes(grid) + Fζⁿ⁻¹ = vorticity_fluxes(grid) + Pζ = vorticity_fluxes(grid) + ζⁿ⁻¹ = Field{Face, Face, Center}(grid) + + P = merge(P, (; ζ = Pζ)) + Fⁿ = merge(Fⁿ, (; ζ = Fζⁿ)) + Fⁿ⁻¹ = merge(Fⁿ⁻¹, (; ζ = Fζⁿ⁻¹)) + cⁿ⁻¹ = merge(cⁿ⁻¹, (; ζ = ζⁿ⁻¹)) + + Kζ = enstrophy_closure_dissipation(grid, diffusivities, closure) + Vζⁿ = enstrophy_closure_dissipation(grid, diffusivities, closure) + Vζⁿ⁻¹ = enstrophy_closure_dissipation(grid, diffusivities, closure) + + K = merge(K, (; ζ = Kζ)) + Vⁿ = merge(Vⁿ, (; ζ = Vζⁿ)) + Vⁿ⁻¹ = merge(Vⁿ⁻¹, (; ζ = Vζⁿ⁻¹)) + end + + previous_state = merge(cⁿ⁻¹, (; Uⁿ⁻¹, Uⁿ)) + advective_fluxes = (; Fⁿ, Fⁿ⁻¹) + diffusive_fluxes = (; Vⁿ, Vⁿ⁻¹) + + gradients = deepcopy(P) + + return VarianceDissipation(P, K, advective_fluxes, diffusive_fluxes, previous_state, gradients) +end + +# Function to call in a callback +# Note: This works only if the callback is called with an IterationInterval(1), if not the +# previous fluxes and velocities will not be correct +# TODO: make sure that the correct velocities and fluxes are used even if +# the callback is not called with an IterationInterval(1) +function (ϵ::VarianceDissipation)(simulation) + + # We first assemble values for Pⁿ⁻¹ + assemble_dissipation!(simulation, ϵ) + + # Then we update the fluxes to be used in the next time step + update_fluxes!(simulation, ϵ) + + return nothing +end + +include("get_dissipation_fields.jl") +include("update_fluxes.jl") +include("advective_fluxes.jl") +include("diffusive_fluxes.jl") +include("assemble_dissipation.jl") + +end diff --git a/src/Models/VarianceDissipationComputation/advective_dissipation.jl b/src/Models/VarianceDissipationComputation/advective_dissipation.jl new file mode 100644 index 0000000000..d28fd941b2 --- /dev/null +++ b/src/Models/VarianceDissipationComputation/advective_dissipation.jl @@ -0,0 +1,61 @@ +# TODO: This is only for AB2, figure out how to generalize this for other timesteppers for example RK3 +@kernel function _assemble_advective_tracer_dissipation!(P, grid, χ, Fⁿ, Fⁿ⁻¹, Uⁿ⁺¹, Uⁿ, Uⁿ⁻¹, cⁿ⁺¹, cⁿ) + i, j, k = @index(Global, NTuple) + + δˣc★ = δxᶠᶜᶜ(i, j, k, grid, c★, cⁿ⁺¹, cⁿ) + δˣc² = δxᶠᶜᶜ(i, j, k, grid, c², cⁿ⁺¹, cⁿ) + + δʸc★ = δyᶜᶠᶜ(i, j, k, grid, c★, cⁿ⁺¹, cⁿ) + δʸc² = δyᶜᶠᶜ(i, j, k, grid, c², cⁿ⁺¹, cⁿ) + + δᶻc★ = δzᶜᶜᶠ(i, j, k, grid, c★, cⁿ⁺¹, cⁿ) + δᶻc² = δzᶜᶜᶠ(i, j, k, grid, c², cⁿ⁺¹, cⁿ) + + C₁ = convert(eltype(grid), 1.5 + χ) + C₂ = convert(eltype(grid), 0.5 + χ) + + @inbounds begin + u₁ = C₁ * Uⁿ.u[i, j, k] + u₂ = C₂ * Uⁿ⁻¹.u[i, j, k] + v₁ = C₁ * Uⁿ.v[i, j, k] + v₂ = C₂ * Uⁿ⁻¹.v[i, j, k] + w₁ = C₁ * Uⁿ.w[i, j, k] + w₂ = C₂ * Uⁿ⁻¹.w[i, j, k] + + fx₁ = C₁ * Fⁿ.x[i, j, k] + fx₂ = C₂ * Fⁿ⁻¹.x[i, j, k] + fy₁ = C₁ * Fⁿ.y[i, j, k] + fy₂ = C₂ * Fⁿ⁻¹.y[i, j, k] + fz₁ = C₁ * Fⁿ.z[i, j, k] + fz₂ = C₂ * Fⁿ⁻¹.z[i, j, k] + + P.x[i, j, k] = 2 * δˣc★ * (fx₁ - fx₂) - δˣc² * (u₁ - u₂) + P.y[i, j, k] = 2 * δʸc★ * (fy₁ - fy₂) - δʸc² * (v₁ - v₂) + P.z[i, j, k] = 2 * δᶻc★ * (fz₁ - fz₂) - δᶻc² * (w₁ - w₂) + end +end + +@kernel function _assemble_advective_vorticity_dissipation!(P, grid, χ, Fⁿ, Fⁿ⁻¹, Uⁿ⁺¹, Uⁿ, Uⁿ⁻¹, c, ζⁿ) + i, j, k = @index(Global, NTuple) + + δˣζ★ = δxᶠᶜᶜ(i, j, k, grid, ζ★, Uⁿ⁺¹.u, Uⁿ⁺¹.v, ζⁿ) + δˣζ² = δxᶠᶜᶜ(i, j, k, grid, ζ², Uⁿ⁺¹.u, Uⁿ⁺¹.v, ζⁿ) + + δʸζ★ = δyᶜᶠᶜ(i, j, k, grid, ζ★, Uⁿ⁺¹.u, Uⁿ⁺¹.v, ζⁿ) + δʸζ² = δyᶜᶠᶜ(i, j, k, grid, ζ², Uⁿ⁺¹.u, Uⁿ⁺¹.v, ζⁿ) + + @inbounds begin + u₁ = C₁ * ℑxyᶜᶠᵃ(i, j, k, grid, Δy_qᶠᶜᶜ, Uⁿ.u) / Δyᶠᶜᶜ(i, j, k, grid) + u₂ = C₂ * ℑxyᶜᶠᵃ(i, j, k, grid, Δy_qᶠᶜᶜ, Uⁿ⁻¹.u) / Δyᶠᶜᶜ(i, j, k, grid) + v₁ = C₁ * ℑxyᶠᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶜ, Uⁿ.v) / Δxᶜᶠᶜ(i, j, k, grid) + v₂ = C₂ * ℑxyᶠᶜᵃ(i, j, k, grid, Δx_qᶜᶠᶜ, Uⁿ⁻¹.v) / Δxᶜᶠᶜ(i, j, k, grid) + + fx₁ = C₁ * Fⁿ.x[i, j, k] + fx₂ = C₂ * Fⁿ⁻¹.x[i, j, k] + fy₁ = C₁ * Fⁿ.y[i, j, k] + fy₂ = C₂ * Fⁿ⁻¹.y[i, j, k] + + P.x[i, j, k] = 2 * δˣζ★ * (fx₁ - fx₂) - δˣζ² * (u₁ - u₂) + P.y[i, j, k] = 2 * δʸζ★ * (fy₁ - fy₂) - δʸζ² * (v₁ - v₂) + end +end diff --git a/src/Models/VarianceDissipationComputation/advective_fluxes.jl b/src/Models/VarianceDissipationComputation/advective_fluxes.jl new file mode 100644 index 0000000000..be11a5444b --- /dev/null +++ b/src/Models/VarianceDissipationComputation/advective_fluxes.jl @@ -0,0 +1,42 @@ +@kernel function _update_advective_tracer_fluxes!(Gⁿ, Fⁿ, Fⁿ⁻¹, cⁿ⁻¹, grid, advection, U, c) + i, j, k = @index(Global, NTuple) + u, v, w = U + + @inbounds begin + # Save previous advective fluxes + Fⁿ⁻¹.x[i, j, k] = Fⁿ.x[i, j, k] + Fⁿ⁻¹.y[i, j, k] = Fⁿ.y[i, j, k] + Fⁿ⁻¹.z[i, j, k] = Fⁿ.z[i, j, k] + + cⁿ⁻¹[i, j, k] = c[i, j, k] + + # Calculate new advective fluxes + Fⁿ.x[i, j, k] = _advective_tracer_flux_x(i, j, k, grid, advection, u, c) + Fⁿ.y[i, j, k] = _advective_tracer_flux_y(i, j, k, grid, advection, v, c) + Fⁿ.z[i, j, k] = _advective_tracer_flux_z(i, j, k, grid, advection, w, c) + + Gⁿ.x[i, j, k] = Axᶠᶜᶜ(i, j, k, grid) * δxᶠᶜᶜ(i, j, k, grid, c)^2 / Δxᶠᶜᶜ(i, j, k, grid) + Gⁿ.y[i, j, k] = Ayᶜᶠᶜ(i, j, k, grid) * δyᶜᶠᶜ(i, j, k, grid, c)^2 / Δyᶜᶠᶜ(i, j, k, grid) + Gⁿ.z[i, j, k] = Azᶜᶜᶠ(i, j, k, grid) * δzᶜᶜᶠ(i, j, k, grid, c)^2 / Δzᶜᶜᶠ(i, j, k, grid) + end +end + +@kernel function _update_advective_vorticity_fluxes!(Gⁿ, Fⁿ, Fⁿ⁻¹, ζⁿ⁻¹, grid, advection, U, c) + i, j, k = @index(Global, NTuple) + u, v, w = U + + @inbounds begin + # Save previous advective fluxes + Fⁿ⁻¹.x[i, j, k] = Fⁿ.x[i, j, k] + Fⁿ⁻¹.y[i, j, k] = Fⁿ.y[i, j, k] + + ζⁿ⁻¹[i, j, k] = ζ₃ᶠᶠᶜ(i, j, k, grid, U.u, U.v) + + # Calculate new advective fluxes + Fⁿ.x[i, j, k] = horizontal_advection_V(i, j, k, grid, advection, u, ζ) * Axᶜᶠᶜ(i, j, k, grid) + Fⁿ.y[i, j, k] = - horizontal_advection_U(i, j, k, grid, advection, v, ζ) * Ayᶠᶜᶜ(i, j, k, grid) + + Gⁿ.x[i, j, k] = Axᶜᶠᶜ(i, j, k, grid) * δxᶜᶠᶜ(i, j, k, grid, ζ₃ᶠᶠᶜ, U.u)^2 / Δxᶜᶠᶜ(i, j, k, grid) + Gⁿ.y[i, j, k] = Ayᶠᶜᶜ(i, j, k, grid) * δyᶠᶜᶜ(i, j, k, grid, ζ₃ᶠᶠᶜ, U.v)^2 / Δyᶠᶜᶜ(i, j, k, grid) + end +end \ No newline at end of file diff --git a/src/Models/VarianceDissipationComputation/assemble_dissipation.jl b/src/Models/VarianceDissipationComputation/assemble_dissipation.jl new file mode 100644 index 0000000000..7e6ec76293 --- /dev/null +++ b/src/Models/VarianceDissipationComputation/assemble_dissipation.jl @@ -0,0 +1,61 @@ +function assemble_dissipation!(simulation, dissipation) + model = simulation.model + + for tracer_name in keys(dissipation.advective_production) + assemble_dissipation!(dissipation, model, tracer_name) + end + + return nothing +end + +@inline c★(i, j, k, grid, cⁿ⁺¹, cⁿ) = @inbounds (cⁿ⁺¹[i, j, k] + cⁿ[i, j, k]) / 2 +@inline c²(i, j, k, grid, cⁿ⁺¹, cⁿ) = @inbounds (cⁿ⁺¹[i, j, k] * cⁿ[i, j, k]) + +@inline ζ★(i, j, k, grid, uⁿ⁺¹, vⁿ⁺¹, ζⁿ) = @inbounds (ζ₃ᶠᶠᶜ(i, j, k, grid, uⁿ⁺¹, vⁿ⁺¹) + ζⁿ[i, j, k]) / 2 +@inline ζ²(i, j, k, grid, uⁿ⁺¹, vⁿ⁺¹, ζⁿ) = @inbounds (ζ₃ᶠᶠᶜ(i, j, k, grid, uⁿ⁺¹, vⁿ⁺¹) * ζⁿ[i, j, k]) + +function assemble_dissipation!(dissipation, model, tracer_name::Symbol) + + + arch = architecture(grid) + χ = simulation.model.timestepper.χ + + # General velocities + Uⁿ⁺¹ = model.velocities + Uⁿ = dissipation.previous_state.Uⁿ + Uⁿ⁻¹ = dissipation.previous_state.Uⁿ⁻¹ + + cⁿ⁺¹ = tracer_symbol == :ζ ? nothing : model.tracers[tracer_name] + cⁿ = dissipation.previous_state[tracer_name] + + + _assemble_advective_dissipation! = assemble_advective_dissipation_kernel(Val(tracer_name)) + _assemble_diffusive_dissipation! = assemble_diffusive_dissipation_kernel(Val(tracer_name)) + + #### + #### Assemble the advective dissipation + #### + + P = dissipation.advective_production[tracer_name] + Fⁿ = dissipation.advective_fluxes.Fⁿ[tracer_name] + Fⁿ⁻¹ = dissipation.advective_fluxes.Fⁿ⁻¹[tracer_name] + + launch!(arch, grid, :xyz, _assemble_advective_dissipation!, P, grid, χ, Fⁿ, Fⁿ⁻¹, Uⁿ⁺¹, Uⁿ, Uⁿ⁻¹, cⁿ⁺¹, cⁿ) + + #### + #### Assemble the diffusive dissipation + #### + + K = dissipation.diffusive_production[tracer_name] + Vⁿ = dissipation.advective_fluxes.Vⁿ[tracer_name] + Vⁿ⁻¹ = dissipation.advective_fluxes.Vⁿ⁻¹[tracer_name] + + launch!(arch, grid, params, _assemble_diffusive_dissipation!, K, grid, χ, Vⁿ, Vⁿ⁻¹, Uⁿ⁺¹, cⁿ⁺¹, cⁿ) + + return nothing +end + +assemble_advective_dissipation_kernel(val_tracer_name) = _assemble_advective_tracer_dissipation! +assemble_advective_dissipation_kernel(::Val{:ζ}) = _assemble_advective_vorticity_dissipation! +assemble_diffusive_dissipation_kernel(val_tracer_name) = _assemble_diffusive_tracer_dissipation! +assemble_diffusive_dissipation_kernel(::Val{:ζ}) = _assemble_diffusive_vorticity_dissipation! diff --git a/src/Models/VarianceDissipationComputation/diffusive_dissipation.jl b/src/Models/VarianceDissipationComputation/diffusive_dissipation.jl new file mode 100644 index 0000000000..5f4215908e --- /dev/null +++ b/src/Models/VarianceDissipationComputation/diffusive_dissipation.jl @@ -0,0 +1,49 @@ +@kernel function _assemble_diffusive_tracer_dissipation!(K, grid, χ, Vⁿ, Vⁿ⁻¹, Uⁿ⁺¹, cⁿ⁺¹, cⁿ) + i, j, k = @index(Global, NTuple) + compute_diffusive_dissipation!(K, i, j, k, grid, Vⁿ, Vⁿ⁻¹, χ, cⁿ⁺¹, cⁿ) +end + +@inline function compute_diffusive_tracer_dissipation!(K::Tuple, i, j, k, grid, Vⁿ, Vⁿ⁻¹, χ, cⁿ⁺¹, cⁿ) + for n in eachindex(K) + compute_diffusive_dissipation!(K[n], i, j, k, grid, Vⁿ[n], Vⁿ⁻¹[n], χ, cⁿ⁺¹, cⁿ) + end +end + +@inline function compute_diffusive_tracer_dissipation!(K, i, j, k, grid, Vⁿ, Vⁿ⁻¹, χ, cⁿ⁺¹, cⁿ) + C₁ = convert(eltype(grid), 1.5 + χ) + C₂ = convert(eltype(grid), 0.5 + χ) + + δˣc★ = δxᶠᶜᶜ(i, j, k, grid, c★, cⁿ⁺¹, cⁿ) + δʸc★ = δyᶜᶠᶜ(i, j, k, grid, c★, cⁿ⁺¹, cⁿ) + δᶻc★ = δzᶜᶜᶠ(i, j, k, grid, c★, cⁿ⁺¹, cⁿ) + + @inbounds begin + K.x[i, j, k] = 2 * δˣc★ * (C₁ * Vⁿ.x[i, j, k] - C₂ * Vⁿ⁻¹.x[i, j, k]) + K.y[i, j, k] = 2 * δʸc★ * (C₁ * Vⁿ.y[i, j, k] - C₂ * Vⁿ⁻¹.y[i, j, k]) + K.z[i, j, k] = 2 * δᶻc★ * (C₁ * Vⁿ.z[i, j, k] - C₂ * Vⁿ⁻¹.z[i, j, k]) + end +end + +@kernel function _assemble_diffusive_vorticity_dissipation!(K, grid, χ, Vⁿ, Vⁿ⁻¹, Uⁿ⁺¹, cⁿ⁺¹, ζⁿ) + i, j, k = @index(Global, NTuple) + compute_diffusive_vorticity_dissipation!(K, i, j, k, grid, Vⁿ, Vⁿ⁻¹, χ, Uⁿ⁺¹, ζⁿ) +end + +@inline function compute_diffusive_vorticity_dissipation!(K::Tuple, i, j, k, grid, Vⁿ, Vⁿ⁻¹, χ, Uⁿ⁺¹, ζⁿ) + for n in eachindex(K) + compute_diffusive_vorticity_dissipation!(K[n], i, j, k, grid, Vⁿ[n], Vⁿ⁻¹[n], χ, Uⁿ⁺¹, ζⁿ) + end +end + +@inline function compute_diffusive_vorticity_dissipation!(K, i, j, k, grid, Vⁿ, Vⁿ⁻¹, χ, Uⁿ⁺¹, ζⁿ) + C₁ = convert(eltype(grid), 1.5 + χ) + C₂ = convert(eltype(grid), 0.5 + χ) + + δˣζ★ = δxᶠᶜᶜ(i, j, k, grid, ζ★, Uⁿ⁺¹.u, Uⁿ⁺¹.v, ζⁿ) + δʸζ★ = δyᶜᶠᶜ(i, j, k, grid, ζ★, Uⁿ⁺¹.u, Uⁿ⁺¹.v, ζⁿ) + + @inbounds begin + K.x[i, j, k] = 2 * δˣζ★ * (C₁ * Vⁿ.x[i, j, k] - C₂ * Vⁿ⁻¹.x[i, j, k]) + K.y[i, j, k] = 2 * δʸζ★ * (C₁ * Vⁿ.y[i, j, k] - C₂ * Vⁿ⁻¹.y[i, j, k]) + end +end \ No newline at end of file diff --git a/src/Models/VarianceDissipationComputation/diffusive_fluxes.jl b/src/Models/VarianceDissipationComputation/diffusive_fluxes.jl new file mode 100644 index 0000000000..ff43c56b86 --- /dev/null +++ b/src/Models/VarianceDissipationComputation/diffusive_fluxes.jl @@ -0,0 +1,40 @@ + +@kenrel function _update_diffusive_tracer_fluxes!(Vⁿ, Vⁿ⁻¹, grid, closure, diffusivity, bouyancy, c, tracer_id, clk, model_fields) + i, j, k = @index(Global, NTuple) + compute_diffusive_tracer_fluxes!(Vⁿ, Vⁿ⁻¹, i, j, k, grid, closure, diffusivity, bouyancy, c, tracer_id, clk, model_fields) +end + +@inline function compute_diffusive_tracer_fluxes!(Vⁿ, Vⁿ⁻¹, i, j, k, grid, closure::Tuple, K, args...) + for n in eachindex(closure) + compute_diffusive_tracer_fluxes!(Vⁿ[n], Vⁿ⁻¹[n], i, j, k, grid, closure[n], K[n], args...) + end +end + +@inline function compute_diffusive_tracer_fluxes!(Vⁿ, Vⁿ⁻¹, i, j, k, grid, clo, K, b, c, c_id, clk, fields) + Vⁿ⁻¹.x[i, j, k] = Vⁿ.x[i, j, k] + Vⁿ⁻¹.y[i, j, k] = Vⁿ.y[i, j, k] + Vⁿ⁻¹.z[i, j, k] = Vⁿ.z[i, j, k] + + Vⁿ.x[i, j, k] = _diffusive_tracer_flux_x(i, j, k, grid, clo, K, Val(c_id), c, clk, fields, b) * Axᶠᶜᶜ(i, j, k, grid) + Vⁿ.y[i, j, k] = _diffusive_tracer_flux_y(i, j, k, grid, clo, K, Val(c_id), c, clk, fields, b) * Ayᶜᶠᶜ(i, j, k, grid) + Vⁿ.z[i, j, k] = _diffusive_tracer_flux_z(i, j, k, grid, clo, K, Val(c_id), c, clk, fields, b) * Azᶜᶜᶠ(i, j, k, grid) +end + +@kenrel function _update_diffusive_vorticity_fluxes!(Vⁿ, Vⁿ⁻¹, grid, closure, diffusivity, bouyancy, c, tracer_id, clk, model_fields) + i, j, k = @index(Global, NTuple) + compute_diffusive_vorticity_fluxes!(Vⁿ, Vⁿ⁻¹, i, j, k, grid, closure, diffusivity, bouyancy, clk, model_fields) +end + +@inline function compute_diffusive_vorticity_fluxes!(Vⁿ, Vⁿ⁻¹, i, j, k, grid, closure::Tuple, K, args...) + for n in eachindex(closure) + compute_diffusive_vorticity_fluxes!(Vⁿ[n], Vⁿ⁻¹[n], i, j, k, grid, closure[n], K[n], args...) + end +end + +@inline function compute_diffusive_vorticity_fluxes!(Vⁿ, Vⁿ⁻¹, i, j, k, grid, clo, K, b, clk, fields) + Vⁿ⁻¹.x[i, j, k] = Vⁿ.x[i, j, k] + Vⁿ⁻¹.y[i, j, k] = Vⁿ.y[i, j, k] + + Vⁿ.x[i, j, k] = ∂ⱼ_τ₂ⱼ(i, j, k, grid, clo, K, clk, fields, b) * Axᶜᶠᶜ(i, j, k, grid) + Vⁿ.y[i, j, k] = - ∂ⱼ_τ₁ⱼ(i, j, k, grid, clo, K, clk, fields, b) * Ayᶠᶜᶜ(i, j, k, grid) +end diff --git a/src/Models/VarianceDissipationComputation/dissipation_utils.jl b/src/Models/VarianceDissipationComputation/dissipation_utils.jl new file mode 100644 index 0000000000..a6f060c58c --- /dev/null +++ b/src/Models/VarianceDissipationComputation/dissipation_utils.jl @@ -0,0 +1,48 @@ +import Oceananigans.Utils: KernelParameters + +tracer_closure_dissipation(grid, K, ::Nothing, tracer_id) = nothing +tracer_closure_dissipation(grid, K, c::Tuple, tracer_id) = + Tuple(tracer_closure_dissipation(grid, K[i], c[i], tracer_id) for i in eachindex(c)) + +function tracer_closure_dissipation(K, c, tracer_id) + κ = diffusivity(c, K, tracer_id) + include_dissipation = !(κ isa Number) || (κ != 0) + return ifelse(include_dissipation, tracer_fluxes(grid), nothing) +end + +enstrophy_closure_dissipation(grid, K, c::Tuple) = + Tuple(enstrophy_closure_dissipation(grid, K[i], c[i]) for i in eachindex(c)) + +# Fallback +enstrophy_closure_dissipation(grid, K, ::Nothing) = nothing + +function enstrophy_closure_dissipation(K, c) + ν = viscosity(c, K) + include_dissipation = !(ν isa Number) || (ν != 0) + return ifelse(include_dissipation, vorticity_fluxes(grid), nothing) +end + +@inline getadvection(advection, tracer_name) = advection +@inline getadvection(advection::NamedTuple, tracer_name) = + @inbounds ifelse(tracer_name == :ζ, advection.momentum, advection[tracer_name]) + +@inline function KernelParameters(f::Field) + sz = size(f.data) + of = f.data.offsets + return KernelParameters(sz, of) +end + +function tracer_fluxes(grid) + x = XFaceField(grid) + y = YFaceField(grid) + z = ZFaceField(grid) + + return (; x, y, z) +end + +function vorticity_fluxes(grid) + x = YFaceField(grid) + y = XFaceField(grid) + + return (; x, y) +end diff --git a/src/Models/VarianceDissipationComputation/get_dissipation_fields.jl b/src/Models/VarianceDissipationComputation/get_dissipation_fields.jl new file mode 100644 index 0000000000..c1c92c0708 --- /dev/null +++ b/src/Models/VarianceDissipationComputation/get_dissipation_fields.jl @@ -0,0 +1,27 @@ +function get_dissipation_fields(t::VarianceDissipation) + f = NamedTuple() + + for name in keys(t.advective_production) + f = merge(f, get_dissipation_fields(t, name)) + end + + return f +end + +function get_dissipation_fields(t::VarianceDissipation, tracer_name) + A = t.advective_production[tracer_name] + D = t.diffusive_production[tracer_name] + G = t.gradient_squared[tracer_name] + + dirs = tracer_name == :ζ ? (:x, :y) : (:x, :y, :z) + + prod_names = Tuple(Symbol(:A, tracer_name, dir) for dir in dirs) + diff_names = Tuple(Symbol(:D, tracer_name, dir) for dir in dirs) + grad_names = Tuple(Symbol(:G, tracer_name, dir) for dir in dirs) + + advective_prod = Tuple(getproperty(A, dir) for dir in dirs) + diffusive_prod = Tuple(getproperty(D, dir) for dir in dirs) + grad = Tuple(getproperty(G, dir) for dir in dirs) + + return NamedTuple{tuple(prod_names..., diff_names..., grad_names...)}(tuple(advective_prod..., diffusive_prod..., grad...)) +end \ No newline at end of file diff --git a/src/Models/VarianceDissipationComputation/update_fluxes.jl b/src/Models/VarianceDissipationComputation/update_fluxes.jl new file mode 100644 index 0000000000..09fe3cc00c --- /dev/null +++ b/src/Models/VarianceDissipationComputation/update_fluxes.jl @@ -0,0 +1,84 @@ +using Oceananigans: fields + +function update_fluxes!(simulation, dissipation) + model = simulation.model + + grid = model.grid + arch = architecture(grid) + params = KernelParameters(model.tracers[1]) + + Uⁿ = dissipation.previous_state.Uⁿ + Uⁿ⁻¹ = dissipation.previous_state.Uⁿ⁻¹ + U = model.velocities + + launch!(arch, grid, params, _update_transport!, Uⁿ, Uⁿ⁻¹, grid, U) + + for (tracer_id, tracer_name) in enumerate(keys(dissipation.advective_production)) + update_fluxes!(dissipation, model, tracer_name, tracer_id) + end + + return nothing +end + +function update_fluxes!(dissipation, model, tracer_name::Symbol, tracer_id) + + # Grab tracer properties + c = tracer_name == :ζ ? nothing : model.tracers[tracer_name] + cⁿ⁻¹ = dissipation.previous_state[tracer_name] + + grid = model.grid + arch = architecture(grid) + + U = model.velocities + params = KernelParameters(model.tracers[1]) + + _update_advective_fluxes! = update_advective_fluxes_kernel(Val(tracer_name)) + _update_diffusive_fluxes! = update_diffusive_fluxes_kernel(Val(tracer_name)) + + #### + #### Update the advective fluxes and compute gradient squared + #### + + Fⁿ = dissipation.advective_fluxes.Fⁿ[tracer_name] + Fⁿ⁻¹ = dissipation.advective_fluxes.Fⁿ⁻¹[tracer_name] + Gⁿ = dissipation.gradient_squared[tracer_name] + advection = getadvection(model.advection, tracer_name) + + launch!(arch, grid, params, _update_advective_fluxes!, Gⁿ, Fⁿ, Fⁿ⁻¹, cⁿ⁻¹, grid, advection, U, c) + + #### + #### Update the diffusive fluxes + #### + + Vⁿ = dissipation.diffusive_fluxes.Vⁿ[tracer_name] + Vⁿ⁻¹ = dissipation.diffusive_fluxes.Vⁿ⁻¹[tracer_name] + + D = model.diffusivity_fields + B = model.buoyancy + clk = model.clock + clo = model.closure + model_fields = fields(model) + + launch!(arch, grid, params, _update_diffusive_fluxes!, Vⁿ, Vⁿ⁻¹, grid, clo, D, B, c, tracer_id, clk, model_fields) + + return nothing +end + +@kernel function _update_transport!(Uⁿ, Uⁿ⁻¹, grid, U) + i, j, k = @index(Global, NTuple) + + @inbounds begin + Uⁿ⁻¹.u[i, j, k] = Uⁿ.u[i, j, k] + Uⁿ⁻¹.v[i, j, k] = Uⁿ.v[i, j, k] + Uⁿ⁻¹.w[i, j, k] = Uⁿ.w[i, j, k] + Uⁿ.u[i, j, k] = U.u[i, j, k] * Axᶠᶜᶜ(i, j, k, grid) + Uⁿ.v[i, j, k] = U.v[i, j, k] * Ayᶜᶠᶜ(i, j, k, grid) + Uⁿ.w[i, j, k] = U.w[i, j, k] * Azᶜᶜᶠ(i, j, k, grid) + end +end + +update_advective_fluxes_kernel(val_tracer_name) = _update_advective_tracer_fluxes! +update_advective_fluxes_kernel(::Val{:ζ}) = _update_advective_vorticity_fluxes! +update_diffusive_fluxes_kernel(val_tracer_name) = _update_diffusive_tracer_fluxes! +update_diffusive_fluxes_kernel(::Val{:ζ}) = _update_diffusive_vorticity_fluxes! + diff --git a/src/Oceananigans.jl b/src/Oceananigans.jl index db93e518ef..4e617e6abf 100644 --- a/src/Oceananigans.jl +++ b/src/Oceananigans.jl @@ -174,6 +174,7 @@ abstract type AbstractOutputWriter end # Callsites for Callbacks struct TimeStepCallsite end struct TendencyCallsite end +struct BeforeTimeStepCallsite end struct UpdateStateCallsite end ##### diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl index ed9f55750b..16fc71b847 100644 --- a/src/TimeSteppers/quasi_adams_bashforth_2.jl +++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl @@ -1,5 +1,6 @@ using Oceananigans.Fields: FunctionField, location using Oceananigans.Utils: @apply_regionally, apply_regionally! +using Oceananigans: BeforeTimeStepCallsite mutable struct QuasiAdamsBashforth2TimeStepper{FT, GT, IT} <: AbstractTimeStepper χ :: FT @@ -78,6 +79,8 @@ function time_step!(model::AbstractModel{<:QuasiAdamsBashforth2TimeStepper}, Δt Δt == 0 && @warn "Δt == 0 may cause model blowup!" + [callback(model) for callback in callbacks if isa(callback.callsite, BeforeTimeStepCallsite)] + # Be paranoid and update state at iteration 0 model.clock.iteration == 0 && update_state!(model, callbacks) diff --git a/validation/implicit_dissipation/implicit_dissipation.jl b/validation/implicit_dissipation/implicit_dissipation.jl new file mode 100644 index 0000000000..6bb7135208 --- /dev/null +++ b/validation/implicit_dissipation/implicit_dissipation.jl @@ -0,0 +1,150 @@ +using Oceananigans +using Oceananigans.Fields: VelocityFields +using Oceananigans.Models.HydrostaticFreeSurfaceModels: PrescribedVelocityFields +using Oceananigans.Models: TracerVarianceDissipation +using Printf +using GLMakie + +# the initial condition +@inline G(x, β, z) = exp(-β*(x - z)^2) +@inline F(x, α, a) = √(max(1 - α^2*(x-a)^2, 0.0)) + +const Z = -0.7 +const δ = 0.005 +const β = log(2)/(36*δ^2) +const a = 0.5 +const α = 10 + +@inline function bᵢ(x) + if x <= -0.6 && x >= -0.8 + return 1/6*(G(x, β, Z-δ) + 4*G(x, β, Z) + G(x, β, Z+δ)) + elseif x <= -0.2 && x >= -0.4 + return 1.0 + elseif x <= 0.2 && x >= 0.0 + return 1.0 - abs(10 * (x - 0.1)) + elseif x <= 0.6 && x >= 0.4 + return 1/6*(F(x, α, a-δ) + 4*F(x, α, a) + F(x, α, a+δ)) + else + return 0.0 + end +end + +function save_previous_buoyancy!(simulation) + bⁿ⁻¹ = simulation.model.auxiliary_fields.bⁿ⁻¹ + bⁿ = simulation.model.tracers.b + + parent(bⁿ⁻¹) .= parent(bⁿ) + + return nothing +end + +function one_dimensional_simulation(grid, advection, label) + + auxiliary_fields = (; bⁿ⁻¹ = CenterField(grid)) + + model = NonhydrostaticModel(; grid, auxiliary_fields, tracers = :b, timestepper = :QuasiAdamsBashforth2, advection) + + set!(model.tracers.b, bᵢ) + set!(model.velocities.u, 1.0) + + Δt = 0.1 * minimum_xspacing(grid) + + dissipation_computation = TracerVarianceDissipation(model; tracers = :b) + simulation = Simulation(model; Δt, stop_time = 10) + simulation.callbacks[:compute_dissipation] = Callback(dissipation_computation, IterationInterval(1)) + simulation.callbacks[:save_previous_buoyancy] = Callback(save_previous_buoyancy!, IterationInterval(1); callsite = Oceananigans.BeforeTimeStepCallsite()) + + Px = dissipation_computation.production.b.x + bⁿ⁻¹ = simulation.model.auxiliary_fields.bⁿ⁻¹ + bⁿ = model.tracers.b + + Σb² = (bⁿ^2 - bⁿ⁻¹^2) / Δt # Not at the same time-step of Px unfortunately + outputs = (; Px, b = model.tracers.b, Σb²) + filename = "dissipation_" * label * ".jld2" + + # Save the dissipation + simulation.output_writers[:dissipation] = JLD2OutputWriter(model, outputs; + filename, + schedule = IterationInterval(10), + overwrite_existing = true) + + run!(simulation) +end + +grid = RectilinearGrid(size = 100, halo = 6, x = (-1, 1), topology = (Periodic, Flat, Flat)) + +advections = [UpwindBiased(; order = 3), + WENO(; order = 5), + WENO(; order = 7)] + +labels = ["upwind3", + "WENO5", + "WENO7"] + +# Run the simulations +for (advection, label) in zip(advections, labels) + one_dimensional_simulation(grid, advection, label) +end + +B_series = [] +b_series = [] +P_series = [] + +iter = Observable(1) + +bn = [] +Pn = [] +Bn = [] + +for (i, label) in enumerate(labels) + filename = "dissipation_" * label * ".jld2" + + push!(b_series, FieldTimeSeries(filename, "b")) + push!(P_series, FieldTimeSeries(filename, "Px")) + push!(B_series, FieldTimeSeries(filename, "Σb²")) + + push!(bn, @lift(interior(b_series[i][$iter], :, 1, 1))) + push!(Pn, @lift(interior(P_series[i][$iter], :, 1, 1))) + push!(Bn, @lift(interior(B_series[i][$iter], :, 1, 1))) +end + +b0 = b_series[1][1] +x = xnodes(b0) + +fig = Figure(size = (1200, 400)) +ax = Axis(fig[1, 1], xlabel = L"x", ylabel = L"tracer") +lines!(ax, x, interior(b0, :, 1, 1), color = :grey, linestyle = :dash, linewidth = 2, label = "initial condition") +lines!(ax, x, bn[1], label = labels[1], color = :red ) +lines!(ax, x, bn[3], label = labels[2], color = :blue) +axislegend(ax, position = :rb) + +ax = Axis(fig[1, 2], xlabel = L"x", ylabel = L"variance dissipation") +lines!(ax, x, Pn[1], color = :red , label = labels[1]) +lines!(ax, x, Pn[3], color = :blue, label = labels[2]) +lines!(ax, x, Bn[1], color = :red , linestyle = :dash) +lines!(ax, x, Bn[3], color = :blue, linestyle = :dash) +ylims!(ax, (-1, 1)) +axislegend(ax, position = :rb) + +record(fig, "implicit_dissipation.mp4", 1:length(b_series[1]), framerate=8) do i + @info "doing iter $i" + iter[] = i +end + +Nt = length(b_series[1]) +time = (0:Nt-1) .* 0.1 .* minimum_xspacing(grid) + +fig = Figure() +ax = Axis(fig[1, 1], xlabel = L"time", ylabel = L"tracer variance budget") +Σb² = [sum(B_series[1][i]) for i in 1:Nt] +ΣP = [sum(P_series[1][i]) for i in 1:Nt] +lines!(ax, time, Σb², label = "Σb²", color = :red) +lines!(ax, time, ΣP / 0.002, label = "ΣP", color = :red, linestyle = :dash) +Σb² = [sum(B_series[2][i]) for i in 1:Nt] +ΣP = [sum(P_series[2][i]) for i in 1:Nt] +lines!(ax, time, Σb², label = "Σb²", color = :blue) +lines!(ax, time, ΣP / 0.002, label = "ΣP", color = :blue, linestyle = :dash) +Σb² = [sum(B_series[3][i]) for i in 1:Nt] +ΣP = [sum(P_series[3][i]) for i in 1:Nt] +lines!(ax, time, Σb², label = "Σb²", color = :green) +lines!(ax, time, ΣP / 0.002, label = "ΣP", color = :green, linestyle = :dash) \ No newline at end of file