Make the 3D stokes drift validation case even better (#3578)

* Make the 3D stokes drift validation case even better

* Better comments

* Apply suggestions from code review

---------

Co-authored-by: Navid C. Constantinou <[email protected]>

validation/stokes_drift/surface_wave_quasi_geostrophic_flow.jl

@@ -2,7 +2,7 @@ using Oceananigans
 using Oceananigans.Units
 using GLMakie
 
-ϵ = 0.1
+ϵ = 0.3
 λ = 100 # meters
 g = 9.81
 
@@ -21,27 +21,31 @@ const Uˢ = ϵ^2 * c
 
 # Write the Stokes drift as
 #
-# uˢ(x, z, t) = A(x, t) * ûˢ(z)
+# uˢ(x, y, z, t) = A(x, y, t) * ûˢ(z)
 #
 # which implies
 
-@inline    ûˢ(z)       = Uˢ * exp(2k * z)
+@inline    ûˢ(z)          = Uˢ * exp(2k * z)
 @inline    uˢ(x, y, z, t) =           A(x - cᵍ * t, y) * ûˢ(z)
 @inline ∂z_uˢ(x, y, z, t) =     2k *  A(x - cᵍ * t, y) * ûˢ(z)
 @inline ∂y_uˢ(x, y, z, t) =        ∂η_A(x - cᵍ * t, y) * ûˢ(z)
 @inline ∂t_uˢ(x, y, z, t) = - cᵍ * ∂ξ_A(x - cᵍ * t, y) * ûˢ(z)
 
+# where we have noted that η = y, and ξ = x - cᵍ t,
+# such that ∂ξ/∂x = 1 and ∂ξ/∂t = - cᵍ.
+#
 # Note that if uˢ represents the solenoidal component of the Stokes drift,
 # then
 #
 # ```math
-# ∂z_wˢ = - ∂x_uˢ = - ∂ξ_A * ûˢ .
+# ∂z_wˢ = - ∂x_uˢ = - ∂ξ_A * ∂ξ/∂x * ûˢ .
+#                 = - ∂ξ_A * ûˢ .
 # ```
 #
 # We therefore find that
 #
 # ```math
-# wˢ = - ∂ξ_A / 2k * ûˢ
+# wˢ = - ∂ξ_A / 2k * ûˢ .
 # ```
 #
 # and
@@ -52,8 +56,8 @@ const Uˢ = ϵ^2 * c
 
 stokes_drift = StokesDrift(; ∂z_uˢ, ∂t_uˢ, ∂y_uˢ, ∂t_wˢ, ∂x_wˢ, ∂y_wˢ)
 
-grid = RectilinearGrid(size = (64, 64, 16),
-                       x = (-5δ, 15δ),
+grid = RectilinearGrid(size = (128, 64, 16),
+                       x = (-10δ, 30δ),
                        y = (-10δ, 10δ),
                        z = (-512, 0),
                        topology = (Periodic, Periodic, Bounded))
@@ -81,12 +85,18 @@ simulation.callbacks[:progress] = Callback(progress, IterationInterval(10))
 
 filename = "surface_wave_quasi_geostrophic_induced_flow.jld2"
 outputs = model.velocities
+
+u, v, w = model.velocities
+e = @at (Center, Center, Center) (u^2 + v^2 + w^2) / 2
+outputs = merge(outputs, (; e))
+
 simulation.output_writers[:jld2] = JLD2OutputWriter(model, outputs; filename,
                                                     schedule = IterationInterval(3),
                                                     overwrite_existing = true)
 
 run!(simulation)
 
+et = FieldTimeSeries(filename, "e")
 ut = FieldTimeSeries(filename, "u")
 wt = FieldTimeSeries(filename, "w")
 
@@ -98,18 +108,57 @@ n = Observable(1)
 
 un = @lift interior(ut[$n], :, :, Nz)
 wn = @lift interior(wt[$n], :, :, Nz)
+en = @lift interior(et[$n], :, :, Nz)
 
 xu, yu, zu = nodes(ut)
 xw, yw, zw = nodes(wt)
+xe, ye, ze = nodes(et)
+
+Nx, Ny, Nz = size(grid)
+xa, ya, za = nodes(grid, Center(), Center(), Center())
+xc, yc, zc = nodes(grid, Center(), Center(), Center())
+xa = reshape(xa, Nx, 1)
+ya = reshape(ya, 1, Ny)
+
+An = @lift begin
+    t = times[$n]
+    ξ = @. xa - cᵍ * t
+    A.(ξ, ya)
+end
+
+# convert m -> km
+xu = xu .* 1e-3
+yu = yu .* 1e-3
+zu = zu .* 1e-3
+xw = xw .* 1e-3
+yw = yw .* 1e-3
+zw = zw .* 1e-3
+xe = xe .* 1e-3
+ye = ye .* 1e-3
+ze = ze .* 1e-3
 
-fig = Figure(size=(800, 400))
+xc = xc .* 1e-3
+yc = yc .* 1e-3
+zc = zc .* 1e-3
 
-axu = Axis(fig[1, 1], xlabel="x (m)", ylabel="z (m)", aspect=1)
-axw = Axis(fig[1, 2], xlabel="x (m)", ylabel="z (m)", aspect=1)
+fig = Figure(size=(1600, 400))
 
-heatmap!(axu, xu, yu, un)
-heatmap!(axw, xw, yw, wn)
+axu = Axis(fig[1, 1], xlabel="x (km)", ylabel="z (km)", aspect=2, title="u")
+axw = Axis(fig[1, 2], xlabel="x (km)", ylabel="z (km)", aspect=2, title="w")
+
+ulim = 1e-2 * Uˢ
+elim = ulim^2
+wlim = 1e-1 * Uˢ / (δ * 2k)
+
+heatmap!(axu, xu, yu, un, colormap=:balance, colorrange=(-ulim, ulim))
+#heatmap!(axu, xe, ye, en, colormap=:solar, colorrange=(0, elim))
+contour!(axu, xc, yc, An, color=:gray, levels=5)
+
+heatmap!(axw, xw, yw, wn, colormap=:balance, colorrange=(-wlim, wlim))
+contour!(axw, xc, yc, An, color=:gray, levels=5)
 
 record(fig, "surface_wave_quasi_geostrophic_induced_flow.mp4", 1:Nt, framerate=8) do nn
+record(fig, "surface_wave_non_rotating_induced_flow.mp4", 1:Nt, framerate=8) do nn
     n[] = nn
 end
+