-
Notifications
You must be signed in to change notification settings - Fork 191
/
shallow_water_diffusion_operators.jl
87 lines (65 loc) · 4.32 KB
/
shallow_water_diffusion_operators.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
using Oceananigans.Operators
using Oceananigans.Architectures: device, device_event
using Oceananigans.TurbulenceClosures: ExplicitTimeDiscretization, ThreeDimensionalFormulation
using Oceananigans.TurbulenceClosures:
AbstractScalarDiffusivity,
convert_diffusivity,
viscosity_location,
viscosity,
ν_σᶜᶜᶜ,
∂ⱼ_τ₁ⱼ,
∂ⱼ_τ₂ⱼ
import Oceananigans.TurbulenceClosures:
DiffusivityFields,
calculate_diffusivities!,
calculate_nonlinear_viscosity!,
viscosity,
with_tracers,
calc_nonlinear_νᶜᶜᶜ,
νᶜᶜᶜ
struct ShallowWaterScalarDiffusivity{N, X} <: AbstractScalarDiffusivity{ExplicitTimeDiscretization, ThreeDimensionalFormulation}
ν :: N
ξ :: X
end
function ShallowWaterScalarDiffusivity(FT::DataType=Float64; ν=0, ξ=0, discrete_form = false)
ν = convert_diffusivity(FT, ν; discrete_form)
ξ = convert_diffusivity(FT, ξ; discrete_form)
return ShallowWaterScalarDiffusivity(ν, ξ)
end
# We have no tracers in the shallow water diffusivity
with_tracers(tracers, closure::ShallowWaterScalarDiffusivity) = closure
viscosity(closure::ShallowWaterScalarDiffusivity, K) = closure.ν
Adapt.adapt_structure(to, closure::ShallowWaterScalarDiffusivity) =
ShallowWaterScalarDiffusivity(Adapt.adapt(to, closure.ν), Adapt.adapt(to, closure.ξ))
# The diffusivity for the shallow water model is calculated as h*ν in order to have a viscous term in the form
# h⁻¹ ∇ ⋅ (hν t) where t is the 2D stress tensor plus a trace => t = ∇u + (∇u)ᵀ - ξI⋅(∇⋅u)
@inline calc_nonlinear_νᶜᶜᶜ(i, j, k, grid, closure::ShallowWaterScalarDiffusivity, clock, fields) =
fields.h[i, j, k] * νᶜᶜᶜ(i, j, k, grid, viscosity_location(closure), closure.ν, clock, fields)
function calculate_diffusivities!(diffusivity_fields, closure::ShallowWaterScalarDiffusivity, model)
arch = model.architecture
grid = model.grid
clock = model.clock
model_fields = shallow_water_fields(model.velocities, model.tracers, model.solution, formulation(model))
event = launch!(arch, grid, :xyz,
calculate_nonlinear_viscosity!,
diffusivity_fields.νₑ, grid, closure, clock, model_fields,
dependencies = device_event(arch))
wait(device(arch), event)
return nothing
end
DiffusivityFields(grid, tracer_names, bcs, ::ShallowWaterScalarDiffusivity) = (; νₑ=CenterField(grid, boundary_conditions=bcs.h))
#####
##### Diffusion flux divergence operators
#####
@inline sw_∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, K, clock, fields, ::ConservativeFormulation) =
∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, K, clock, fields, nothing) + trace_term_x(i, j, k, grid, closure, K, clock, fields)
@inline sw_∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, K, clock, fields, ::ConservativeFormulation) =
∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, K, clock, fields, nothing) + trace_term_y(i, j, k, grid, closure, K, clock, fields)
@inline sw_∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, K, clock, fields, ::VectorInvariantFormulation) =
(∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, K, clock, fields, nothing) + trace_term_x(i, j, k, grid, closure, K, clock, fields) ) / ℑxᶠᵃᵃ(i, j, k, grid, fields.h)
@inline sw_∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, K, clock, fields, ::VectorInvariantFormulation) =
(∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, K, clock, fields, nothing) + trace_term_y(i, j, k, grid, closure, K, clock, fields) ) / ℑyᵃᶠᵃ(i, j, k, grid, fields.h)
@inline trace_term_x(i, j, k, grid, clo, K, clk, fields) = - δxᶠᵃᵃ(i, j, k, grid, ν_σᶜᶜᶜ, clo, K, clk, fields, div_xyᶜᶜᶜ, fields.u, fields.v) * clo.ξ / Azᶠᶜᶜ(i, j, k, grid)
@inline trace_term_y(i, j, k, grid, clo, K, clk, fields) = - δyᵃᶠᵃ(i, j, k, grid, ν_σᶜᶜᶜ, clo, K, clk, fields, div_xyᶜᶜᶜ, fields.u, fields.v) * clo.ξ / Azᶠᶜᶜ(i, j, k, grid)
@inline trace_term_x(i, j, k, grid, ::Nothing, args...) = zero(grid)
@inline trace_term_y(i, j, k, grid, ::Nothing, args...) = zero(grid)