Skip to content

Commit 0e06f28

Browse files
glwagnernavidcy
glwagner
and
navidcy
authored
Field Stokes shear and tendency with UniformStokesDrift (#2320)
* Adds new UniformStokesDrift constructor for field-based stokes drift * Fix bug and add a docstring * Adds a docstring * adds show(io::IO, ::UniformStokesDrift) + fix doctest * Bump version and add two tests * Makes show consistent with other obj, add to docstring * fix docstring examples * Update StokesDrift.jl * Update test_time_stepping.jl * Disambiguate --------- Co-authored-by: Navid C. Constantinou <[email protected]>
1 parent 949d0c7 commit 0e06f28

File tree

3 files changed

+68
-18
lines changed

3 files changed

+68
-18
lines changed

‎src/StokesDrifts.jl

Lines changed: 47 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -125,28 +125,57 @@ UniformStokesDrift with parameters (uˢ=0.005, h=20):
125125
UniformStokesDrift(; ∂z_uˢ=zerofunction, ∂z_vˢ=zerofunction, ∂t_uˢ=zerofunction, ∂t_vˢ=zerofunction, parameters=nothing) =
126126
UniformStokesDrift(∂z_uˢ, ∂z_vˢ, ∂t_uˢ, ∂t_vˢ, parameters)
127127

128+
function UniformStokesDrift(grid::AbstractGrid;
129+
∂z_uˢ = Field{Nothing, Nothing, Face}(grid),
130+
∂z_vˢ = Field{Nothing, Nothing, Face}(grid),
131+
∂t_uˢ = Field{Nothing, Nothing, Center}(grid),
132+
∂t_vˢ = Field{Nothing, Nothing, Center}(grid),
133+
parameters = nothing)
134+
135+
return UniformStokesDrift(∂z_uˢ, ∂z_vˢ, ∂t_uˢ, ∂t_vˢ, parameters)
136+
end
137+
128138
const USD = UniformStokesDrift
129139
const USDnoP = UniformStokesDrift{<:Nothing}
130-
const f = Face()
131140
const c = Center()
141+
const f = Face()
132142

133-
@inline ∂t_uˢ(i, j, k, grid, sw::USD, time) = sw.∂t_uˢ(znode(k, grid, c), time, sw.parameters)
134-
@inline ∂t_vˢ(i, j, k, grid, sw::USD, time) = sw.∂t_vˢ(znode(k, grid, c), time, sw.parameters)
135-
@inline ∂t_wˢ(i, j, k, grid, sw::USD, time) = zero(grid)
136-
137-
@inline x_curl_Uˢ_cross_U(i, j, k, grid, sw::USD, U, time) = ℑxzᶠᵃᶜ(i, j, k, grid, U.w) * sw.∂z_uˢ(znode(k, grid, c), time, sw.parameters)
138-
@inline y_curl_Uˢ_cross_U(i, j, k, grid, sw::USD, U, time) = ℑyzᵃᶠᶜ(i, j, k, grid, U.w) * sw.∂z_vˢ(znode(k, grid, c), time, sw.parameters)
139-
@inline z_curl_Uˢ_cross_U(i, j, k, grid, sw::USD, U, time) = (- ℑxzᶜᵃᶠ(i, j, k, grid, U.u) * sw.∂z_uˢ(znode(k, grid, f), time, sw.parameters)
140-
- ℑyzᵃᶜᶠ(i, j, k, grid, U.v) * sw.∂z_vˢ(znode(k, grid, f), time, sw.parameters))
141-
142-
# Methods for when `parameters == nothing`
143-
@inline ∂t_uˢ(i, j, k, grid, sw::USDnoP, time) = sw.∂t_uˢ(znode(k, grid, c), time)
144-
@inline ∂t_vˢ(i, j, k, grid, sw::USDnoP, time) = sw.∂t_vˢ(znode(k, grid, c), time)
145-
146-
@inline x_curl_Uˢ_cross_U(i, j, k, grid, sw::USDnoP, U, time) = ℑxzᶠᵃᶜ(i, j, k, grid, U.w) * sw.∂z_uˢ(znode(k, grid, c), time)
147-
@inline y_curl_Uˢ_cross_U(i, j, k, grid, sw::USDnoP, U, time) = ℑyzᵃᶠᶜ(i, j, k, grid, U.w) * sw.∂z_vˢ(znode(k, grid, c), time)
148-
@inline z_curl_Uˢ_cross_U(i, j, k, grid, sw::USDnoP, U, time) = (- ℑxzᶜᵃᶠ(i, j, k, grid, U.u) * sw.∂z_uˢ(znode(k, grid, f), time)
149-
- ℑyzᵃᶜᶠ(i, j, k, grid, U.v) * sw.∂z_vˢ(znode(k, grid, f), time))
143+
# Some helpers for three cases: Nothing, AbstractArray, or fallback (function)
144+
@inline ∂z_Uᵃᵃᶜ(i, j, k, grid, sd::USDnoP, ∂z_Uˢ, time) = ∂z_Uˢ(znode(k, grid, c), time)
145+
@inline ∂z_Uᵃᵃᶜ(i, j, k, grid, sd::USD, ∂z_Uˢ, time) = ∂z_Uˢ(znode(k, grid, c), time, sd.parameters)
146+
@inline ∂z_Uᵃᵃᶜ(i, j, k, grid, sd::USD, ∂z_Uˢ::AbstractArray, time) = ℑzᵃᵃᶜ(i, j, k, grid, ∂z_Uˢ)
147+
@inline ∂z_Uᵃᵃᶜ(i, j, k, grid, sd::USDnoP, ∂z_Uˢ::AbstractArray, time) = ℑzᵃᵃᶜ(i, j, k, grid, ∂z_Uˢ)
148+
@inline ∂z_Uᵃᵃᶜ(i, j, k, grid, sd::USD, ::Nothing, time) = zero(grid)
149+
@inline ∂z_Uᵃᵃᶜ(i, j, k, grid, sd::USDnoP, ::Nothing, time) = zero(grid)
150+
151+
@inline ∂z_Uᵃᵃᶠ(i, j, k, grid, sd::USDnoP, ∂z_Uˢ, time) = ∂z_Uˢ(znode(k, grid, f), time)
152+
@inline ∂z_Uᵃᵃᶠ(i, j, k, grid, sd::USD, ∂z_Uˢ, time) = ∂z_Uˢ(znode(k, grid, f), time, sd.parameters)
153+
@inline ∂z_Uᵃᵃᶠ(i, j, k, grid, sd::USD, ∂z_Uˢ::AbstractArray, time) = @inbounds ∂z_Uˢ[i, j, k]
154+
@inline ∂z_Uᵃᵃᶠ(i, j, k, grid, sd::USDnoP, ∂z_Uˢ::AbstractArray, time) = @inbounds ∂z_Uˢ[i, j, k]
155+
@inline ∂z_Uᵃᵃᶠ(i, j, k, grid, sd::USD, ::Nothing, time) = zero(grid)
156+
@inline ∂z_Uᵃᵃᶠ(i, j, k, grid, sd::USDnoP, ::Nothing, time) = zero(grid)
157+
158+
@inline ∂t_U(i, j, k, grid, sd::USDnoP, ∂t_Uˢ, time) = ∂t_Uˢ(znode(k, grid, c), time)
159+
@inline ∂t_U(i, j, k, grid, sd::USD, ∂t_Uˢ, time) = ∂t_Uˢ(znode(k, grid, c), time, sd.parameters)
160+
@inline ∂t_U(i, j, k, grid, sd::USD, ∂t_Uˢ::AbstractArray, time) = @inbounds ∂t_Uˢ[i, j, k]
161+
@inline ∂t_U(i, j, k, grid, sd::USDnoP, ∂t_Uˢ::AbstractArray, time) = @inbounds ∂t_Uˢ[i, j, k]
162+
@inline ∂t_U(i, j, k, grid, sd::USD, ::Nothing, time) = zero(grid)
163+
@inline ∂t_U(i, j, k, grid, sd::USDnoP, ::Nothing, time) = zero(grid)
164+
165+
# Kernel functions
166+
@inline ∂t_uˢ(i, j, k, grid, sd::USD, time) = ∂t_U(i, j, k, grid, sd, sd.∂t_uˢ, time)
167+
@inline ∂t_vˢ(i, j, k, grid, sd::USD, time) = ∂t_U(i, j, k, grid, sd, sd.∂t_vˢ, time)
168+
@inline ∂t_wˢ(i, j, k, grid, sd::USD, time) = zero(grid)
169+
170+
@inline x_curl_Uˢ_cross_U(i, j, k, grid, sd::USD, U, time) =
171+
ℑxzᶠᵃᶜ(i, j, k, grid, U.w) * ∂z_Uᵃᵃᶜ(i, j, k, grid, sd, sd.∂z_uˢ, time)
172+
173+
@inline y_curl_Uˢ_cross_U(i, j, k, grid, sd::USD, U, time) =
174+
ℑyzᵃᶠᶜ(i, j, k, grid, U.w) * ∂z_Uᵃᵃᶜ(i, j, k, grid, sd, sd.∂z_vˢ, time)
175+
176+
@inline z_curl_Uˢ_cross_U(i, j, k, grid, sd::USD, U, time) = (
177+
- ℑxzᶜᵃᶠ(i, j, k, grid, U.u) * ∂z_Uᵃᵃᶠ(i, j, k, grid, sd, sd.∂z_uˢ, time)
178+
- ℑyzᵃᶜᶠ(i, j, k, grid, U.v) * ∂z_Uᵃᵃᶠ(i, j, k, grid, sd, sd.∂z_vˢ, time))
150179

151180
struct StokesDrift{P, VX, WX, UY, WY, UZ, VZ, UT, VT, WT}
152181
∂x_vˢ :: VX

‎test/test_stokes_drift.jl

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,5 +46,15 @@ end
4646
@testset "Stokes drift" begin
4747
@test instantiate_uniform_stokes_drift()
4848
@test instantiate_stokes_drift()
49+
50+
for arch in archs
51+
grid = RectilinearGrid(arch, size=(3, 3, 3), extent=(1, 1, 1))
52+
stokes_drift = UniformStokesDrift(grid)
53+
@test location(stokes_drift.∂z_uˢ) === (Nothing, Nothing, Face)
54+
@test location(stokes_drift.∂z_vˢ) === (Nothing, Nothing, Face)
55+
@test location(stokes_drift.∂t_uˢ) === (Nothing, Nothing, Center)
56+
@test location(stokes_drift.∂t_vˢ) === (Nothing, Nothing, Center)
57+
end
4958
end
5059
end
60+

‎test/test_time_stepping.jl

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -402,6 +402,17 @@ timesteppers = (:QuasiAdamsBashforth2, :RungeKutta3)
402402
end
403403
end
404404

405+
@testset "UniformStokesDrift" begin
406+
for arch in archs
407+
grid = RectilinearGrid(arch, size=(1, 1, 1), extent=(1, 1, 1))
408+
# Cover three cases:
409+
stokes_drift = UniformStokesDrift(grid, ∂z_vˢ=nothing, ∂t_uˢ= (z, t) -> exp(z/20))
410+
model = NonhydrostaticModel(; grid, stokes_drift)
411+
time_step!(model, 1)
412+
@test true
413+
end
414+
end
415+
405416
@testset "Idealized nonlinear equation of state" begin
406417
for arch in archs, FT in [Float64]
407418
for eos_type in (SeawaterPolynomials.RoquetEquationOfState, SeawaterPolynomials.TEOS10EquationOfState)

0 commit comments

Comments
 (0)