Close u/v momentum budget on HydrostaticFreeSurfaceModel

docs/src/momentum_equation.md

@@ -66,8 +66,9 @@ KernelFunctionOperation at (Center, Face, Center)
 
 The combined RHS — `Tendency(model)` — wraps Oceananigans' internal `*_velocity_tendency`
 kernel directly. The sum of the individual term diagnostics, taken with the correct
-signs from Oceananigans' source, equals `Tendency` to machine precision. For
-`NonhydrostaticModel` the u/v formula is
+signs from Oceananigans' source, equals `Tendency` to machine precision.
+
+### `NonhydrostaticModel` u/v budget
 
 ```math
 \text{Tendency} = -\text{Advection} + \text{BuoyancyAcceleration}
@@ -77,30 +78,53 @@ signs from Oceananigans' source, equals `Tendency` to machine precision. For
                   +\text{StokesShear} + \text{StokesTendency} + \text{Forcing}
 ```
 
-The w-momentum equation is similar but with one subtlety: Oceananigans' `w_velocity_tendency`
-has no explicit `-\partial_z p` line. The vertical hydrostatic balance is treated as a
-property of the pressure-projection step, and whether the buoyancy term appears in the
-tendency depends on whether the model uses hydrostatic/nonhydrostatic pressure splitting:
+### `NonhydrostaticModel` w budget
+
+The w-momentum equation has one subtlety: Oceananigans' `w_velocity_tendency` has no
+explicit `-\partial_z p` line. The vertical hydrostatic balance is treated as a property
+of the pressure-projection step, and whether the buoyancy term appears in the tendency
+depends on whether the model uses hydrostatic/nonhydrostatic pressure splitting:
 
 | `hydrostatic_pressure_anomaly` | Buoyancy term in `Tendency` | Budget formula |
 | --- | --- | --- |
 | `Field` (default, splitting on) | absorbed by ``pHY'`` | `-ADV - COR - TVISC + SS + ST + FORC` |
 | `nothing` (splitting off) | included as ``+b`` | `-ADV + BUOY - COR - TVISC + SS + ST + FORC` |
 
+### `HydrostaticFreeSurfaceModel` u/v budget
+
+HFS uses `hydrostatic_free_surface_{u,v}_velocity_tendency`, which differs from the NH
+tendency: there are no Stokes terms, no buoyancy term (absorbed into the hydrostatic
+pressure anomaly `pHY′`), and the free-surface contribution to the horizontal pressure
+gradient appears as a separate `BarotropicPressureGradient` term:
+
+```math
+\text{Tendency} = -\text{Advection}
+                  -\text{BarotropicPressureGradient}
+                  -\text{CoriolisAcceleration}
+                  -\text{PressureGradient}
+                  -\text{TotalViscousDissipation}
+                  +\text{Forcing}
+```
+
+For `ImplicitFreeSurface` (the default) and `SplitExplicitFreeSurface`,
+`BarotropicPressureGradient` returns zero (the contribution is handled inside the pressure
+solve), so the budget closes without it. For `ExplicitFreeSurface` it returns `g ∂_i η`
+and must be included for closure.
+
 ## `HydrostaticFreeSurfaceModel` caveats
 
 - `WMomentumEquation.Tendency`, `WMomentumEquation.Forcing(model, Val(:w))`, and the
   U/V/W `StokesShear` and `StokesTendency` throw an `ArgumentError` on
   `HydrostaticFreeSurfaceModel`: `w` is diagnosed from continuity rather than evolved by a
   prognostic equation, and HFS has no `stokes_drift` field.
-- `Advection(::HydrostaticFreeSurfaceModel)` wraps `div_𝐯u`/`div_𝐯v` (the flux-form
-  kernel). HFS's default `VectorInvariant` momentum advection uses `U_dot_∇u` instead, so
-  to use the `Advection` diagnostic with HFS you must build the model with a flux-form
-  scheme, e.g. `momentum_advection = Centered()`.
-- For HFS, `Tendency` wraps `hydrostatic_free_surface_*_velocity_tendency`, which differs
-  from the NH tendency: it has no Stokes terms, includes a barotropic free-surface
-  pressure gradient, and a grid-slope contribution. Budget closure for HFS is not yet
-  supported by this module — it would require additional diagnostics for those terms.
+- `Advection(::HydrostaticFreeSurfaceModel)` (and the explicit-args form
+  `Advection(::HydrostaticFreeSurfaceModel, u, v, w, scheme)`) both wrap
+  `U_dot_∇u`/`U_dot_∇v` — the vector-invariant kernel that HFS's tendency actually uses.
+  This works with HFS's default `VectorInvariant` momentum advection as well as any
+  flux-form scheme. The diagnostic returns a KFO whose kernel function is
+  `U_dot_∇u`/`U_dot_∇v`, whereas the NH `Advection` diagnostic returns a KFO whose kernel
+  function is `div_𝐯u`/`div_𝐯v` — both pass the `isa Advection` type check via the
+  `Advection` Union alias.
 
 ## `Forcing` dispatch quirk
 
@@ -125,6 +149,7 @@ Oceanostics.UMomentumEquation.Advection
 Oceanostics.UMomentumEquation.BuoyancyAcceleration
 Oceanostics.UMomentumEquation.CoriolisAcceleration
 Oceanostics.UMomentumEquation.PressureGradient
+Oceanostics.UMomentumEquation.BarotropicPressureGradient
 Oceanostics.UMomentumEquation.ViscousDissipation
 Oceanostics.UMomentumEquation.ImmersedViscousDissipation
 Oceanostics.UMomentumEquation.TotalViscousDissipation
@@ -141,6 +166,7 @@ Oceanostics.VMomentumEquation.Advection
 Oceanostics.VMomentumEquation.BuoyancyAcceleration
 Oceanostics.VMomentumEquation.CoriolisAcceleration
 Oceanostics.VMomentumEquation.PressureGradient
+Oceanostics.VMomentumEquation.BarotropicPressureGradient
 Oceanostics.VMomentumEquation.ViscousDissipation
 Oceanostics.VMomentumEquation.ImmersedViscousDissipation
 Oceanostics.VMomentumEquation.TotalViscousDissipation

src/Oceanostics.jl

@@ -14,13 +14,13 @@ export TracerAdvection, TracerDiffusion, TracerImmersedDiffusion, TracerTotalDif
 #---
 
 #+++ UMomentumEquation exports
-export UAdvection, UBuoyancyAcceleration, UCoriolisAcceleration, UPressureGradient,
+export UAdvection, UBuoyancyAcceleration, UCoriolisAcceleration, UPressureGradient, UBarotropicPressureGradient,
        UViscousDissipation, UImmersedViscousDissipation, UTotalViscousDissipation,
        UStokesShear, UStokesTendency, UForcing, UTendency
 #---
 
 #+++ VMomentumEquation exports
-export VAdvection, VBuoyancyAcceleration, VCoriolisAcceleration, VPressureGradient,
+export VAdvection, VBuoyancyAcceleration, VCoriolisAcceleration, VPressureGradient, VBarotropicPressureGradient,
        VViscousDissipation, VImmersedViscousDissipation, VTotalViscousDissipation,
        VStokesShear, VStokesTendency, VForcing, VTendency
 #---

src/UMomentumEquation.jl

@@ -5,19 +5,20 @@ using Oceananigans: fields, Face, Center, KernelFunctionOperation, AbstractModel
 using Oceananigans.Models: HydrostaticFreeSurfaceModel
 using Oceananigans.Models.NonhydrostaticModels: u_velocity_tendency
 using Oceananigans.Models.HydrostaticFreeSurfaceModels: hydrostatic_free_surface_u_velocity_tendency
-using Oceananigans.Advection: div_𝐯u
+using Oceananigans.Advection: div_𝐯u, U_dot_∇u
 using Oceananigans.BuoyancyFormulations: x_dot_g_bᶠᶜᶜ
 using Oceananigans.Coriolis: x_f_cross_U
 using Oceananigans.TurbulenceClosures: ∂ⱼ_τ₁ⱼ, immersed_∂ⱼ_τ₁ⱼ
 using Oceananigans.StokesDrifts: x_curl_Uˢ_cross_U, ∂t_uˢ
 using Oceananigans.Operators: ∂xᶠᶜᶜ
+using Oceananigans.Models.HydrostaticFreeSurfaceModels: explicit_barotropic_pressure_x_gradient
 
 using Oceanostics: validate_location, CustomKFO
 
-export Advection, BuoyancyAcceleration, CoriolisAcceleration, PressureGradient,
+export Advection, BuoyancyAcceleration, CoriolisAcceleration, PressureGradient, BarotropicPressureGradient,
        ViscousDissipation, ImmersedViscousDissipation, TotalViscousDissipation,
        StokesShear, StokesTendency, Forcing, Tendency,
-       UAdvection, UBuoyancyAcceleration, UCoriolisAcceleration, UPressureGradient,
+       UAdvection, UBuoyancyAcceleration, UCoriolisAcceleration, UPressureGradient, UBarotropicPressureGradient,
        UViscousDissipation, UImmersedViscousDissipation, UTotalViscousDissipation,
        UStokesShear, UStokesTendency, UForcing, UTendency
 
@@ -30,24 +31,27 @@ export Advection, BuoyancyAcceleration, CoriolisAcceleration, PressureGradient,
 @inline hydrostatic_pressure_gradient_x(i, j, k, grid, hydrostatic_pressure) = ∂xᶠᶜᶜ(i, j, k, grid, hydrostatic_pressure)
 @inline hydrostatic_pressure_gradient_x(i, j, k, grid, ::Nothing) = zero(grid)
 
-# Type aliases for major functions
-const Advection = CustomKFO{<:typeof(div_𝐯u)}
+# Type aliases for major functions. `Advection` is a Union so that the NH flux-form
+# (`div_𝐯u`) and the HFS vector-invariant (`U_dot_∇u`) wrappers share the same alias.
+const Advection = Union{CustomKFO{<:typeof(div_𝐯u)}, CustomKFO{<:typeof(U_dot_∇u)}}
 const BuoyancyAcceleration = CustomKFO{<:typeof(x_dot_g_bᶠᶜᶜ)}
 const CoriolisAcceleration = CustomKFO{<:typeof(x_f_cross_U)}
 const PressureGradient = CustomKFO{<:typeof(hydrostatic_pressure_gradient_x)}
+const BarotropicPressureGradient = CustomKFO{<:typeof(explicit_barotropic_pressure_x_gradient)}
 const ViscousDissipation = CustomKFO{<:typeof(∂ⱼ_τ₁ⱼ)}
 const ImmersedViscousDissipation = CustomKFO{<:typeof(immersed_∂ⱼ_τ₁ⱼ)}
 const TotalViscousDissipation = CustomKFO{<:typeof(total_∂ⱼ_τ₁ⱼ)}
 const StokesShear = CustomKFO{<:typeof(x_curl_Uˢ_cross_U)}
 const StokesTendency = CustomKFO{<:typeof(∂t_uˢ)}
 const Forcing = KernelFunctionOperation{<:Any, <:Any, <:Any, <:Any, <:Any, <:Any}
 const Tendency = Union{CustomKFO{<:typeof(u_velocity_tendency)},
-                            CustomKFO{<:typeof(hydrostatic_free_surface_u_velocity_tendency)}}
+                       CustomKFO{<:typeof(hydrostatic_free_surface_u_velocity_tendency)}}
 
 const UAdvection = Advection
 const UBuoyancyAcceleration = BuoyancyAcceleration
 const UCoriolisAcceleration = CoriolisAcceleration
 const UPressureGradient = PressureGradient
+const UBarotropicPressureGradient = BarotropicPressureGradient
 const UViscousDissipation = ViscousDissipation
 const UImmersedViscousDissipation = ImmersedViscousDissipation
 const UTotalViscousDissipation = TotalViscousDissipation
@@ -60,11 +64,16 @@ const UTendency = Tendency
 """
     $(SIGNATURES)
 
-Calculate the advection of u-momentum as
+Calculate the advection of u-momentum. For `NonhydrostaticModel` this wraps the flux-form
+divergence kernel [`div_𝐯u`](https://clima.github.io/OceananigansDocumentation/stable/appendix/library/#Oceananigans.Advection.div_𝐯u-NTuple{7,%20Any})
+as
 
     ADV = ∂ⱼ (uⱼ u)
 
-using Oceananigans' kernel [`div_𝐯u`.](https://clima.github.io/OceananigansDocumentation/stable/appendix/library/#Oceananigans.Advection.div_𝐯u-NTuple{7,%20Any})
+For `HydrostaticFreeSurfaceModel` this wraps the vector-invariant kernel `U_dot_∇u` (which
+is what `hydrostatic_free_surface_u_velocity_tendency` actually uses) as
+
+    ADV = (u · ∇) u
 
 ```jldoctest
 julia> using Oceananigans, Oceanostics
@@ -86,8 +95,18 @@ function Advection(model, u, v, w, advection_scheme; location = (Face, Center, C
     return KernelFunctionOperation{Face, Center, Center}(div_𝐯u, model.grid, advection_scheme, total_velocities, u)
 end
 
+function Advection(model::HydrostaticFreeSurfaceModel, advection_scheme, velocities; location = (Face, Center, Center))
+    validate_location(location, "Advection", (Face, Center, Center))
+    return KernelFunctionOperation{Face, Center, Center}(U_dot_∇u, model.grid, advection_scheme, velocities)
+end
+
+# HFS-typed 5-arg overload so that explicit calls `Advection(hfs_model, u, v, w, scheme)` also use
+# `U_dot_∇u` rather than dispatching to the generic NH `div_𝐯u` form.
+Advection(model::HydrostaticFreeSurfaceModel, u, v, w, advection_scheme; kwargs...) =
+    Advection(model, advection_scheme, (; u, v, w); kwargs...)
+
 Advection(model; kwargs...)                              = Advection(model, model.velocities..., model.advection; kwargs...)
-Advection(model::HydrostaticFreeSurfaceModel; kwargs...) = Advection(model, model.velocities..., model.advection.momentum; kwargs...)
+Advection(model::HydrostaticFreeSurfaceModel; kwargs...) = Advection(model, model.advection.momentum, model.velocities; kwargs...)
 #---
 
 #+++ Buoyancy acceleration
@@ -198,6 +217,41 @@ function PressureGradient(model; kwargs...)
 end
 #---
 
+#+++ Barotropic pressure gradient
+"""
+    $(SIGNATURES)
+
+Calculate the explicit barotropic pressure gradient in the x-direction. For
+`HydrostaticFreeSurfaceModel` this is the contribution of the free surface ``η`` to the
+horizontal pressure gradient that's treated explicitly during time-stepping (i.e., not
+solved implicitly). For `ExplicitFreeSurface` it equals `g ∂x η`; for `ImplicitFreeSurface`
+or `SplitExplicitFreeSurface` the contribution is handled inside the pressure solve so
+this kernel returns zero. For `NonhydrostaticModel` (`free_surface == nothing`) it also
+returns zero. This term completes the u-momentum budget closure on HFS.
+
+```jldoctest
+julia> using Oceananigans, Oceanostics
+
+julia> grid = RectilinearGrid(size=(4, 4, 4), extent=(1, 1, 1));
+
+julia> model = NonhydrostaticModel(grid);
+
+julia> BARO = UMomentumEquation.BarotropicPressureGradient(model)
+KernelFunctionOperation at (Face, Center, Center)
+├── grid: 4×4×4 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 3×3×3 halo
+├── kernel_function: explicit_barotropic_pressure_x_gradient (generic function with 4 methods)
+└── arguments: ("Nothing",)
+```
+"""
+function BarotropicPressureGradient(model, free_surface; location = (Face, Center, Center))
+    validate_location(location, "BarotropicPressureGradient", (Face, Center, Center))
+    return KernelFunctionOperation{Face, Center, Center}(explicit_barotropic_pressure_x_gradient, model.grid, free_surface)
+end
+
+BarotropicPressureGradient(model; kwargs...) =
+    BarotropicPressureGradient(model, hasfield(typeof(model), :free_surface) ? model.free_surface : nothing; kwargs...)
+#---
+
 #+++ Viscous dissipation
 """
     $(SIGNATURES)