From 0fc5708f0c4dd953f2433b8db2ff22f2dbab2639 Mon Sep 17 00:00:00 2001
From: Gregory Wagner <wagner.greg@gmail.com>
Date: Wed, 18 Dec 2024 15:58:39 -0700
Subject: [PATCH] Allow divergent Lagrangian-mean velocity by subtracting
 Stokes drift before pressure correction

---
 .../pressure_correction.jl                    | 42 +++++++++++++++++--
 src/StokesDrifts.jl                           | 20 ++++++---
 2 files changed, 54 insertions(+), 8 deletions(-)

diff --git a/src/Models/NonhydrostaticModels/pressure_correction.jl b/src/Models/NonhydrostaticModels/pressure_correction.jl
index c657995799..da0297b340 100644
--- a/src/Models/NonhydrostaticModels/pressure_correction.jl
+++ b/src/Models/NonhydrostaticModels/pressure_correction.jl
@@ -1,5 +1,40 @@
 import Oceananigans.TimeSteppers: calculate_pressure_correction!, pressure_correct_velocities!
 
+const c = Center()
+const f = Face()
+
+using Oceananigans.Grids: node
+using Oceananigans.StokesDrifts: StokesDrift, parameters_tuple
+
+function modulate_by_stokes_drift!(model, sgn, t=time(model))
+    stokes_drift = model.stokes_drift
+    if stokes_drift isa StokesDrift
+        grid = model.grid
+        arch = architecture(grid)
+        u, v, w = model.velocities
+        launch!(arch, grid, :xyz, _modulate_by_stokes_drift, u, v, w, sgn, grid, t, stokes_drift)
+    end
+    return nothing
+end
+
+@kernel function _modulate_by_stokes_drift(u, v, w, sgn, grid, time, sd)
+    i, j, k = @index(Global, NTuple)
+
+    pt = parameters_tuple(sd)
+    Xu = node(i, j, k, grid, f, c, c)
+    Xv = node(i, j, k, grid, c, f, c)
+    Xw = node(i, j, k, grid, c, c, f)
+
+    @inbounds begin
+        u[i, j, k] = u[i, j, k] + sgn * sd.uˢ(Xu..., time, pt...)
+        v[i, j, k] = v[i, j, k] + sgn * sd.vˢ(Xv..., time, pt...)
+        w[i, j, k] = w[i, j, k] + sgn * sd.wˢ(Xw..., time, pt...)
+    end
+end
+
+subtract_stokes_drift!(model, t=time(model)) = modulate_by_stokes_drift!(model, -1, t)
+add_stokes_drift!(model, t=time(model)) = modulate_by_stokes_drift!(model, +1, t)
+
 """
     calculate_pressure_correction!(model::NonhydrostaticModel, Δt)
 
@@ -7,15 +42,16 @@ Calculate the (nonhydrostatic) pressure correction associated `tendencies`, `vel
 """
 function calculate_pressure_correction!(model::NonhydrostaticModel, Δt)
 
+    subtract_stokes_drift!(model)
+
     # Mask immersed velocities
     foreach(mask_immersed_field!, model.velocities)
-
     fill_halo_regions!(model.velocities, model.clock, fields(model))
-
     solve_for_pressure!(model.pressures.pNHS, model.pressure_solver, Δt, model.velocities)
-
     fill_halo_regions!(model.pressures.pNHS)
 
+    add_stokes_drift!(model, time(model) + Δt)
+
     return nothing
 end
 
diff --git a/src/StokesDrifts.jl b/src/StokesDrifts.jl
index e812b02f66..04285161da 100644
--- a/src/StokesDrifts.jl
+++ b/src/StokesDrifts.jl
@@ -18,6 +18,8 @@ using Oceananigans.Operators
 using Oceananigans.Grids: AbstractGrid, node
 using Oceananigans.Utils: prettysummary
 
+using KernelAbstractions: @kernel, @index
+
 import Adapt: adapt_structure
 
 #####
@@ -148,7 +150,10 @@ const c = Center()
 @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)
                                                                  - ℑyzᵃᶜᶠ(i, j, k, grid, U.v) * sw.∂z_vˢ(znode(k, grid, f), time))
 
-struct StokesDrift{P, VX, WX, UY, WY, UZ, VZ, UT, VT, WT}
+struct StokesDrift{P, US, VS, WS, VX, WX, UY, WY, UZ, VZ, UT, VT, WT}
+    uˢ    :: US
+    vˢ    :: VS
+    wˢ    :: WS
     ∂x_vˢ :: VX
     ∂x_wˢ :: WX
     ∂y_uˢ :: UY
@@ -161,7 +166,10 @@ struct StokesDrift{P, VX, WX, UY, WY, UZ, VZ, UT, VT, WT}
     parameters :: P
 end
 
-adapt_structure(to, sd::StokesDrift) = StokesDrift(adapt(to, sd.∂x_vˢ),
+adapt_structure(to, sd::StokesDrift) = StokesDrift(adapt(to, sd.uˢ),
+                                                   adapt(to, sd.vˢ),
+                                                   adapt(to, sd.wˢ),
+                                                   adapt(to, sd.∂x_wˢ),
                                                    adapt(to, sd.∂x_wˢ),
                                                    adapt(to, sd.∂y_uˢ),
                                                    adapt(to, sd.∂y_wˢ),
@@ -283,7 +291,10 @@ StokesDrift{Nothing}:
 └── ∂t_wˢ: ∂t_wˢ
 ```
 """
-function StokesDrift(; ∂x_vˢ = zerofunction,
+function StokesDrift(; uˢ = zerofunction,
+                       vˢ = zerofunction,
+                       wˢ = zerofunction,
+                       ∂x_vˢ = zerofunction,
                        ∂x_wˢ = zerofunction,
                        ∂y_uˢ = zerofunction,
                        ∂y_wˢ = zerofunction,
@@ -294,7 +305,7 @@ function StokesDrift(; ∂x_vˢ = zerofunction,
                        ∂t_wˢ = zerofunction,
                        parameters = nothing)
 
-    return StokesDrift(∂x_vˢ, ∂x_wˢ, ∂y_uˢ, ∂y_wˢ, ∂z_uˢ, ∂z_vˢ, ∂t_uˢ, ∂t_vˢ, ∂t_wˢ, parameters)
+    return StokesDrift(uˢ, vˢ, wˢ, ∂x_vˢ, ∂x_wˢ, ∂y_uˢ, ∂y_wˢ, ∂z_uˢ, ∂z_vˢ, ∂t_uˢ, ∂t_vˢ, ∂t_wˢ, parameters)
 end
 
 const SD = StokesDrift
@@ -325,7 +336,6 @@ const SDnoP = StokesDrift{<:Nothing}
     return wᶠᶜᶜ * (∂z_uˢ - ∂x_wˢ) - vᶠᶜᶜ * (∂x_vˢ - ∂y_uˢ)
 end
 
-
 @inline function y_curl_Uˢ_cross_U(i, j, k, grid, sw::SD, U, time)
     wᶜᶠᶜ = ℑyzᵃᶠᶜ(i, j, k, grid, U.w)
     uᶜᶠᶜ = ℑxyᶜᶠᵃ(i, j, k, grid, U.u)