(0.90.4) Implement three-dimensional StokesDrift 🌊

[glwagner]

This PR implements a StokesDrift abstraction in which the Stokes shear and tendency terms may be functions of x, y, z, t rather than only z, t as in UniformStokesDrift.

I'll post some results from a modified version of the langmuir_instability.jl example here when we have them.

Could be applicable to simulations of wave-affected turbulence in sea ice leads.

With @BrodiePearson and Ara Lee.

[glwagner]

@LeeAra0

[glwagner]

@BrodiePearson posted the following on #3392:

I also added an example code that utilizes the new StokesDrift function (and retains the previous UniformStokesDrift function). It is a little convoluted because I wanted to implement a Stokes drift that varied in all 3 dimensions to ensure that the code was well-behaved. The example is as follows:

• Is based on the Langmuir turbulence example
• The new example imposes a spatially-varying x-aligned Stokes drift that is a separable function of the three spatial dimensions

$$u^s(x, y, z) = f(x)g'(y)h'(z)$$

• The vertical variation is a classic exponential decay with depth over a vertical length scale $\delta$. Note this function is defined as the z-derivative of another function $h(z)$ to make the description of the vertical Stokes drift easier below.

$$h'(z) = U^s e^{(z/\delta)}$$

• The x-variation is sinusoidal modulation on top of a background that peaks in the middle of the x-domain.

$$f(x) = \frac{1}{2}\left( 1 + 0.1 \cos\left(\frac{\pi(2x-L_x)}{L_x}\right)\right)$$

• The y-variation represents a piece-wise function consisting of a jet region of width $L_{jet}$ centered at $y_{jet}$ where $u^s$ is constant, transition regions of width $L_{transition}$ on either side of the jet where Stokes drift follows a sinusoidal decay from its jet value to zero, and a no wave region where there is no Stokes drift outside of the jet and transition zones.

$$g(y) = 1 + \cos\left(\pi \frac{y - y_{jet} + L_{jet}/2 }{L_{transition}} \right) \quad \text{if } y_{jet} - L_{jet}/2 > y > y_{jet} - L_{jet}/2 - L_{transition}$$ $$g(y) = 1 + \cos\left(\pi \frac{y - y_{jet} - L_{jet}/2 }{L_{transition}} \right) \quad \text{if } y_{jet} + L_{jet}/2 + L_{transition} > y > y_{jet} + L_{jet}/2$$ $$g(y) = 0 \quad \text{at all other } y$$

• This x-aligned Stokes drift has $\partial_x u^s\neq0$ so incompressibility ($\nabla \cdot \mathbf{u^s}$) requires another Stokes drift component. Since we do not impose a y-oriented wave field, we achieve this through a vertical Stokes drift component $w^s(x,y,z)$. We diagnose this from the $u^s$ field and by imposing a surface boundary condition of no vertical motion: $w^s(x,y,0)=0$ . It then follows that:

$$w^s(x, y, z) = f'(x)g(y) \left[h(z)-h(0) \right]$$

[BrodiePearson]

Correction to my last message, I accidentally included a second prime. The form of the Stokes drift is actually:

$u^s(x,y,z)=f(x)g(y)h′(z)$

[glwagner]

And I want to add that we can only use $\partial_z w^S = - \partial_x u^S - \partial_y v^S$ if $(u^S, v^S, w^S)$ is the solenoidal component of the Stokes drift. In other words, the typical definition of Stokes drift contains a divergent component. However, it is possible to compute only the solenoidal component of the Stokes drift; moreover in that case it enters the Craik-Leibovich equations in the same way as the ordinary Stokes drift. This is the Stokes drift that should be used with Oceananigans, because our pressure solver enforces $\nabla \cdot u^L = 0$ --- something that is only true if the Stokes drift is also solenoidal.

(I guess, users in fact do not have a choice but to use a solenoidal Stokes drift due to the way our code is written...)

As for an example, it could be nice to come up with an example that illustrates the flow associated with a traveling wave packet. It should look like a dipole.

[glwagner]

@BrodiePearson I refactored StokesDrifts.jl very slightly --- I agree with the new terms you added!

I put together this script that reproduces the "deep Eulerian return flow" from McIntyre (1981):

using Oceananigans
using Oceananigans.Units
using GLMakie

ϵ = 0.1
λ = 60 # meters
g = 9.81

const k = 2π / λ

c = sqrt(g / k)
const δ = 1kilometer
const cᵍ = c / 2
const Uˢ = ϵ^2 * c

@inline A(ξ) = exp(- ξ^2 / 2δ^2)
@inline A′(ξ) = - ξ / δ^2 * A(ξ)
@inline A′′(ξ) = (ξ^2 / δ^2 - 1) * A(ξ) / δ^2

# Write the Stokes drift as
#
# uˢ(x, z, t) = A(x, t) * ûˢ(z)
#
# which implies

@inline    ûˢ(z)       = Uˢ * exp(2k * z)
@inline    uˢ(x, z, t) =         A(x - cᵍ * t) * ûˢ(z)
@inline ∂z_uˢ(x, z, t) =   2k *  A(x - cᵍ * t) * ûˢ(z)
@inline ∂t_uˢ(x, z, t) = - cᵍ * A′(x - cᵍ * t) * ûˢ(z)

# Note that if uˢ represents the solenoidal component of the Stokes drift,
# then
#
# ```math
# ∂z_wˢ = - ∂x_uˢ = - A′ * ûˢ .
# ```
#
# We therefore find that
#
# ```math
# wˢ = - A′ / 2k * ûˢ
# ```
#
# and

@inline ∂x_wˢ(x, z, t) = -  1 / 2k * A′′(x - cᵍ * t) * ûˢ(z)
@inline ∂t_wˢ(x, z, t) = + cᵍ / 2k * A′′(x - cᵍ * t) * ûˢ(z)

stokes_drift = StokesDrift(; ∂z_uˢ, ∂t_uˢ, ∂t_wˢ, ∂x_wˢ)

grid = RectilinearGrid(size = (256, 64),
                       x = (-5kilometers, 15kilometers),
                       z = (-512, 0),
                       topology = (Periodic, Flat, Bounded))

model = NonhydrostaticModel(; grid, stokes_drift,
                            tracers = :b,
                            buoyancy = BuoyancyTracer(),
                            timestepper = :RungeKutta3)

# Set Lagrangian-mean flow equal to uˢ,
uᵢ(x, z) = uˢ(x, z, 0)

# And put in a stable stratification,
N² = 0
bᵢ(x, z) = N² * z
set!(model, u=uᵢ, b=bᵢ)

Δx = xspacings(grid, Center())
Δt = 0.2 * Δx / cᵍ
simulation = Simulation(model; Δt, stop_iteration = 600)

progress(sim) = @info string("Iter: ", iteration(sim), ", time: ", prettytime(sim))
simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))

filename = "surface_wave_induced_flow.jld2"
outputs = model.velocities
simulation.output_writers[:jld2] = JLD2OutputWriter(model, outputs; filename,
                                                    schedule = IterationInterval(10),
                                                    overwrite_existing = true)

run!(simulation)

ut = FieldTimeSeries(filename, "u")
wt = FieldTimeSeries(filename, "w")

times = ut.times
Nt = length(times)

n = Observable(1)

un = @lift interior(ut[$n], :, 1, :)
wn = @lift interior(wt[$n], :, 1, :)

xu, yu, zu = nodes(ut)
xw, yw, zw = nodes(wt)

fig = Figure(resolution=(800, 300))

axu = Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)")
axw = Axis(fig[1, 2], xlabel="x (m)", ylabel="z (m)")

heatmap!(axu, xu, zu, un)
heatmap!(axw, xw, zw, wn)

record(fig, "surface_wave_induced_flow.mp4", 1:Nt, framerate=12) do nn
    n[] = nn
end

The result is

surface_wave_induced_flow.mp4

where the left panel is u and the right panel is w.

@BrodiePearson I don't think we should add a new example for this feature (examples are expensive, because they have to run every time we run CI / build the documentation). However, another avenue to keep some code around is to add a "validation" case. If you're up for that, I'll move your code there, along with the above example. Let me know and then we can possibly merge this great addition.

[glwagner]

@BrodiePearson I've also added a docstring for StokesDrift. Check it out and let me know what you think.

[BrodiePearson]

@glwagner These changes look great!

I fixed some typos in the refactored file, which were terms that would not have affected your new validation case. I moved my original example to validation and provided some minor tweaks to your first validation case. I did not look much at the surface_wave_quasi_geostrophic_flow.jl case you created

[glwagner]

nice catch!

[glwagner]

maybe we should add two validation tests, one for stokes drift in u, w and one for v, w.

[BrodiePearson]

I just added a Langmuir validation test, forced by a Stokes drift jet in the y-direction, which is simply a rotated version of the validation test forced by an x-oriented Stokes drift jet. The resulting velocities and Reynolds stress components mirror each other as expected.

x-oriented Stokes drift jet

y-oriented Stokes drift jet

[BrodiePearson]

If you prefer, it would be more simple to switch your 2D validation case to be in the y-z direction, but the Langmuir cases together utilize all horizontal and vertical derivative terms in StokesDrifts.jl.

[glwagner]

Ok great then I think this is close. I'll just add a Project.toml to the new validation directory (hopefully eventually we will transition all validation directories to this more maintainable state).

[navidcy]

@glwagner merge at will, I won't do anything else

[glwagner]

I think the example for 3D Stokes drift can be improved to be more realistic. Note that it is confusing to use cos(k * x), because in that case the k is the wavelength of the envelope, where typically we use k as the wavelength of the underlying surface wave (which provides the vertical scale of the Stokes drift, but not the horizontal scale). This can be something for the future.