From 278ed6b1a5ba98bc42d178ce10badf86c5f58d14 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 28 Oct 2024 17:07:37 +0100 Subject: [PATCH 01/10] first try --- src/Models/dissipation_computation.jl | 193 ++++++++++++++++++ .../implicit_dissipation.jl | 0 2 files changed, 193 insertions(+) create mode 100644 src/Models/dissipation_computation.jl create mode 100644 validation/implicit_dissipation/implicit_dissipation.jl diff --git a/src/Models/dissipation_computation.jl b/src/Models/dissipation_computation.jl new file mode 100644 index 0000000000..fd28f52f89 --- /dev/null +++ b/src/Models/dissipation_computation.jl @@ -0,0 +1,193 @@ +using Oceananigans.Grids: architecture +using Oceananigans.Utils +using Oceananigans.Fields: Field, VelocityFields +using Oceananigans.Operators +using Oceananigans.BoundaryConditions +using Oceananigans.Advection: _advective_tracer_flux_x, _advective_tracer_flux_y, _advective_tracer_flux_z +using Oceananigans.Advection: horizontal_advection_U, horizontal_advection_V +using Oceananigans.Models.HydrostaticFreeSurfaceModels: vertical_scaling, previous_vertical_scaling + +import Oceananigans.Utils: KernelParameters + +const c = Center() +const f = Face() + +struct DissipationComputation{P, A, S} + production :: P + advective_fluxes :: A + previous_state :: S +end + +function fluxes_field(grid) + x = XFaceField(grid) + y = YFaceField(grid) + z = ZFaceField(grid) + + return (; x, y, z) +end + +function DissipationComputation(model; tracers = tracer_names(model.tracers)) + + if !(model.timestepper isa QuasiAdamsBashforth2Timestepper) + throw(ArgumentError("DissipationComputation requires a QuasiAdamsBashforth2Timestepper")) + end + + grid = model.grid + P = NamedTuple{tracers}(fluxes_field(grid) for tracer in tracers) + Fⁿ = NamedTuple{tracers}(fluxes_field(grid) for tracer in tracers) + Fⁿ⁻¹ = NamedTuple{tracers}(fluxes_field(grid) for tracer in tracers) + + Uⁿ⁻¹ = VelocityFields(grid) + Uⁿ = VelocityFields(grid) + + cⁿ⁻¹ = NamedTuple{tracers}(CenterField(grid) for tracer in tracers) + + previous_state = merge(cⁿ⁻¹, (; Uⁿ⁻¹, Uⁿ)) + advective_fluxes = (; Fⁿ, Fⁿ⁻¹) + + return DissipationComputation(P, advective_fluxes, previous_state) +end + +# Function to call in a callback +function (dc::DissipationComputation)(simulation) + # We first assemble values for Pⁿ⁻¹ + assemble_P_values!(simulation, dc) + + # Then we update the fluxes to be used in the next time step + update_fluxes!(simulation, dc) + + return nothing +end + +@inline function KernelParameters(f::Field) + sz = size(f.data) + of = f.data.offsets + return KernelParameters(sz, of) +end + +function update_fluxes!(simulation, dissipation_computation) + model = simulation.model + + grid = model.grid + arch = architecture(grid) + + params = KernelParameters(model.tracers[1]) + + Uⁿ = dissipation_computation.previous_state.Uⁿ + Uⁿ⁻¹ = dissipation_computation.previous_state.Uⁿ⁻¹ + + U = model.velocities + + launch!(architecture(grid), grid, params, _update_transport!, Uⁿ, Uⁿ⁻¹, grid, U) + + for tracer_name in keys(dissipation_computation.production) + c = model.tracers[tracer_name] + cⁿ⁻¹ = dissipation_computation.previous_state.tracers[tracer_name] + Fⁿ = dissipation_computation.advective_fluxes.Fⁿ[tracer_fluxes] + Fⁿ⁻¹ = dissipation_computation.advective_fluxes.Fⁿ⁻¹[tracer_name] + A = model.advection[tracer_name] + + launch!(arch, grid, params, _update_fluxes!, Fⁿ, Fⁿ⁻¹, cⁿ⁻¹, grid, A, U, c) + end + + return nothing +end + +@kenrel 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] = Vⁿ.v[i, j, k] + Uⁿ⁻¹.w[i, j, k] = Wⁿ.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 + +@kernel function _update_fluxes!( 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) + end +end + +function assemble_P_values!(simulation, dissipation_computation) + model = simulation.model + grid = model.grid + arch = architecture(grid) + + χ = simulation.model.timestepper.χ + + # General velocities + Uⁿ = dissipation_computation.previous_state.Uⁿ + Uⁿ⁻¹ = dissipation_computation.previous_state.Uⁿ⁻¹ + + for tracer_names in keys(dissipation_computation.production) + c = model.tracers[tracer_name] + cⁿ⁻¹ = dissipation_computation.previous_state.tracers[tracer_name] + P = dissipation_computation.production[tracer_name] + Fⁿ = dissipation_computation.advective_fluxes.Fⁿ[tracer_name] + Fⁿ⁻¹ = dissipation_computation.advective_fluxes.Fⁿ⁻¹[tracer_name] + + launch!(arch, grid, :xyz, _compute_dissipation!, + P, grid, χ, Fⁿ, Fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, c, cⁿ⁻¹) + + 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]) + +@kernel function _compute_dissipation!(P, + grid, χ, + Fⁿ, Fⁿ⁻¹, + 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ⁿ⁻¹) + + @inbounds P.x[i, j, k] = compute_dissipation(i, j, k, grid, Uⁿ.u, Uⁿ⁻¹.u, χ, Fⁿ.x, Fⁿ⁻¹.x, δˣc★, δˣc²) + @inbounds P.y[i, j, k] = compute_dissipation(i, j, k, grid, Uⁿ.v, Uⁿ⁻¹.v, χ, Fⁿ.y, Fⁿ⁻¹.y, δʸc★, δʸc²) + @inbounds P.z[i, j, k] = compute_dissipation(i, j, k, grid, Uⁿ.w, Uⁿ⁻¹.w, χ, Fⁿ.z, Fⁿ⁻¹.z, δᶻc★, δᶻc²) +end + +@inline function compute_dissipation(i, j, k, grid, Uⁿ⁻¹, Uⁿ⁻², χ, fⁿ⁻¹, fⁿ⁻², δc★, δc²) + + C₁ = convert(eltype(grid), 1.5 + χ) + C₂ = convert(eltype(grid), 0.5 + χ) + + @inbounds begin + 𝒰ⁿ⁻¹ = C₁ * Uⁿ⁻¹[i, j, k] + 𝒰ⁿ⁻² = C₂ * Uⁿ⁻²[i, j, k] + Fⁿ⁻¹ = C₁ * fⁿ⁻¹[i, j, k] + Fⁿ⁻² = C₂ * fⁿ⁻²[i, j, k] + 𝒜 = Fⁿ⁻¹ - Fⁿ⁻² + 𝒟 = 𝒰ⁿ⁻¹ - 𝒰ⁿ⁻² + end + + return 2 * δc★ * 𝒜 - δc² * 𝒟 +end diff --git a/validation/implicit_dissipation/implicit_dissipation.jl b/validation/implicit_dissipation/implicit_dissipation.jl new file mode 100644 index 0000000000..e69de29bb2 From 6e6c8a963aded35a6ab7cdfbb376a0b004f4c8d1 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 28 Oct 2024 17:38:58 +0100 Subject: [PATCH 02/10] add a testcase --- src/Models/Models.jl | 3 + src/Models/dissipation_computation.jl | 85 ++++++++++--------- .../implicit_dissipation.jl | 73 ++++++++++++++++ 3 files changed, 123 insertions(+), 38 deletions(-) diff --git a/src/Models/Models.jl b/src/Models/Models.jl index 9399a5b503..d1e1fa8b3f 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -188,4 +188,7 @@ 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("dissipation_computation.jl") + end # module diff --git a/src/Models/dissipation_computation.jl b/src/Models/dissipation_computation.jl index fd28f52f89..2a46301888 100644 --- a/src/Models/dissipation_computation.jl +++ b/src/Models/dissipation_computation.jl @@ -1,24 +1,22 @@ 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.Advection: _advective_tracer_flux_x, _advective_tracer_flux_y, _advective_tracer_flux_z -using Oceananigans.Advection: horizontal_advection_U, horizontal_advection_V -using Oceananigans.Models.HydrostaticFreeSurfaceModels: vertical_scaling, previous_vertical_scaling +using KernelAbstractions: @kernel, @index import Oceananigans.Utils: KernelParameters -const c = Center() -const f = Face() - -struct DissipationComputation{P, A, S} +struct VarianceDissipationComputation{P, A, S} production :: P advective_fluxes :: A previous_state :: S end -function fluxes_field(grid) +function fluxes_fields(grid) x = XFaceField(grid) y = YFaceField(grid) z = ZFaceField(grid) @@ -26,16 +24,16 @@ function fluxes_field(grid) return (; x, y, z) end -function DissipationComputation(model; tracers = tracer_names(model.tracers)) +function VarianceDissipationComputation(model; tracers = propertynames(model.tracers)) - if !(model.timestepper isa QuasiAdamsBashforth2Timestepper) - throw(ArgumentError("DissipationComputation requires a QuasiAdamsBashforth2Timestepper")) + if !(model.timestepper isa QuasiAdamsBashforth2TimeStepper) + throw(ArgumentError("DissipationComputation requires a QuasiAdamsBashforth2TimeStepper")) end grid = model.grid - P = NamedTuple{tracers}(fluxes_field(grid) for tracer in tracers) - Fⁿ = NamedTuple{tracers}(fluxes_field(grid) for tracer in tracers) - Fⁿ⁻¹ = NamedTuple{tracers}(fluxes_field(grid) for tracer in tracers) + P = NamedTuple{tracers}(fluxes_fields(grid) for tracer in tracers) + Fⁿ = NamedTuple{tracers}(fluxes_fields(grid) for tracer in tracers) + Fⁿ⁻¹ = NamedTuple{tracers}(fluxes_fields(grid) for tracer in tracers) Uⁿ⁻¹ = VelocityFields(grid) Uⁿ = VelocityFields(grid) @@ -45,11 +43,15 @@ function DissipationComputation(model; tracers = tracer_names(model.tracers)) previous_state = merge(cⁿ⁻¹, (; Uⁿ⁻¹, Uⁿ)) advective_fluxes = (; Fⁿ, Fⁿ⁻¹) - return DissipationComputation(P, advective_fluxes, previous_state) + return VarianceDissipationComputation(P, advective_fluxes, previous_state) end # Function to call in a callback -function (dc::DissipationComputation)(simulation) +# 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 (dc::VarianceDissipationComputation)(simulation) # We first assemble values for Pⁿ⁻¹ assemble_P_values!(simulation, dc) @@ -65,6 +67,9 @@ end return KernelParameters(sz, of) end +@inline getadvection(advection, tracer_name) = advection +@inline getadvection(advection::NamedTuple, tracer_name) = @inbounds advection[tracer_name] + function update_fluxes!(simulation, dissipation_computation) model = simulation.model @@ -82,10 +87,10 @@ function update_fluxes!(simulation, dissipation_computation) for tracer_name in keys(dissipation_computation.production) c = model.tracers[tracer_name] - cⁿ⁻¹ = dissipation_computation.previous_state.tracers[tracer_name] - Fⁿ = dissipation_computation.advective_fluxes.Fⁿ[tracer_fluxes] + cⁿ⁻¹ = dissipation_computation.previous_state[tracer_name] + Fⁿ = dissipation_computation.advective_fluxes.Fⁿ[tracer_name] Fⁿ⁻¹ = dissipation_computation.advective_fluxes.Fⁿ⁻¹[tracer_name] - A = model.advection[tracer_name] + A = getadvection(model.advection, tracer_name) launch!(arch, grid, params, _update_fluxes!, Fⁿ, Fⁿ⁻¹, cⁿ⁻¹, grid, A, U, c) end @@ -93,13 +98,13 @@ function update_fluxes!(simulation, dissipation_computation) return nothing end -@kenrel function _update_transport!(Uⁿ, Uⁿ⁻¹, grid, U) +@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] = Vⁿ.v[i, j, k] - Uⁿ⁻¹.w[i, j, k] = Wⁿ.w[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) @@ -136,15 +141,19 @@ function assemble_P_values!(simulation, dissipation_computation) Uⁿ = dissipation_computation.previous_state.Uⁿ Uⁿ⁻¹ = dissipation_computation.previous_state.Uⁿ⁻¹ - for tracer_names in keys(dissipation_computation.production) + for tracer_name in keys(dissipation_computation.production) c = model.tracers[tracer_name] - cⁿ⁻¹ = dissipation_computation.previous_state.tracers[tracer_name] + cⁿ⁻¹ = dissipation_computation.previous_state[tracer_name] P = dissipation_computation.production[tracer_name] Fⁿ = dissipation_computation.advective_fluxes.Fⁿ[tracer_name] Fⁿ⁻¹ = dissipation_computation.advective_fluxes.Fⁿ⁻¹[tracer_name] - launch!(arch, grid, :xyz, _compute_dissipation!, - P, grid, χ, Fⁿ, Fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, c, cⁿ⁻¹) + launch!(arch, grid, :xyz, _compute_dissipation!, P, grid, χ, + Fⁿ, Fⁿ⁻¹, + Uⁿ, Uⁿ⁻¹, + c, cⁿ⁻¹) + + end return nothing end @@ -170,24 +179,24 @@ end δᶻc★ = δzᶜᶜᶠ(i, j, k, grid, c★, cⁿ, cⁿ⁻¹) δᶻc² = δzᶜᶜᶠ(i, j, k, grid, c², cⁿ, cⁿ⁻¹) - @inbounds P.x[i, j, k] = compute_dissipation(i, j, k, grid, Uⁿ.u, Uⁿ⁻¹.u, χ, Fⁿ.x, Fⁿ⁻¹.x, δˣc★, δˣc²) - @inbounds P.y[i, j, k] = compute_dissipation(i, j, k, grid, Uⁿ.v, Uⁿ⁻¹.v, χ, Fⁿ.y, Fⁿ⁻¹.y, δʸc★, δʸc²) - @inbounds P.z[i, j, k] = compute_dissipation(i, j, k, grid, Uⁿ.w, Uⁿ⁻¹.w, χ, Fⁿ.z, Fⁿ⁻¹.z, δᶻc★, δᶻc²) + @inbounds P.x[i, j, k] = compute_dissipation(i, j, k, grid, χ, Fⁿ.x, Fⁿ⁻¹.x, Uⁿ.u, Uⁿ⁻¹.u, δˣc★, δˣc²) + @inbounds P.y[i, j, k] = compute_dissipation(i, j, k, grid, χ, Fⁿ.y, Fⁿ⁻¹.y, Uⁿ.v, Uⁿ⁻¹.v, δʸc★, δʸc²) + @inbounds P.z[i, j, k] = compute_dissipation(i, j, k, grid, χ, Fⁿ.z, Fⁿ⁻¹.z, Uⁿ.w, Uⁿ⁻¹.w, δᶻc★, δᶻc²) end -@inline function compute_dissipation(i, j, k, grid, Uⁿ⁻¹, Uⁿ⁻², χ, fⁿ⁻¹, fⁿ⁻², δc★, δc²) +@inline function compute_dissipation(i, j, k, grid, χ, fⁿ, fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, δc★, δc²) C₁ = convert(eltype(grid), 1.5 + χ) C₂ = convert(eltype(grid), 0.5 + χ) @inbounds begin - 𝒰ⁿ⁻¹ = C₁ * Uⁿ⁻¹[i, j, k] - 𝒰ⁿ⁻² = C₂ * Uⁿ⁻²[i, j, k] - Fⁿ⁻¹ = C₁ * fⁿ⁻¹[i, j, k] - Fⁿ⁻² = C₂ * fⁿ⁻²[i, j, k] - 𝒜 = Fⁿ⁻¹ - Fⁿ⁻² - 𝒟 = 𝒰ⁿ⁻¹ - 𝒰ⁿ⁻² + 𝒰ⁿ = C₁ * Uⁿ[i, j, k] + 𝒰ⁿ⁻¹ = C₂ * Uⁿ⁻¹[i, j, k] + Fⁿ = C₁ * fⁿ[i, j, k] + Fⁿ⁻¹ = C₂ * fⁿ⁻¹[i, j, k] + A = Fⁿ - Fⁿ⁻¹ + D = 𝒰ⁿ - 𝒰ⁿ⁻¹ end - - return 2 * δc★ * 𝒜 - δc² * 𝒟 -end + + return 2 * δc★ * A - δc² * D +end diff --git a/validation/implicit_dissipation/implicit_dissipation.jl b/validation/implicit_dissipation/implicit_dissipation.jl index e69de29bb2..edd232fd74 100644 --- a/validation/implicit_dissipation/implicit_dissipation.jl +++ b/validation/implicit_dissipation/implicit_dissipation.jl @@ -0,0 +1,73 @@ +using Oceananigans +using Oceananigans.Fields: VelocityFields +using Oceananigans.Models.HydrostaticFreeSurfaceModels: PrescribedVelocityFields +using Oceananigans.Models: VarianceDissipationComputation +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 + + +grid = RectilinearGrid(size = N, halo = 6, x = (-1, 1), topology = (Periodic, Flat, Flat)) + +model = NonhydrostaticModel(; grid, tracers = :b, timestepper = :QuasiAdamsBashforth2, advection) + +set!(model.tracers.b, bᵢ) +set!(model.velocities.u, 1.0) + +b₀ = deepcopy(model.tracers.b) +Δt = 0.1 * minimum_xspacing(grid) + +dissipation_computation = VarianceDissipationComputation(model) +simulation = Simulation(model; Δt, stop_time = 10) +simulation.callbacks[:compute_dissipation] = Callback(dissipation_computation, IterationInterval(1)) + +outputs = (; Px = dissipation_computation.production.b.x, b = model.tracers.b) + +# Save the dissipation +simulation.output_writers[:dissipation] = JLD2OutputWriter(model, dissipation_computation.production.b, + filename = "dissipation.jld2", + schedule = IterationInterval(10)) + +run!(simulation) +b = FieldTimesSeries("dissipation.jld2", "b") +P = FieldTimesSeries("dissipation.jld2", "Px") + +iter = Observable(1) + +bn = @lift(interior(b[$iter], :, 1, 1)) +Pn = @lift(interior(b[$iter], :, 1, 1)) + +fig = Figure(size = (800, 400)) +ax = Axis(fig[1, 1], xlabel = L"x", ylabel = L"tracer") +lines!(ax, xnodes(b[1]), b0, color = :grey, linestyle = :dash, linewidth = 2) +lines!(ax, xnodes(b[1]), bn, color = :blue) + +ax = Axis(fig[1, 1], xlabel = L"x", ylabel = L"variance dissipation") +lines!(ax, xnodes(b[1]), Pn, color = :red) + +record(fig, "implicit_dissipation.mp4", 1:length(b), framerate=8) do i + n[] = i +end \ No newline at end of file From b74ee6888401bd43bb44434597e38f0db0cc6e4e Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 28 Oct 2024 17:42:51 +0100 Subject: [PATCH 03/10] this should work --- src/Models/dissipation_computation.jl | 4 +++- .../implicit_dissipation.jl | 17 +++++++++-------- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/Models/dissipation_computation.jl b/src/Models/dissipation_computation.jl index 2a46301888..c5b905437b 100644 --- a/src/Models/dissipation_computation.jl +++ b/src/Models/dissipation_computation.jl @@ -25,10 +25,12 @@ function fluxes_fields(grid) end function VarianceDissipationComputation(model; tracers = propertynames(model.tracers)) - + if !(model.timestepper isa QuasiAdamsBashforth2TimeStepper) throw(ArgumentError("DissipationComputation requires a QuasiAdamsBashforth2TimeStepper")) end + + tracers = tupleit(tracers) grid = model.grid P = NamedTuple{tracers}(fluxes_fields(grid) for tracer in tracers) diff --git a/validation/implicit_dissipation/implicit_dissipation.jl b/validation/implicit_dissipation/implicit_dissipation.jl index edd232fd74..1bd14bd729 100644 --- a/validation/implicit_dissipation/implicit_dissipation.jl +++ b/validation/implicit_dissipation/implicit_dissipation.jl @@ -30,9 +30,9 @@ const α = 10 end -grid = RectilinearGrid(size = N, halo = 6, x = (-1, 1), topology = (Periodic, Flat, Flat)) - -model = NonhydrostaticModel(; grid, tracers = :b, timestepper = :QuasiAdamsBashforth2, advection) +grid = RectilinearGrid(size = 100, halo = 6, x = (-1, 1), topology = (Periodic, Flat, Flat)) +advection = WENO(; order = 7) +model = NonhydrostaticModel(; grid, tracers = :b, timestepper = :QuasiAdamsBashforth2, advection) set!(model.tracers.b, bᵢ) set!(model.velocities.u, 1.0) @@ -40,20 +40,21 @@ set!(model.velocities.u, 1.0) b₀ = deepcopy(model.tracers.b) Δt = 0.1 * minimum_xspacing(grid) -dissipation_computation = VarianceDissipationComputation(model) +dissipation_computation = VarianceDissipationComputation(model; tracers = :b) simulation = Simulation(model; Δt, stop_time = 10) simulation.callbacks[:compute_dissipation] = Callback(dissipation_computation, IterationInterval(1)) outputs = (; Px = dissipation_computation.production.b.x, b = model.tracers.b) # Save the dissipation -simulation.output_writers[:dissipation] = JLD2OutputWriter(model, dissipation_computation.production.b, +simulation.output_writers[:dissipation] = JLD2OutputWriter(model, outputs, filename = "dissipation.jld2", - schedule = IterationInterval(10)) + schedule = IterationInterval(1), + overwrite_existing = true) run!(simulation) -b = FieldTimesSeries("dissipation.jld2", "b") -P = FieldTimesSeries("dissipation.jld2", "Px") +b = FieldTimeSeries("dissipation.jld2", "b") +P = FieldTimeSeries("dissipation.jld2", "Px") iter = Observable(1) From 281ce84d1019a331fb272865e50cf59ab405cfad Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 28 Oct 2024 17:53:04 +0100 Subject: [PATCH 04/10] add a test --- .../implicit_dissipation.jl | 78 +++++++++++++------ 1 file changed, 53 insertions(+), 25 deletions(-) diff --git a/validation/implicit_dissipation/implicit_dissipation.jl b/validation/implicit_dissipation/implicit_dissipation.jl index 1bd14bd729..97105103d8 100644 --- a/validation/implicit_dissipation/implicit_dissipation.jl +++ b/validation/implicit_dissipation/implicit_dissipation.jl @@ -29,46 +29,74 @@ const α = 10 end end +grid = RectilinearGrid(size = 100, halo = 6, x = (-1, 1), topology = (Periodic, Flat, Flat)) -grid = RectilinearGrid(size = 100, halo = 6, x = (-1, 1), topology = (Periodic, Flat, Flat)) -advection = WENO(; order = 7) -model = NonhydrostaticModel(; grid, tracers = :b, timestepper = :QuasiAdamsBashforth2, advection) +advections = [UpwindBiased(; order = 3), + WENO(; order = 5), + WENO(; order = 7)] -set!(model.tracers.b, bᵢ) -set!(model.velocities.u, 1.0) +labels = ["upwind3", + "WENO5", + "WENO7"] -b₀ = deepcopy(model.tracers.b) -Δt = 0.1 * minimum_xspacing(grid) +for (advection, label) in zip(advections, labels) + model = NonhydrostaticModel(; grid, tracers = :b, timestepper = :QuasiAdamsBashforth2, advection) -dissipation_computation = VarianceDissipationComputation(model; tracers = :b) -simulation = Simulation(model; Δt, stop_time = 10) -simulation.callbacks[:compute_dissipation] = Callback(dissipation_computation, IterationInterval(1)) + set!(model.tracers.b, bᵢ) + set!(model.velocities.u, 1.0) -outputs = (; Px = dissipation_computation.production.b.x, b = model.tracers.b) + Δt = 0.1 * minimum_xspacing(grid) -# Save the dissipation -simulation.output_writers[:dissipation] = JLD2OutputWriter(model, outputs, - filename = "dissipation.jld2", - schedule = IterationInterval(1), - overwrite_existing = true) + dissipation_computation = VarianceDissipationComputation(model; tracers = :b) + simulation = Simulation(model; Δt, stop_time = 10) + simulation.callbacks[:compute_dissipation] = Callback(dissipation_computation, IterationInterval(1)) -run!(simulation) -b = FieldTimeSeries("dissipation.jld2", "b") -P = FieldTimeSeries("dissipation.jld2", "Px") + outputs = (; Px = dissipation_computation.production.b.x, b = model.tracers.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 + +b_series = [] +P_series = [] iter = Observable(1) -bn = @lift(interior(b[$iter], :, 1, 1)) -Pn = @lift(interior(b[$iter], :, 1, 1)) +bn = [] +Pn = [] + +for label in labels + filename = "dissipation_" * label * ".jld2" + + push!(b_series, FieldTimeSeries(filename, "b")) + push!(P_series, FieldTimeSeries(filename, "Px")) + + push!(bn, @lift(interior(b_series[end][$iter], :, 1, 1))) + push!(Pn, @lift(interior(P_series[end][$iter], :, 1, 1))) +end + +x = xnodes(b_series[1][1]) fig = Figure(size = (800, 400)) ax = Axis(fig[1, 1], xlabel = L"x", ylabel = L"tracer") -lines!(ax, xnodes(b[1]), b0, color = :grey, linestyle = :dash, linewidth = 2) -lines!(ax, xnodes(b[1]), bn, color = :blue) +lines!(ax, xnodes(b[1]), interior(b[1], :, 1, 1), color = :grey, linestyle = :dash, linewidth = 2) +for case in eachindex(labels) + lines!(ax, x, bn[case], label = labels[case]) +end ax = Axis(fig[1, 1], xlabel = L"x", ylabel = L"variance dissipation") -lines!(ax, xnodes(b[1]), Pn, color = :red) +for case in eachindex(labels) + lines!(ax, x, Pn[case], label = labels[case]) +end record(fig, "implicit_dissipation.mp4", 1:length(b), framerate=8) do i - n[] = i + @info "doing iter $i" + iter[] = i end \ No newline at end of file From 5abcd269de1698425de656bd88c230d2f1941346 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 28 Oct 2024 18:09:36 +0100 Subject: [PATCH 05/10] new dissipation --- src/Advection/adapt_advection_order.jl | 2 +- src/Models/dissipation_computation.jl | 11 +++-- .../implicit_dissipation.jl | 47 ++++++++++++------- 3 files changed, 38 insertions(+), 22 deletions(-) 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/dissipation_computation.jl b/src/Models/dissipation_computation.jl index c5b905437b..60b8c7b7b8 100644 --- a/src/Models/dissipation_computation.jl +++ b/src/Models/dissipation_computation.jl @@ -1,3 +1,4 @@ +using Oceananigans: instantiated_location using Oceananigans.Grids: architecture using Oceananigans.Utils using Oceananigans.TimeSteppers @@ -6,6 +7,7 @@ using Oceananigans.Fields: Field, VelocityFields using Oceananigans.Operators using Oceananigans.BoundaryConditions using Oceananigans.Advection: _advective_tracer_flux_x, _advective_tracer_flux_y, _advective_tracer_flux_z +using Oceananigans.Operators: volume using KernelAbstractions: @kernel, @index import Oceananigans.Utils: KernelParameters @@ -29,7 +31,7 @@ function VarianceDissipationComputation(model; tracers = propertynames(model.tra if !(model.timestepper isa QuasiAdamsBashforth2TimeStepper) throw(ArgumentError("DissipationComputation requires a QuasiAdamsBashforth2TimeStepper")) end - + tracers = tupleit(tracers) grid = model.grid @@ -188,8 +190,9 @@ end @inline function compute_dissipation(i, j, k, grid, χ, fⁿ, fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, δc★, δc²) - C₁ = convert(eltype(grid), 1.5 + χ) - C₂ = convert(eltype(grid), 0.5 + χ) + C₁ = convert(eltype(grid), 1.5 + χ) + C₂ = convert(eltype(grid), 0.5 + χ) + loc = instantiated_location(Uⁿ) @inbounds begin 𝒰ⁿ = C₁ * Uⁿ[i, j, k] @@ -200,5 +203,5 @@ end D = 𝒰ⁿ - 𝒰ⁿ⁻¹ end - return 2 * δc★ * A - δc² * D + return (2 * δc★ * A - δc² * D) / volume(i, j, k, grid, loc...) end diff --git a/validation/implicit_dissipation/implicit_dissipation.jl b/validation/implicit_dissipation/implicit_dissipation.jl index 97105103d8..243dfeeb78 100644 --- a/validation/implicit_dissipation/implicit_dissipation.jl +++ b/validation/implicit_dissipation/implicit_dissipation.jl @@ -29,17 +29,8 @@ const α = 10 end 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"] +function one_dimensional_simulation(grid, advection, label) -for (advection, label) in zip(advections, labels) model = NonhydrostaticModel(; grid, tracers = :b, timestepper = :QuasiAdamsBashforth2, advection) set!(model.tracers.b, bᵢ) @@ -51,8 +42,12 @@ for (advection, label) in zip(advections, labels) simulation = Simulation(model; Δt, stop_time = 10) simulation.callbacks[:compute_dissipation] = Callback(dissipation_computation, IterationInterval(1)) - outputs = (; Px = dissipation_computation.production.b.x, b = model.tracers.b) + Px = dissipation_computation.production.b.x + bⁿ⁻¹ = dissipation_computation.previous_state.b + bⁿ = model.tracers.b + Σb² = @at((Center, Center, Center), Px) + (bⁿ^2 - bⁿ⁻¹^2) / Δt + outputs = (; Px, b = model.tracers.b, Σb²) filename = "dissipation_" * label * ".jld2" # Save the dissipation @@ -64,6 +59,22 @@ for (advection, label) in zip(advections, labels) 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 = [] @@ -72,31 +83,33 @@ iter = Observable(1) bn = [] Pn = [] -for label in labels +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[end][$iter], :, 1, 1))) - push!(Pn, @lift(interior(P_series[end][$iter], :, 1, 1))) + push!(bn, @lift(interior(b_series[i][$iter], :, 1, 1))) + push!(Pn, @lift(interior(P_series[i][$iter], :, 1, 1))) end x = xnodes(b_series[1][1]) -fig = Figure(size = (800, 400)) +fig = Figure(size = (1200, 400)) ax = Axis(fig[1, 1], xlabel = L"x", ylabel = L"tracer") lines!(ax, xnodes(b[1]), interior(b[1], :, 1, 1), color = :grey, linestyle = :dash, linewidth = 2) for case in eachindex(labels) lines!(ax, x, bn[case], label = labels[case]) end -ax = Axis(fig[1, 1], xlabel = L"x", ylabel = L"variance dissipation") +ax = Axis(fig[1, 2], xlabel = L"x", ylabel = L"variance dissipation") for case in eachindex(labels) lines!(ax, x, Pn[case], label = labels[case]) end +ylims!(ax, (-1e-2, 1e-2)) -record(fig, "implicit_dissipation.mp4", 1:length(b), framerate=8) do i +record(fig, "implicit_dissipation.mp4", 1:length(b_series[1]), framerate=8) do i @info "doing iter $i" iter[] = i end \ No newline at end of file From cb12eaa69e7976024f399c0b7c00373f643441dd Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 28 Oct 2024 18:20:13 +0100 Subject: [PATCH 06/10] fix a bit the implicit dissipation validation --- .../implicit_dissipation.jl | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/validation/implicit_dissipation/implicit_dissipation.jl b/validation/implicit_dissipation/implicit_dissipation.jl index 243dfeeb78..45c6208d3c 100644 --- a/validation/implicit_dissipation/implicit_dissipation.jl +++ b/validation/implicit_dissipation/implicit_dissipation.jl @@ -46,7 +46,7 @@ function one_dimensional_simulation(grid, advection, label) bⁿ⁻¹ = dissipation_computation.previous_state.b bⁿ = model.tracers.b - Σb² = @at((Center, Center, Center), Px) + (bⁿ^2 - bⁿ⁻¹^2) / Δt + Σb² = @at((Center, Center, Center), Px) + (bⁿ^2 - bⁿ⁻¹^2) / Δt # Actually not at the same time-step outputs = (; Px, b = model.tracers.b, Σb²) filename = "dissipation_" * label * ".jld2" @@ -82,6 +82,7 @@ iter = Observable(1) bn = [] Pn = [] +Bn = [] for (i, label) in enumerate(labels) filename = "dissipation_" * label * ".jld2" @@ -92,22 +93,25 @@ for (i, label) in enumerate(labels) 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 x = xnodes(b_series[1][1]) fig = Figure(size = (1200, 400)) ax = Axis(fig[1, 1], xlabel = L"x", ylabel = L"tracer") -lines!(ax, xnodes(b[1]), interior(b[1], :, 1, 1), color = :grey, linestyle = :dash, linewidth = 2) -for case in eachindex(labels) - lines!(ax, x, bn[case], label = labels[case]) -end +lines!(ax, xnodes(b[1]), interior(b[1], :, 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") -for case in eachindex(labels) - lines!(ax, x, Pn[case], label = labels[case]) -end -ylims!(ax, (-1e-2, 1e-2)) +lines!(ax, x, Pn[1], color = :red , label = labels[1]) +lines!(ax, x, Pn[1], 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" From 00a1b8fcc494a8f745c5c3e5e909b1bb9f2955e5 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Mon, 28 Oct 2024 18:24:28 +0100 Subject: [PATCH 07/10] some name change --- src/Models/dissipation_computation.jl | 8 ++++---- validation/implicit_dissipation/implicit_dissipation.jl | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/Models/dissipation_computation.jl b/src/Models/dissipation_computation.jl index 60b8c7b7b8..44ae6218e4 100644 --- a/src/Models/dissipation_computation.jl +++ b/src/Models/dissipation_computation.jl @@ -12,7 +12,7 @@ using KernelAbstractions: @kernel, @index import Oceananigans.Utils: KernelParameters -struct VarianceDissipationComputation{P, A, S} +struct TracerVarianceDissipation{P, A, S} production :: P advective_fluxes :: A previous_state :: S @@ -26,7 +26,7 @@ function fluxes_fields(grid) return (; x, y, z) end -function VarianceDissipationComputation(model; tracers = propertynames(model.tracers)) +function TracerVarianceDissipation(model; tracers = propertynames(model.tracers)) if !(model.timestepper isa QuasiAdamsBashforth2TimeStepper) throw(ArgumentError("DissipationComputation requires a QuasiAdamsBashforth2TimeStepper")) @@ -47,7 +47,7 @@ function VarianceDissipationComputation(model; tracers = propertynames(model.tra previous_state = merge(cⁿ⁻¹, (; Uⁿ⁻¹, Uⁿ)) advective_fluxes = (; Fⁿ, Fⁿ⁻¹) - return VarianceDissipationComputation(P, advective_fluxes, previous_state) + return TracerVarianceDissipation(P, advective_fluxes, previous_state) end # Function to call in a callback @@ -55,7 +55,7 @@ end # 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 (dc::VarianceDissipationComputation)(simulation) +function (dc::TracerVarianceDissipation)(simulation) # We first assemble values for Pⁿ⁻¹ assemble_P_values!(simulation, dc) diff --git a/validation/implicit_dissipation/implicit_dissipation.jl b/validation/implicit_dissipation/implicit_dissipation.jl index 45c6208d3c..d2d758903c 100644 --- a/validation/implicit_dissipation/implicit_dissipation.jl +++ b/validation/implicit_dissipation/implicit_dissipation.jl @@ -1,7 +1,7 @@ using Oceananigans using Oceananigans.Fields: VelocityFields using Oceananigans.Models.HydrostaticFreeSurfaceModels: PrescribedVelocityFields -using Oceananigans.Models: VarianceDissipationComputation +using Oceananigans.Models: TracerVarianceDissipation using Printf using GLMakie @@ -38,7 +38,7 @@ function one_dimensional_simulation(grid, advection, label) Δt = 0.1 * minimum_xspacing(grid) - dissipation_computation = VarianceDissipationComputation(model; tracers = :b) + dissipation_computation = TracerVarianceDissipation(model; tracers = :b) simulation = Simulation(model; Δt, stop_time = 10) simulation.callbacks[:compute_dissipation] = Callback(dissipation_computation, IterationInterval(1)) @@ -46,7 +46,7 @@ function one_dimensional_simulation(grid, advection, label) bⁿ⁻¹ = dissipation_computation.previous_state.b bⁿ = model.tracers.b - Σb² = @at((Center, Center, Center), Px) + (bⁿ^2 - bⁿ⁻¹^2) / Δt # Actually not at the same time-step + Σ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" @@ -107,7 +107,7 @@ 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[1], color = :blue, label = labels[2]) +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)) From 945215654351fa6e67e3f885bee5ff2164a40640 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Tue, 29 Oct 2024 13:18:05 +0100 Subject: [PATCH 08/10] testing tracer variance budget --- src/Oceananigans.jl | 1 + src/TimeSteppers/quasi_adams_bashforth_2.jl | 3 ++ .../implicit_dissipation.jl | 43 ++++++++++++++++--- 3 files changed, 41 insertions(+), 6 deletions(-) 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 index d2d758903c..6bb7135208 100644 --- a/validation/implicit_dissipation/implicit_dissipation.jl +++ b/validation/implicit_dissipation/implicit_dissipation.jl @@ -29,9 +29,20 @@ const α = 10 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) - model = NonhydrostaticModel(; grid, tracers = :b, timestepper = :QuasiAdamsBashforth2, advection) + 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) @@ -40,10 +51,11 @@ function one_dimensional_simulation(grid, advection, label) dissipation_computation = TracerVarianceDissipation(model; tracers = :b) simulation = Simulation(model; Δt, stop_time = 10) - simulation.callbacks[:compute_dissipation] = Callback(dissipation_computation, IterationInterval(1)) + 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ⁿ⁻¹ = dissipation_computation.previous_state.b + 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 @@ -96,11 +108,12 @@ for (i, label) in enumerate(labels) push!(Bn, @lift(interior(B_series[i][$iter], :, 1, 1))) end -x = xnodes(b_series[1][1]) +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, xnodes(b[1]), interior(b[1], :, 1, 1), color = :grey, linestyle = :dash, linewidth = 2, label = "initial condition") +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) @@ -116,4 +129,22 @@ 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 \ No newline at end of file +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 From 1c193dea2a5563bb7065f5588ca42565c27008e0 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 7 Nov 2024 14:59:49 +0100 Subject: [PATCH 09/10] new implementation that includes closures --- src/Models/Models.jl | 4 +- .../VarianceDissipationComputation.jl | 118 ++++++++++++++++++ .../advective_dissipation.jl | 61 +++++++++ .../advective_fluxes.jl | 42 +++++++ .../assemble_dissipation.jl | 61 +++++++++ .../diffusive_dissipation.jl | 49 ++++++++ .../diffusive_fluxes.jl | 40 ++++++ .../dissipation_utils.jl | 48 +++++++ .../get_dissipation_fields.jl | 27 ++++ .../update_fluxes.jl | 84 +++++++++++++ 10 files changed, 533 insertions(+), 1 deletion(-) create mode 100644 src/Models/VarianceDissipationComputation/VarianceDissipationComputation.jl create mode 100644 src/Models/VarianceDissipationComputation/advective_dissipation.jl create mode 100644 src/Models/VarianceDissipationComputation/advective_fluxes.jl create mode 100644 src/Models/VarianceDissipationComputation/assemble_dissipation.jl create mode 100644 src/Models/VarianceDissipationComputation/diffusive_dissipation.jl create mode 100644 src/Models/VarianceDissipationComputation/diffusive_fluxes.jl create mode 100644 src/Models/VarianceDissipationComputation/dissipation_utils.jl create mode 100644 src/Models/VarianceDissipationComputation/get_dissipation_fields.jl create mode 100644 src/Models/VarianceDissipationComputation/update_fluxes.jl diff --git a/src/Models/Models.jl b/src/Models/Models.jl index d1e1fa8b3f..78aab8b8bf 100644 --- a/src/Models/Models.jl +++ b/src/Models/Models.jl @@ -189,6 +189,8 @@ default_nan_checker(::OnlyParticleTrackingModel) = nothing include("seawater_density.jl") # Implementation of the diagnostic for computing the dissipation rate -include("dissipation_computation.jl") +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..1a7f0b2539 --- /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 TracerVarianceDissipation(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 (ϵ::TracerVarianceDissipation)(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..af717a61a3 --- /dev/null +++ b/src/Models/VarianceDissipationComputation/get_dissipation_fields.jl @@ -0,0 +1,27 @@ +function get_dissipation_fields(t::TracerVarianceDissipation) + 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::TracerVarianceDissipation, 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! + From b09bc35b2ae18242c5bf58659a028326657bc06f Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Thu, 7 Nov 2024 15:00:29 +0100 Subject: [PATCH 10/10] new implementation --- .../VarianceDissipationComputation.jl | 4 +- .../get_dissipation_fields.jl | 4 +- src/Models/dissipation_computation.jl | 207 ------------------ 3 files changed, 4 insertions(+), 211 deletions(-) delete mode 100644 src/Models/dissipation_computation.jl diff --git a/src/Models/VarianceDissipationComputation/VarianceDissipationComputation.jl b/src/Models/VarianceDissipationComputation/VarianceDissipationComputation.jl index 1a7f0b2539..7138682e47 100644 --- a/src/Models/VarianceDissipationComputation/VarianceDissipationComputation.jl +++ b/src/Models/VarianceDissipationComputation/VarianceDissipationComputation.jl @@ -90,7 +90,7 @@ function VarianceDissipation(model; gradients = deepcopy(P) - return TracerVarianceDissipation(P, K, advective_fluxes, diffusive_fluxes, previous_state, gradients) + return VarianceDissipation(P, K, advective_fluxes, diffusive_fluxes, previous_state, gradients) end # Function to call in a callback @@ -98,7 +98,7 @@ end # 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 (ϵ::TracerVarianceDissipation)(simulation) +function (ϵ::VarianceDissipation)(simulation) # We first assemble values for Pⁿ⁻¹ assemble_dissipation!(simulation, ϵ) diff --git a/src/Models/VarianceDissipationComputation/get_dissipation_fields.jl b/src/Models/VarianceDissipationComputation/get_dissipation_fields.jl index af717a61a3..c1c92c0708 100644 --- a/src/Models/VarianceDissipationComputation/get_dissipation_fields.jl +++ b/src/Models/VarianceDissipationComputation/get_dissipation_fields.jl @@ -1,4 +1,4 @@ -function get_dissipation_fields(t::TracerVarianceDissipation) +function get_dissipation_fields(t::VarianceDissipation) f = NamedTuple() for name in keys(t.advective_production) @@ -8,7 +8,7 @@ function get_dissipation_fields(t::TracerVarianceDissipation) return f end -function get_dissipation_fields(t::TracerVarianceDissipation, tracer_name) +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] diff --git a/src/Models/dissipation_computation.jl b/src/Models/dissipation_computation.jl deleted file mode 100644 index 44ae6218e4..0000000000 --- a/src/Models/dissipation_computation.jl +++ /dev/null @@ -1,207 +0,0 @@ -using Oceananigans: instantiated_location -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.Advection: _advective_tracer_flux_x, _advective_tracer_flux_y, _advective_tracer_flux_z -using Oceananigans.Operators: volume -using KernelAbstractions: @kernel, @index - -import Oceananigans.Utils: KernelParameters - -struct TracerVarianceDissipation{P, A, S} - production :: P - advective_fluxes :: A - previous_state :: S -end - -function fluxes_fields(grid) - x = XFaceField(grid) - y = YFaceField(grid) - z = ZFaceField(grid) - - return (; x, y, z) -end - -function TracerVarianceDissipation(model; tracers = propertynames(model.tracers)) - - if !(model.timestepper isa QuasiAdamsBashforth2TimeStepper) - throw(ArgumentError("DissipationComputation requires a QuasiAdamsBashforth2TimeStepper")) - end - - tracers = tupleit(tracers) - - grid = model.grid - P = NamedTuple{tracers}(fluxes_fields(grid) for tracer in tracers) - Fⁿ = NamedTuple{tracers}(fluxes_fields(grid) for tracer in tracers) - Fⁿ⁻¹ = NamedTuple{tracers}(fluxes_fields(grid) for tracer in tracers) - - Uⁿ⁻¹ = VelocityFields(grid) - Uⁿ = VelocityFields(grid) - - cⁿ⁻¹ = NamedTuple{tracers}(CenterField(grid) for tracer in tracers) - - previous_state = merge(cⁿ⁻¹, (; Uⁿ⁻¹, Uⁿ)) - advective_fluxes = (; Fⁿ, Fⁿ⁻¹) - - return TracerVarianceDissipation(P, advective_fluxes, previous_state) -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 (dc::TracerVarianceDissipation)(simulation) - # We first assemble values for Pⁿ⁻¹ - assemble_P_values!(simulation, dc) - - # Then we update the fluxes to be used in the next time step - update_fluxes!(simulation, dc) - - return nothing -end - -@inline function KernelParameters(f::Field) - sz = size(f.data) - of = f.data.offsets - return KernelParameters(sz, of) -end - -@inline getadvection(advection, tracer_name) = advection -@inline getadvection(advection::NamedTuple, tracer_name) = @inbounds advection[tracer_name] - -function update_fluxes!(simulation, dissipation_computation) - model = simulation.model - - grid = model.grid - arch = architecture(grid) - - params = KernelParameters(model.tracers[1]) - - Uⁿ = dissipation_computation.previous_state.Uⁿ - Uⁿ⁻¹ = dissipation_computation.previous_state.Uⁿ⁻¹ - - U = model.velocities - - launch!(architecture(grid), grid, params, _update_transport!, Uⁿ, Uⁿ⁻¹, grid, U) - - for tracer_name in keys(dissipation_computation.production) - c = model.tracers[tracer_name] - cⁿ⁻¹ = dissipation_computation.previous_state[tracer_name] - Fⁿ = dissipation_computation.advective_fluxes.Fⁿ[tracer_name] - Fⁿ⁻¹ = dissipation_computation.advective_fluxes.Fⁿ⁻¹[tracer_name] - A = getadvection(model.advection, tracer_name) - - launch!(arch, grid, params, _update_fluxes!, Fⁿ, Fⁿ⁻¹, cⁿ⁻¹, grid, A, U, c) - end - - 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 - -@kernel function _update_fluxes!( 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) - end -end - -function assemble_P_values!(simulation, dissipation_computation) - model = simulation.model - grid = model.grid - arch = architecture(grid) - - χ = simulation.model.timestepper.χ - - # General velocities - Uⁿ = dissipation_computation.previous_state.Uⁿ - Uⁿ⁻¹ = dissipation_computation.previous_state.Uⁿ⁻¹ - - for tracer_name in keys(dissipation_computation.production) - c = model.tracers[tracer_name] - cⁿ⁻¹ = dissipation_computation.previous_state[tracer_name] - P = dissipation_computation.production[tracer_name] - Fⁿ = dissipation_computation.advective_fluxes.Fⁿ[tracer_name] - Fⁿ⁻¹ = dissipation_computation.advective_fluxes.Fⁿ⁻¹[tracer_name] - - launch!(arch, grid, :xyz, _compute_dissipation!, P, grid, χ, - Fⁿ, Fⁿ⁻¹, - Uⁿ, Uⁿ⁻¹, - c, cⁿ⁻¹) - - 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]) - -@kernel function _compute_dissipation!(P, - grid, χ, - Fⁿ, Fⁿ⁻¹, - 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ⁿ⁻¹) - - @inbounds P.x[i, j, k] = compute_dissipation(i, j, k, grid, χ, Fⁿ.x, Fⁿ⁻¹.x, Uⁿ.u, Uⁿ⁻¹.u, δˣc★, δˣc²) - @inbounds P.y[i, j, k] = compute_dissipation(i, j, k, grid, χ, Fⁿ.y, Fⁿ⁻¹.y, Uⁿ.v, Uⁿ⁻¹.v, δʸc★, δʸc²) - @inbounds P.z[i, j, k] = compute_dissipation(i, j, k, grid, χ, Fⁿ.z, Fⁿ⁻¹.z, Uⁿ.w, Uⁿ⁻¹.w, δᶻc★, δᶻc²) -end - -@inline function compute_dissipation(i, j, k, grid, χ, fⁿ, fⁿ⁻¹, Uⁿ, Uⁿ⁻¹, δc★, δc²) - - C₁ = convert(eltype(grid), 1.5 + χ) - C₂ = convert(eltype(grid), 0.5 + χ) - loc = instantiated_location(Uⁿ) - - @inbounds begin - 𝒰ⁿ = C₁ * Uⁿ[i, j, k] - 𝒰ⁿ⁻¹ = C₂ * Uⁿ⁻¹[i, j, k] - Fⁿ = C₁ * fⁿ[i, j, k] - Fⁿ⁻¹ = C₂ * fⁿ⁻¹[i, j, k] - A = Fⁿ - Fⁿ⁻¹ - D = 𝒰ⁿ - 𝒰ⁿ⁻¹ - end - - return (2 * δc★ * A - δc² * D) / volume(i, j, k, grid, loc...) -end