Implicit free surface formulation for NonhydrostaticModel (#4740)

* Add Combination boundary condition classification

* adapt regularize_boundary_condition to work with CombinationBoundaryCondition

* better names

* validate -> regularize

* Update test_boundary_conditions.jl

* Combination -> Mixed

* formatting

* add InhomogeneousFormulation to tridiagonal pressure solver

* add free_surface property to NonhydrostaticModel

* WIP nonhydrostatic free surface pressure solver

* fix bug launching internal_work_layout?

* Write a routine to compute the RHS of the pressure poisson equation for implicit free surface

* persist poisson eigenvalues in FourierTridiagonalPoissonSolver

* update tridiagonal term for implicit free surface

* fix

* bugfix

* bugfix

* bugfixes

* fix rhs + diagonal computation for ifs

* add MixedBoundaryCondition to pressure for ifs case

* WIP (code runs, but sign errors may remain)

* Fixed typos, added w update at surface, commented out no-penetration bc for w in field tuples, commented out subtraction of mean in solver

* fix top w bc for non-nothing free surface in nonhydrostaticmodel

* Update src/Models/NonhydrostaticModels/nonhydrostatic_model.jl

* use InhomogeneousFormulation for pressure solver

* fixes

* Update src/Fields/field_tuples.jl

* Update src/Fields/field_tuples.jl

* Update src/Solvers/fourier_tridiagonal_poisson_solver.jl

* move surface velocity computation to separate kernel

* fix indexing

* disambiguate

* fix function signatures

* Update src/Models/NonhydrostaticModels/pressure_correction.jl

* Update src/Models/NonhydrostaticModels/pressure_correction.jl

* fix test utility

* fix

* fix solve_for_pressure! signatures

* throw warning when using naive solver properly

---------

Co-authored-by: Shriya Fruitwala <[email protected]>

Author: glwagner

src/Fields/field_tuples.jl

@@ -136,7 +136,7 @@ function VelocityFields(grid::AbstractGrid, user_bcs = NamedTuple())
     u = XFaceField(grid, boundary_conditions=bcs.u)
     v = YFaceField(grid, boundary_conditions=bcs.v)
     w = ZFaceField(grid, boundary_conditions=bcs.w)
-
+    
     return (u=u, v=v, w=w)
 end
 

src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl

@@ -17,6 +17,7 @@ using DocStringExtensions
 
 import Oceananigans: fields, prognostic_fields, initialize!
 import Oceananigans.Advection: cell_advection_timescale
+import Oceananigans.Models: materialize_free_surface
 import Oceananigans.TimeSteppers: step_lagrangian_particles!
 import Oceananigans.Architectures: on_architecture
 

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl

@@ -179,8 +179,7 @@ function HydrostaticFreeSurfaceModel(; grid,
 
     # Next, we form a list of default boundary conditions:
     field_names = constructor_field_names(velocities, tracers, free_surface, auxiliary_fields, biogeochemistry, grid)
-    default_boundary_conditions = NamedTuple{field_names}(Tuple(FieldBoundaryConditions()
-                                                          for name in field_names))
+    default_boundary_conditions = NamedTuple{field_names}(FieldBoundaryConditions() for name in field_names)
 
     # Then we merge specified, embedded, and default boundary conditions. Specified boundary conditions
     # have precedence, followed by embedded, followed by default.

src/Models/Models.jl

@@ -89,6 +89,9 @@ validate_tracer_advection(tracer_advection_tuple::NamedTuple, grid) = Centered()
 validate_tracer_advection(tracer_advection::AbstractAdvectionScheme, grid) = tracer_advection, NamedTuple()
 validate_tracer_advection(tracer_advection::Nothing, grid) = nothing, NamedTuple()
 
+# Used in both NonhydrostaticModels and HydrostaticFreeSurfaceModels
+function materialize_free_surface end
+
 # Communication - Computation overlap in distributed models
 include("interleave_communication_and_computation.jl")
 

src/Models/NonhydrostaticModels/NonhydrostaticModels.jl

@@ -16,32 +16,42 @@ using Oceananigans.DistributedComputations: DistributedFFTBasedPoissonSolver, Di
 using Oceananigans.Grids: XYRegularRG, XZRegularRG, YZRegularRG, XYZRegularRG
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
 using Oceananigans.Solvers: GridWithFFTSolver, GridWithFourierTridiagonalSolver, ConjugateGradientPoissonSolver
+using Oceananigans.Solvers: InhomogeneousFormulation, ZDirection
 using Oceananigans.Utils: sum_of_velocities
 
 import Oceananigans: fields, prognostic_fields
 import Oceananigans.Advection: cell_advection_timescale
 import Oceananigans.TimeSteppers: step_lagrangian_particles!
 
-function nonhydrostatic_pressure_solver(::Distributed, local_grid::XYZRegularRG)
+function nonhydrostatic_pressure_solver(::Distributed, local_grid::XYZRegularRG, ::Nothing)
     global_grid = reconstruct_global_grid(local_grid)
     return DistributedFFTBasedPoissonSolver(global_grid, local_grid)
 end
 
-function nonhydrostatic_pressure_solver(::Distributed, local_grid::GridWithFourierTridiagonalSolver)
+function nonhydrostatic_pressure_solver(::Distributed, local_grid::GridWithFourierTridiagonalSolver, ::Nothing)
     global_grid = reconstruct_global_grid(local_grid)
     return DistributedFourierTridiagonalPoissonSolver(global_grid, local_grid)
 end
 
-nonhydrostatic_pressure_solver(arch, grid::XYZRegularRG) = FFTBasedPoissonSolver(grid)
-nonhydrostatic_pressure_solver(arch, grid::GridWithFourierTridiagonalSolver) =
-    FourierTridiagonalPoissonSolver(grid)
+nonhydrostatic_pressure_solver(arch, grid::XYZRegularRG, ::Nothing) = FFTBasedPoissonSolver(grid)
 
-# fallback
-nonhydrostatic_pressure_solver(arch, grid) = ConjugateGradientPoissonSolver(grid)
+function nonhydrostatic_pressure_solver(arch, grid::GridWithFourierTridiagonalSolver, ::Nothing)
+    return FourierTridiagonalPoissonSolver(grid)
+end
+
+function nonhydrostatic_pressure_solver(arch, grid::Union{GridWithFourierTridiagonalSolver, XYZRegularRG}, free_surface)
+    tridiagonal_formulation = InhomogeneousFormulation(ZDirection())
+    return FourierTridiagonalPoissonSolver(grid; tridiagonal_formulation)
+end
 
-const IBGWithFFTSolver = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:GridWithFFTSolver}
+# fallback
+nonhydrostatic_pressure_solver(arch, grid, ::Nothing) = ConjugateGradientPoissonSolver(grid)
 
-function nonhydrostatic_pressure_solver(arch, ibg::IBGWithFFTSolver)
+const IBGWithFFT = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:GridWithFFTSolver}
+nonhydrostatic_pressure_solver(arch, ibg::IBGWithFFT, ::Nothing) = naive_solver_with_warning(arch, ibg, nothing)
+nonhydrostatic_pressure_solver(arch, ibg::IBGWithFFT, fs) = naive_solver_with_warning(arch, ibg, fs)
+    
+function naive_solver_with_warning(arch, ibg, free_surface)
     msg = """The FFT-based pressure_solver for NonhydrostaticModels on ImmersedBoundaryGrid
           is approximate and will probably produce velocity fields that are divergent
           adjacent to the immersed boundary. An experimental but improved pressure_solver
@@ -54,10 +64,11 @@ function nonhydrostatic_pressure_solver(arch, ibg::IBGWithFFTSolver)
           """
     @warn msg
 
-    return nonhydrostatic_pressure_solver(arch, ibg.underlying_grid)
+    return nonhydrostatic_pressure_solver(arch, ibg.underlying_grid, free_surface)
 end
 
-nonhydrostatic_pressure_solver(grid) = nonhydrostatic_pressure_solver(architecture(grid), grid)
+
+nonhydrostatic_pressure_solver(grid, free_surface) = nonhydrostatic_pressure_solver(architecture(grid), grid, free_surface)
 
 #####
 ##### NonhydrostaticModel definition

src/Models/NonhydrostaticModels/nonhydrostatic_model.jl

@@ -4,13 +4,14 @@ using Oceananigans.Architectures: AbstractArchitecture
 using Oceananigans.DistributedComputations: Distributed
 using Oceananigans.Advection: Centered, adapt_advection_order
 using Oceananigans.BuoyancyFormulations: validate_buoyancy, regularize_buoyancy, SeawaterBuoyancy
+using Oceananigans.BoundaryConditions: MixedBoundaryCondition
 using Oceananigans.Biogeochemistry: validate_biogeochemistry, AbstractBiogeochemistry, biogeochemical_auxiliary_fields
 using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions
 using Oceananigans.Fields: Field, tracernames, VelocityFields, TracerFields, CenterField
 using Oceananigans.Forcings: model_forcing
 using Oceananigans.Grids: inflate_halo_size, with_halo, architecture
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
-using Oceananigans.Models: AbstractModel, extract_boundary_conditions
+using Oceananigans.Models: AbstractModel, extract_boundary_conditions, materialize_free_surface
 using Oceananigans.Solvers: FFTBasedPoissonSolver
 using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state!, AbstractLagrangianParticles
 using Oceananigans.TurbulenceClosures: validate_closure, with_tracers, build_diffusivity_fields, time_discretization, implicit_diffusion_solver
@@ -29,7 +30,7 @@ const BFOrNamedTuple = Union{BackgroundFields, NamedTuple}
 # but for now we use it only for hydrostatic pressure anomalies for now.
 struct DefaultHydrostaticPressureAnomaly end
 
-mutable struct NonhydrostaticModel{TS, E, A<:AbstractArchitecture, G, T, B, R, SD, U, C, Φ, F,
+mutable struct NonhydrostaticModel{TS, E, A<:AbstractArchitecture, G, T, B, R, SD, U, C, Φ, F, FS,
                                    V, S, K, BG, P, BGC, AF, BM} <: AbstractModel{TS, A}
 
          architecture :: A        # Computer `Architecture` on which `Model` is run
@@ -41,6 +42,7 @@ mutable struct NonhydrostaticModel{TS, E, A<:AbstractArchitecture, G, T, B, R, S
          stokes_drift :: SD       # Set of parameters for surfaces waves via the Craik-Leibovich approximation
               forcing :: F        # Container for forcing functions defined by the user
               closure :: E        # Diffusive 'turbulence closure' for all model fields
+         free_surface :: FS       # Linearized free surface representation (Nothing for rigid lid)
     background_fields :: BG       # Background velocity and tracer fields
             particles :: P        # Particle set for Lagrangian tracking
       biogeochemistry :: BGC      # Biogeochemistry for Oceananigans tracers
@@ -120,6 +122,7 @@ function NonhydrostaticModel(; grid,
                              stokes_drift = nothing,
                              forcing::NamedTuple = NamedTuple(),
                              closure = nothing,
+                             free_surface = nothing,
                              boundary_conditions::NamedTuple = NamedTuple(),
                              tracers = (),
                              timestepper = :RungeKutta3,
@@ -128,7 +131,7 @@ function NonhydrostaticModel(; grid,
                              biogeochemistry::AbstractBGCOrNothing = nothing,
                              velocities = nothing,
                              hydrostatic_pressure_anomaly = DefaultHydrostaticPressureAnomaly(),
-                             nonhydrostatic_pressure = CenterField(grid),
+                             nonhydrostatic_pressure = nothing,
                              diffusivity_fields = nothing,
                              pressure_solver = nothing,
                              auxiliary_fields = NamedTuple())
@@ -137,6 +140,20 @@ function NonhydrostaticModel(; grid,
 
     tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example)
 
+    if isnothing(nonhydrostatic_pressure)
+        if isnothing(free_surface)
+            nonhydrostatic_pressure = CenterField(grid)
+        else
+        # elseif free_surface isa ImplicitFreeSurface
+            # Use a MixedBoundaryCondition for the top bc for pressure
+            coefficient = Ref(zero(grid))
+            combination = Field{Center, Center, Nothing}(grid)
+            top_bc = MixedBoundaryCondition(coefficient, combination)
+            pressure_bcs = FieldBoundaryConditions(grid, (Center(), Center(), Center()), top=top_bc)
+            nonhydrostatic_pressure = CenterField(grid; boundary_conditions=pressure_bcs)
+        end
+    end
+
     # Validate pressure fields
     nonhydrostatic_pressure isa Field{Center, Center, Center} ||
         throw(ArgumentError("nonhydrostatic_pressure must be CenterField(grid)."))
@@ -196,25 +213,43 @@ function NonhydrostaticModel(; grid,
 
     # Next, we form a list of default boundary conditions:
     field_names = (:u, :v, :w, tracernames(tracers)..., keys(auxiliary_fields)...)
-    default_boundary_conditions = NamedTuple{field_names}(FieldBoundaryConditions()
-                                                          for name in field_names)
+    default_boundary_conditions = NamedTuple{field_names}(FieldBoundaryConditions() for name in field_names)
 
     # Finally, we merge specified, embedded, and default boundary conditions. Specified boundary conditions
     # have precedence, followed by embedded, followed by default.
     boundary_conditions = merge(default_boundary_conditions, embedded_boundary_conditions, boundary_conditions)
+
+    if !isnothing(free_surface) # replace top boundary condition for `w` with `nothing`
+        w_bcs = boundary_conditions.w
+        w_bcs = FieldBoundaryConditions(top = nothing,
+                                        bottom = w_bcs.bottom,
+                                        east = w_bcs.east,
+                                        west = w_bcs.west,
+                                        south = w_bcs.south,
+                                        north = w_bcs.north,
+                                        immersed = w_bcs.immersed)
+
+        boundary_conditions  =  merge(boundary_conditions, (; w = w_bcs))
+    end
+
     boundary_conditions = regularize_field_boundary_conditions(boundary_conditions, grid, field_names)
 
     # Ensure `closure` describes all tracers
     closure = with_tracers(tracernames(tracers), closure)
 
+    # TODO: limit free surface to `nothing` (rigid lid) or ImplicitFreeSurface
+    if !isnothing(free_surface)
+        free_surface = materialize_free_surface(free_surface, velocities, grid)
+    end
+
     # Either check grid-correctness, or construct tuples of fields
     velocities         = VelocityFields(velocities, grid, boundary_conditions)
     tracers            = TracerFields(tracers,      grid, boundary_conditions)
     pressures          = (pNHS=nonhydrostatic_pressure, pHY′=hydrostatic_pressure_anomaly)
     diffusivity_fields = build_diffusivity_fields(diffusivity_fields, grid, clock, tracernames(tracers), boundary_conditions, closure)
 
     if isnothing(pressure_solver)
-        pressure_solver = nonhydrostatic_pressure_solver(grid)
+        pressure_solver = nonhydrostatic_pressure_solver(grid, free_surface)
     end
 
     # Materialize background fields
@@ -224,7 +259,7 @@ function NonhydrostaticModel(; grid,
 
     # Instantiate timestepper if not already instantiated
     implicit_solver = implicit_diffusion_solver(time_discretization(closure), grid)
-    timestepper = TimeStepper(timestepper, grid, prognostic_fields; implicit_solver=implicit_solver)
+    timestepper = TimeStepper(timestepper, grid, prognostic_fields; implicit_solver)
 
     # Regularize forcing for model tracer and velocity fields.
     forcing = model_forcing(model_fields; forcing...)
@@ -235,7 +270,7 @@ function NonhydrostaticModel(; grid,
     !isnothing(particles) && arch isa Distributed && error("LagrangianParticles are not supported on Distributed architectures.")
 
     model = NonhydrostaticModel(arch, grid, clock, advection, buoyancy, coriolis, stokes_drift,
-                                forcing, closure, background_fields, particles, biogeochemistry, velocities, tracers,
+                                forcing, closure, free_surface, background_fields, particles, biogeochemistry, velocities, tracers,
                                 pressures, diffusivity_fields, timestepper, pressure_solver, auxiliary_fields, boundary_mass_fluxes)
 
     update_state!(model; compute_tendencies = false)

src/Models/NonhydrostaticModels/pressure_correction.jl

@@ -10,15 +10,53 @@ function compute_pressure_correction!(model::NonhydrostaticModel, Δt)
     # Mask immersed velocities
     foreach(mask_immersed_field!, model.velocities)
     fill_halo_regions!(model.velocities, model.clock, fields(model))
-
     enforce_open_boundary_mass_conservation!(model, model.boundary_mass_fluxes)
 
-    solve_for_pressure!(model.pressures.pNHS, model.pressure_solver, Δt, model.velocities)
-    fill_halo_regions!(model.pressures.pNHS)
+    p_Δt = model.pressures.pNHS
+    solve_for_pressure!(p_Δt, model.pressure_solver, model.free_surface, model.velocities, Δt)
+
+    set_top_pressure_boundary_condition!(p_Δt, model.free_surface, model.velocities.w, Δt)
+    fill_halo_regions!(p_Δt)
+
+    return nothing
+end
+
+# TODO: make these fallbacks
+# Routines for rigid lid
+set_top_pressure_boundary_condition!(p_Δt, ::Nothing, w̃, Δt) = nothing
+
+function set_top_pressure_boundary_condition!(p_Δt, free_surface, w̃, Δt)
+    top_bc = p_Δt.boundary_conditions.top
+    g = free_surface.gravitational_acceleration
+
+    # Set the "coefficient" of the MixedBoundaryCondition
+    top_bc.condition.coefficient[] = 1 / (g * Δt^2)
 
+    # Set the "combination" of the MixedBoundaryCondition
+    η = free_surface.η
+    bc_rhs = top_bc.condition.inhomogeneity
+    grid = p_Δt.grid
+    arch = grid.architecture
+    launch!(arch, grid, :xy, _set_top_pressure_boundary_condition!, bc_rhs, grid, w̃, Δt, η, g, p_Δt)
     return nothing
 end
 
+@kernel function _set_top_pressure_boundary_condition!(bc_rhs, grid, w̃, Δt, η, g, p_Δt)
+    i, j = @index(Global, NTuple)
+    Nz = grid.Nz
+
+    @inbounds bc_rhs[i, j, 1] = η[i, j, Nz+1] / Δt + w̃[i, j, Nz+1]
+
+    # Compute ηⁿ⁺¹
+    Δzᶠ = Δzᵃᵃᶠ(i, j, Nz+1, grid)
+    a = 1 / Δzᶠ - 1 / (2g * Δt^2)
+    b = 1 / Δzᶠ + 1 / (2g * Δt^2)
+    @inbounds begin
+        η★ = η[i, j, Nz+1] + Δt * w̃[i, j, Nz+1]
+        η[i, j, Nz+1] = (p_Δt[i, j, Nz] / Δt + (p_Δt[i, j, Nz] * a / (b * Δt) + η★ / (b * Δt^2))) / 2g
+    end
+end
+
 #####
 ##### Fractional and time stepping
 #####
@@ -36,15 +74,32 @@ Update the predictor velocities u, v, and w with the non-hydrostatic pressure mu
     @inbounds U.w[i, j, k] -= ∂zᶜᶜᶠ(i, j, k, grid, pNHSΔt)
 end
 
+@kernel function _compute_surface_vertical_velocity!(w, grid, pNHSΔt)
+    i, j = @index(Global, NTuple)
+    k = grid.Nz + 1
+    @inbounds w[i, j, k] -= ∂zᶜᶜᶠ(i, j, k, grid, pNHSΔt)
+end
+
 "Update the solution variables (velocities and tracers)."
 function make_pressure_correction!(model::NonhydrostaticModel, Δt)
 
-    launch!(model.architecture, model.grid, :xyz,
+    grid = model.grid
+    arch = grid.architecture
+
+    launch!(arch, grid, :xyz,
             _make_pressure_correction!,
             model.velocities,
             model.grid,
             model.pressures.pNHS)
 
+    if !isnothing(model.free_surface)
+        launch!(arch, grid, :xy,
+                _compute_surface_vertical_velocity!,
+                model.velocities.w,
+                model.grid,
+                model.pressures.pNHS)
+    end
+
     ϵ = eps(eltype(model.pressures.pNHS))
     Δt⁺ = max(ϵ, Δt)
     model.pressures.pNHS ./= Δt⁺

src/Models/NonhydrostaticModels/set_nonhydrostatic_model.jl

@@ -36,6 +36,8 @@ function set!(model::NonhydrostaticModel; enforce_incompressibility=true, kwargs
             ϕ = getproperty(model.velocities, fldname)
         elseif fldname ∈ propertynames(model.tracers)
             ϕ = getproperty(model.tracers, fldname)
+        elseif !isnothing(model.free_surface) && fldname ∈ propertynames(model.free_surface)
+            ϕ = getproperty(model.free_surface, fldname)
         else
             throw(ArgumentError("name $fldname not found in model.velocities or model.tracers."))
         end