advective form of the GM skew diffusivity (#3971)

* this should work

* remove all the @show

* restart tests

* will this work?

* test more than one position

* bugfix

* this should work

* remove unused truncate

* new interpolate

* I think this works

* try option 2

* no need for flip anymore

* fix gpu tests

* change comment and restore docstring

* fix the interpolate

* add (nothing nothing nothing) interpolate

* do not reuse names

* slogging along

* implementation is super simple

* works very well

* works, nice

* starting the change

* make it work

* works

* works but it's slow

* works nicely

* works nicely

* formatting

* including one adapt

* change to κ and test difference

* bugfix

* using ZeroField

* import correct stuff

* add docstring

* flat dimension

* add validarion case

* does this work?

* assumption

* formatting

* incorporate in ISSD

* correct tests

* only one skew diffusivity

* bugfix

* better comment

* change docstring

* more fixes

* fix testing

* keep same config for diffusive formulation

* error if we input a namedtuple

* bugfix

* extend for prescribed velocities

* correct immersed baroclinic adjustment

* implement suggestions

src/Biogeochemistry.jl

@@ -27,7 +27,7 @@ Update biogeochemical state variables. Called at the end of update_state!.
 """
 update_biogeochemical_state!(bgc, model) = nothing
 
-@inline biogeochemical_drift_velocity(bgc, val_tracer_name) = (u = ZeroField(), v = ZeroField(), w = ZeroField())
+@inline biogeochemical_drift_velocity(bgc, val_tracer_name) = nothing
 @inline biogeochemical_auxiliary_fields(bgc) = NamedTuple()
 
 """

src/Forcings/advective_forcing.jl

@@ -1,6 +1,6 @@
 using Oceananigans.Advection: div_Uc, div_𝐯u, div_𝐯v, div_𝐯w
 using Oceananigans.Fields: ZeroField, ConstantField
-using Oceananigans.Utils: SumOfArrays
+using Oceananigans.Utils: sum_of_velocities
 using Adapt
 
 maybe_constant_field(u) = u
@@ -72,14 +72,13 @@ Adapt.adapt_structure(to, af::AdvectiveForcing) =
 on_architecture(to, af::AdvectiveForcing) =
     AdvectiveForcing(on_architecture(to, af.u), on_architecture(to, af.v), on_architecture(to, af.w))
 
+@inline velocities(forcing::AdvectiveForcing) = (u=forcing.u, v=forcing.v, w=forcing.w)
 
 # fallback
 @inline with_advective_forcing(forcing, total_velocities) = total_velocities
 
 @inline with_advective_forcing(forcing::AdvectiveForcing, total_velocities) = 
-    (u = SumOfArrays{2}(forcing.u, total_velocities.u),
-     v = SumOfArrays{2}(forcing.v, total_velocities.v),
-     w = SumOfArrays{2}(forcing.w, total_velocities.w))
+    sum_of_velocities(velocities(forcing), total_velocities)
 
 # Unwrap the tuple within MultipleForcings
 @inline with_advective_forcing(mf::MultipleForcings, total_velocities) =

src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl

@@ -9,7 +9,7 @@ using KernelAbstractions: @index, @kernel
 using KernelAbstractions.Extras.LoopInfo: @unroll
 
 using Oceananigans.Utils
-using Oceananigans.Utils: launch!, SumOfArrays
+using Oceananigans.Utils: launch!
 using Oceananigans.Grids: AbstractGrid
 
 using DocStringExtensions

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl

@@ -7,8 +7,8 @@ using Oceananigans.Biogeochemistry: biogeochemical_transition, biogeochemical_dr
 using Oceananigans.TurbulenceClosures: immersed_∂ⱼ_τ₁ⱼ, immersed_∂ⱼ_τ₂ⱼ, immersed_∂ⱼ_τ₃ⱼ, immersed_∇_dot_qᶜ
 using Oceananigans.Advection: div_Uc, U_dot_∇u, U_dot_∇v
 using Oceananigans.Forcings: with_advective_forcing
-using Oceananigans.TurbulenceClosures: shear_production, buoyancy_flux, dissipation
-using Oceananigans.Utils: SumOfArrays
+using Oceananigans.TurbulenceClosures: shear_production, buoyancy_flux, dissipation, closure_turbulent_velocity
+using Oceananigans.Utils: sum_of_velocities
 using KernelAbstractions: @private
 
 import Oceananigans.TurbulenceClosures: hydrostatic_turbulent_kinetic_energy_tendency
@@ -123,11 +123,9 @@ where `c = C[tracer_index]`.
     model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields)
 
     biogeochemical_velocities = biogeochemical_drift_velocity(biogeochemistry, val_tracer_name)
+    closure_velocities = closure_turbulent_velocity(closure, diffusivities, val_tracer_name)
 
-    total_velocities = (u = SumOfArrays{2}(velocities.u, biogeochemical_velocities.u),
-                        v = SumOfArrays{2}(velocities.v, biogeochemical_velocities.v),
-                        w = SumOfArrays{2}(velocities.w, biogeochemical_velocities.w))
-
+    total_velocities = sum_of_velocities(velocities, biogeochemical_velocities, closure_velocities)
     total_velocities = with_advective_forcing(forcing, total_velocities)
 
     return ( - div_Uc(i, j, k, grid, advection, total_velocities, c)

src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl

@@ -8,7 +8,7 @@ using Oceananigans.TimeSteppers: tick!, step_lagrangian_particles!
 
 import Oceananigans.BoundaryConditions: fill_halo_regions!
 import Oceananigans.Models: extract_boundary_conditions
-import Oceananigans.Utils: datatuple
+import Oceananigans.Utils: datatuple, sum_of_velocities
 import Oceananigans.TimeSteppers: time_step!
 
 using Adapt
@@ -87,6 +87,15 @@ end
 @inline fill_halo_regions!(::FunctionField, args...) = nothing
 
 @inline datatuple(obj::PrescribedVelocityFields) = (; u = datatuple(obj.u), v = datatuple(obj.v), w = datatuple(obj.w))
+@inline velocities(obj::PrescribedVelocityFields) = (u = obj.u, v = obj.v, w = obj.w)
+
+# Extend sum_of_velocities for `PrescribedVelocityFields`
+@inline sum_of_velocities(U1::PrescribedVelocityFields, U2) = sum_of_velocities(velocities(U1), U2)
+@inline sum_of_velocities(U1, U2::PrescribedVelocityFields) = sum_of_velocities(U1, velocities(U2))
+
+@inline sum_of_velocities(U1::PrescribedVelocityFields, U2, U3) = sum_of_velocities(velocities(U1), U2, U3)
+@inline sum_of_velocities(U1, U2::PrescribedVelocityFields, U3) = sum_of_velocities(U1, velocities(U2), U3)
+@inline sum_of_velocities(U1, U2, U3::PrescribedVelocityFields) = sum_of_velocities(U1, U2, velocities(U3))
 
 ab2_step_velocities!(::PrescribedVelocityFields, args...) = nothing
 rk3_substep_velocities!(::PrescribedVelocityFields, args...) = nothing
@@ -132,4 +141,4 @@ function time_step!(model::OnlyParticleTrackingModel, Δt; callbacks = [], kwarg
 end
 
 update_state!(model::OnlyParticleTrackingModel, callbacks) =
-    [callback(model) for callback in callbacks if callback.callsite isa UpdateStateCallsite]
+    [callback(model) for callback in callbacks if callback.callsite isa UpdateStateCallsite]

src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

@@ -70,20 +70,19 @@ function mask_immersed_model_fields!(model, grid)
     return nothing
 end
 
-function compute_auxiliaries!(model::HydrostaticFreeSurfaceModel; w_parameters = tuple(w_kernel_parameters(model.grid)),
-                                                                  p_parameters = tuple(p_kernel_parameters(model.grid)),
-                                                                  κ_parameters = tuple(:xyz))
+function compute_auxiliaries!(model::HydrostaticFreeSurfaceModel; w_parameters = w_kernel_parameters(model.grid),
+                                                                  p_parameters = p_kernel_parameters(model.grid),
+                                                                  κ_parameters = :xyz)
 
     grid = model.grid
     closure = model.closure
     diffusivity = model.diffusivity_fields
 
-    for (wpar, ppar, κpar) in zip(w_parameters, p_parameters, κ_parameters)
-        compute_w_from_continuity!(model; parameters = wpar)
-        compute_diffusivities!(diffusivity, closure, model; parameters = κpar)
-        update_hydrostatic_pressure!(model.pressure.pHY′, architecture(grid),
-                                     grid, model.buoyancy, model.tracers; 
-                                     parameters = ppar)
-    end
+    compute_w_from_continuity!(model; parameters = w_parameters)    
+    compute_diffusivities!(diffusivity, closure, model; parameters = κ_parameters)
+    update_hydrostatic_pressure!(model.pressure.pHY′, architecture(grid),
+                                 grid, model.buoyancy, model.tracers; 
+                                 parameters = p_parameters)
+
     return nothing
 end

src/Models/LagrangianParticleTracking/LagrangianParticleTracking.jl

@@ -19,7 +19,7 @@ using Oceananigans.Architectures: device, architecture
 using Oceananigans.Fields: interpolate, datatuple, compute!, location
 using Oceananigans.Fields: fractional_indices
 using Oceananigans.TimeSteppers: AbstractLagrangianParticles
-using Oceananigans.Utils: prettysummary, launch!, SumOfArrays
+using Oceananigans.Utils: prettysummary, launch!
 
 import Oceananigans.TimeSteppers: step_lagrangian_particles!
 

src/Models/NonhydrostaticModels/NonhydrostaticModels.jl

@@ -16,7 +16,7 @@ using Oceananigans.DistributedComputations: DistributedFFTBasedPoissonSolver, Di
 using Oceananigans.Grids: XYRegularRG, XZRegularRG, YZRegularRG, XYZRegularRG
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
 using Oceananigans.Solvers: GridWithFFTSolver, GridWithFourierTridiagonalSolver 
-using Oceananigans.Utils: SumOfArrays
+using Oceananigans.Utils: sum_of_velocities
 
 import Oceananigans: fields, prognostic_fields
 import Oceananigans.Advection: cell_advection_timescale

src/Models/NonhydrostaticModels/nonhydrostatic_model.jl

@@ -256,7 +256,5 @@ end
 
 # return the total advective velocities
 @inline total_velocities(m::NonhydrostaticModel) =
-    (u = SumOfArrays{2}(m.velocities.u, m.background_fields.velocities.u),
-     v = SumOfArrays{2}(m.velocities.v, m.background_fields.velocities.v),
-     w = SumOfArrays{2}(m.velocities.w, m.background_fields.velocities.w))
+    sum_of_velocities(m.velocities, m.background_fields.velocities) 
 

src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl

@@ -62,10 +62,7 @@ pressure anomaly.
 
     model_fields = merge(velocities, tracers, auxiliary_fields)
 
-    total_velocities = (u = SumOfArrays{2}(velocities.u, background_fields.velocities.u),
-                        v = SumOfArrays{2}(velocities.v, background_fields.velocities.v),
-                        w = SumOfArrays{2}(velocities.w, background_fields.velocities.w))
-
+    total_velocities = sum_of_velocities(velocities, background_fields.velocities)
     total_velocities = with_advective_forcing(forcing, total_velocities)
 
     return ( - div_𝐯u(i, j, k, grid, advection, total_velocities, velocities.u)
@@ -125,10 +122,7 @@ pressure anomaly.
 
     model_fields = merge(velocities, tracers, auxiliary_fields)
 
-    total_velocities = (u = SumOfArrays{2}(velocities.u, background_fields.velocities.u),
-                        v = SumOfArrays{2}(velocities.v, background_fields.velocities.v),
-                        w = SumOfArrays{2}(velocities.w, background_fields.velocities.w))
-
+    total_velocities = sum_of_velocities(velocities, background_fields.velocities)
     total_velocities = with_advective_forcing(forcing, total_velocities)
 
     return ( - div_𝐯v(i, j, k, grid, advection, total_velocities, velocities.v)
@@ -191,10 +185,7 @@ velocity components, tracer fields, and precalculated diffusivities where applic
 
     model_fields = merge(velocities, tracers, auxiliary_fields)
 
-    total_velocities = (u = SumOfArrays{2}(velocities.u, background_fields.velocities.u),
-                        v = SumOfArrays{2}(velocities.v, background_fields.velocities.v),
-                        w = SumOfArrays{2}(velocities.w, background_fields.velocities.w))
-
+    total_velocities = sum_of_velocities(velocities, background_fields.velocities)
     total_velocities = with_advective_forcing(forcing, total_velocities)
 
     return ( - div_𝐯w(i, j, k, grid, advection, total_velocities, velocities.w)
@@ -256,10 +247,7 @@ velocity components, tracer fields, and precalculated diffusivities where applic
 
     biogeochemical_velocities = biogeochemical_drift_velocity(biogeochemistry, val_tracer_name)
 
-    total_velocities = (u = SumOfArrays{3}(velocities.u, background_fields.velocities.u, biogeochemical_velocities.u),
-                        v = SumOfArrays{3}(velocities.v, background_fields.velocities.v, biogeochemical_velocities.v),
-                        w = SumOfArrays{3}(velocities.w, background_fields.velocities.w, biogeochemical_velocities.w))
-
+    total_velocities = sum_of_velocities(velocities, background_fields.velocities, biogeochemical_velocities)
     total_velocities = with_advective_forcing(forcing, total_velocities)
 
     return ( - div_Uc(i, j, k, grid, advection, total_velocities, c)

src/TurbulenceClosures/TurbulenceClosures.jl

@@ -186,6 +186,7 @@ include("turbulence_closure_implementations/ri_based_vertical_diffusivity.jl")
 # Special non-abstracted diffusivities:
 # TODO: introduce abstract typing for these
 include("turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl")
+include("turbulence_closure_implementations/advective_skew_diffusion.jl")
 include("turbulence_closure_implementations/leith_enstrophy_diffusivity.jl")
 
 using .TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity, TKEDissipationVerticalDiffusivity

src/TurbulenceClosures/abstract_scalar_diffusivity_closure.jl

@@ -263,6 +263,8 @@ end
 @inline κ_σᶠᶜᶜ(i, j, k, grid, closure, K, id, clock, fields, σᶠᶜᶜ, args...) = κᶠᶜᶜ(i, j, k, grid, closure, K, id, clock, fields) * σᶠᶜᶜ(i, j, k, grid, args...)
 @inline κ_σᶜᶠᶜ(i, j, k, grid, closure, K, id, clock, fields, σᶜᶠᶜ, args...) = κᶜᶠᶜ(i, j, k, grid, closure, K, id, clock, fields) * σᶜᶠᶜ(i, j, k, grid, args...)
 @inline κ_σᶜᶜᶠ(i, j, k, grid, closure, K, id, clock, fields, σᶜᶜᶠ, args...) = κᶜᶜᶠ(i, j, k, grid, closure, K, id, clock, fields) * σᶜᶜᶠ(i, j, k, grid, args...)
+@inline κ_σᶠᶜᶠ(i, j, k, grid, closure, K, id, clock, fields, σᶠᶜᶠ, args...) = κᶠᶜᶠ(i, j, k, grid, closure, K, id, clock, fields) * σᶠᶜᶠ(i, j, k, grid, args...)
+@inline κ_σᶜᶠᶠ(i, j, k, grid, closure, K, id, clock, fields, σᶜᶠᶠ, args...) = κᶜᶠᶠ(i, j, k, grid, closure, K, id, clock, fields) * σᶜᶠᶠ(i, j, k, grid, args...)
 
 #####
 ##### Viscosity "extractors"
@@ -278,6 +280,9 @@ end
 @inline κᶠᶜᶜ(i, j, k, grid, loc, κ::Number, args...) = κ
 @inline κᶜᶠᶜ(i, j, k, grid, loc, κ::Number, args...) = κ
 @inline κᶜᶜᶠ(i, j, k, grid, loc, κ::Number, args...) = κ
+@inline κᶠᶜᶠ(i, j, k, grid, loc, κ::Number, args...) = κ
+@inline κᶜᶠᶠ(i, j, k, grid, loc, κ::Number, args...) = κ
+
 
 # Array / Field at `Center, Center, Center`
 const Lᶜᶜᶜ = Tuple{Center, Center, Center}
@@ -289,6 +294,8 @@ const Lᶜᶜᶜ = Tuple{Center, Center, Center}
 @inline κᶠᶜᶜ(i, j, k, grid, ::Lᶜᶜᶜ, κ::AbstractArray, args...) = ℑxᶠᵃᵃ(i, j, k, grid, κ)
 @inline κᶜᶠᶜ(i, j, k, grid, ::Lᶜᶜᶜ, κ::AbstractArray, args...) = ℑyᵃᶠᵃ(i, j, k, grid, κ)
 @inline κᶜᶜᶠ(i, j, k, grid, ::Lᶜᶜᶜ, κ::AbstractArray, args...) = ℑzᵃᵃᶠ(i, j, k, grid, κ)
+@inline κᶠᶜᶠ(i, j, k, grid, ::Lᶜᶜᶜ, κ::AbstractArray, args...) = ℑxzᶠᵃᶠ(i, j, k, grid, κ)
+@inline κᶜᶠᶠ(i, j, k, grid, ::Lᶜᶜᶜ, κ::AbstractArray, args...) = ℑyzᵃᶠᶠ(i, j, k, grid, κ)
 
 # Array / Field at `Center, Center, Face`
 const Lᶜᶜᶠ = Tuple{Center, Center, Face}
@@ -300,6 +307,8 @@ const Lᶜᶜᶠ = Tuple{Center, Center, Face}
 @inline κᶠᶜᶜ(i, j, k, grid, ::Lᶜᶜᶠ, κ::AbstractArray, args...) = ℑxzᶠᵃᶠ(i, j, k, grid, κ)
 @inline κᶜᶠᶜ(i, j, k, grid, ::Lᶜᶜᶠ, κ::AbstractArray, args...) = ℑyzᵃᶠᶠ(i, j, k, grid, κ)
 @inline κᶜᶜᶠ(i, j, k, grid, ::Lᶜᶜᶠ, κ::AbstractArray, args...) = @inbounds κ[i, j, k]
+@inline κᶠᶜᶠ(i, j, k, grid, ::Lᶜᶜᶠ, κ::AbstractArray, args...) = ℑxᶠᵃᵃ(i, j, k, grid, κ)
+@inline κᶜᶠᶠ(i, j, k, grid, ::Lᶜᶜᶠ, κ::AbstractArray, args...) = ℑyᵃᶠᵃ(i, j, k, grid, κ)
 
 # Function
 
@@ -314,6 +323,8 @@ const f = Face()
 @inline κᶠᶜᶜ(i, j, k, grid, loc, κ::F, clock, args...) where F<:Function = κ(node(i, j, k, grid, f, c, c)..., clock.time)
 @inline κᶜᶠᶜ(i, j, k, grid, loc, κ::F, clock, args...) where F<:Function = κ(node(i, j, k, grid, c, f, c)..., clock.time)
 @inline κᶜᶜᶠ(i, j, k, grid, loc, κ::F, clock, args...) where F<:Function = κ(node(i, j, k, grid, c, c, f)..., clock.time)
+@inline κᶠᶜᶠ(i, j, k, grid, loc, κ::F, clock, args...) where F<:Function = κ(node(i, j, k, grid, f, c, f)..., clock.time)
+@inline κᶜᶠᶠ(i, j, k, grid, loc, κ::F, clock, args...) where F<:Function = κ(node(i, j, k, grid, c, f, f)..., clock.time)
 
 # "DiscreteDiffusionFunction"
 @inline νᶜᶜᶜ(i, j, k, grid, loc, ν::DiscreteDiffusionFunction, clock, fields) = getdiffusivity(ν, i, j, k, grid, (c, c, c), clock, fields)
@@ -324,5 +335,8 @@ const f = Face()
 @inline κᶠᶜᶜ(i, j, k, grid, loc, κ::DiscreteDiffusionFunction, clock, fields) = getdiffusivity(κ, i, j, k, grid, (f, c, c), clock, fields)
 @inline κᶜᶠᶜ(i, j, k, grid, loc, κ::DiscreteDiffusionFunction, clock, fields) = getdiffusivity(κ, i, j, k, grid, (c, f, c), clock, fields)
 @inline κᶜᶜᶠ(i, j, k, grid, loc, κ::DiscreteDiffusionFunction, clock, fields) = getdiffusivity(κ, i, j, k, grid, (c, c, f), clock, fields)
+@inline κᶠᶜᶠ(i, j, k, grid, loc, κ::DiscreteDiffusionFunction, clock, fields) = getdiffusivity(κ, i, j, k, grid, (f, c, f), clock, fields)
+@inline κᶜᶠᶠ(i, j, k, grid, loc, κ::DiscreteDiffusionFunction, clock, fields) = getdiffusivity(κ, i, j, k, grid, (c, f, f), clock, fields)
+