Implement an implicit free surface solver in the NonhydrostaticModel

src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl

@@ -37,6 +37,7 @@ free_surface_displacement_field(velocities, ::Nothing, grid) = nothing
 initialize_free_surface!(free_surface, grid, velocities) = nothing
 
 include("compute_w_from_continuity.jl")
+include("update_hydrostatic_pressure.jl")
 
 # No free surface
 include("nothing_free_surface.jl")

src/Models/HydrostaticFreeSurfaceModels/barotropic_pressure_correction.jl

@@ -46,3 +46,4 @@ end
         U.v[i, j, k] -= g * Δt * ∂yᶜᶠᶠ(i, j, grid.Nz+1, grid, η)
     end
 end
+

src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_buffers.jl

@@ -3,10 +3,10 @@ import Oceananigans.Models: compute_buffer_tendencies!
 using Oceananigans.Grids: halo_size
 using Oceananigans.DistributedComputations: DistributedActiveCellsIBG
 using Oceananigans.ImmersedBoundaries: retrieve_interior_active_cells_map
-using Oceananigans.Models.NonhydrostaticModels: buffer_tendency_kernel_parameters,
-                                                buffer_p_kernel_parameters, 
-                                                buffer_κ_kernel_parameters,
-                                                buffer_parameters
+using Oceananigans.Models: buffer_tendency_kernel_parameters,
+                           buffer_p_kernel_parameters, 
+                           buffer_κ_kernel_parameters,
+                           buffer_parameters
 
 # We assume here that top/bottom BC are always synchronized (no partitioning in z)
 function compute_buffer_tendencies!(model::HydrostaticFreeSurfaceModel)

src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl

@@ -13,9 +13,8 @@ function ab2_step!(model::HydrostaticFreeSurfaceModel, Δt)
 
     compute_free_surface_tendency!(model.grid, model, model.free_surface)
 
-    χ = model.timestepper.χ
-
     # Step locally velocity and tracers
+    χ = model.timestepper.χ
     @apply_regionally local_ab2_step!(model, Δt, χ)
 
     step_free_surface!(model.free_surface, model, model.timestepper, Δt)
@@ -26,7 +25,6 @@ end
 function local_ab2_step!(model, Δt, χ)
     ab2_step_velocities!(model.velocities, model, Δt, χ)
     ab2_step_tracers!(model.tracers, model, Δt, χ)
-
     return nothing
 end
 

src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl

@@ -124,12 +124,12 @@ build_implicit_step_solver(::Val{:Default}, grid, settings, gravitational_accele
 Implicitly step forward η.
 """
 function step_free_surface!(free_surface::ImplicitFreeSurface, model, timestepper, Δt)
-    η      = free_surface.η
-    g      = free_surface.gravitational_acceleration
-    rhs    = free_surface.implicit_step_solver.right_hand_side
-    ∫ᶻQ    = free_surface.barotropic_volume_flux
+    η = free_surface.η
+    g = free_surface.gravitational_acceleration
+    rhs = free_surface.implicit_step_solver.right_hand_side
+    ∫ᶻQ = free_surface.barotropic_volume_flux
     solver = free_surface.implicit_step_solver
-    arch   = model.architecture
+    arch = model.architecture
 
     fill_halo_regions!(model.velocities, model.clock, fields(model))
 
@@ -160,3 +160,4 @@ function local_compute_integrated_volume_flux!(∫ᶻQ, velocities, arch)
 
     return nothing
 end
+

src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_free_surface_model_state.jl

@@ -7,10 +7,9 @@ using Oceananigans.BoundaryConditions: update_boundary_condition!
 using Oceananigans.TurbulenceClosures: compute_diffusivities!
 using Oceananigans.ImmersedBoundaries: mask_immersed_field!, mask_immersed_field_xy!, inactive_node
 using Oceananigans.Models: update_model_field_time_series!
-using Oceananigans.Models.NonhydrostaticModels: update_hydrostatic_pressure!, p_kernel_parameters
 using Oceananigans.Fields: replace_horizontal_vector_halos!
 
-import Oceananigans.Models.NonhydrostaticModels: compute_auxiliaries!
+import Oceananigans.Models: compute_auxiliaries!
 import Oceananigans.TimeSteppers: update_state!
 
 compute_auxiliary_fields!(auxiliary_fields) = Tuple(compute!(a) for a in auxiliary_fields)

src/Models/NonhydrostaticModels/update_hydrostatic_pressure.jl → src/Models/HydrostaticFreeSurfaceModels/update_hydrostatic_pressure.jl

@@ -1,8 +1,10 @@
 using Oceananigans.Operators: Δzᶜᶜᶜ, Δzᶜᶜᶠ
 using Oceananigans.ImmersedBoundaries: PartialCellBottom, ImmersedBoundaryGrid
-using Oceananigans.Grids: topology
+using Oceananigans.Grids: topology, Flat
 using Oceananigans.Grids: XFlatGrid, YFlatGrid
 
+import Oceananigans.Models: p_kernel_parameters
+
 """
 Update the hydrostatic pressure perturbation pHY′. This is done by integrating
 the `buoyancy_perturbationᶜᶜᶜ` downwards:

src/Models/Models.jl

@@ -13,10 +13,10 @@ using Oceananigans: AbstractModel, fields, prognostic_fields
 using Oceananigans.AbstractOperations: AbstractOperation
 using Oceananigans.Advection: AbstractAdvectionScheme, Centered, VectorInvariant
 using Oceananigans.Fields: AbstractField, Field, flattened_unique_values, boundary_conditions
-using Oceananigans.Grids: AbstractGrid, halo_size, inflate_halo_size
+using Oceananigans.Grids: AbstractGrid, halo_size, inflate_halo_size, topology
 using Oceananigans.OutputReaders: update_field_time_series!, extract_field_time_series
 using Oceananigans.TimeSteppers: AbstractTimeStepper, Clock
-using Oceananigans.Utils: Time
+using Oceananigans.Utils: Time, KernelParameters
 
 import Oceananigans: initialize!
 import Oceananigans.Architectures: architecture
@@ -94,8 +94,12 @@ include("interleave_communication_and_computation.jl")
 ##### All the code
 #####
 
-include("NonhydrostaticModels/NonhydrostaticModels.jl")
+function p_kernel_parameters end
+function compute_auxiliaries! end
+
+include("buffer_tendency_kernel_parameters.jl")
 include("HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl")
+include("NonhydrostaticModels/NonhydrostaticModels.jl")
 include("ShallowWaterModels/ShallowWaterModels.jl")
 include("LagrangianParticleTracking/LagrangianParticleTracking.jl")
 

src/Models/NonhydrostaticModels/NonhydrostaticModels.jl

@@ -15,6 +15,8 @@ using Oceananigans.DistributedComputations: reconstruct_global_grid, Distributed
 using Oceananigans.DistributedComputations: DistributedFFTBasedPoissonSolver, DistributedFourierTridiagonalPoissonSolver
 using Oceananigans.Grids: XYRegularRG, XZRegularRG, YZRegularRG, XYZRegularRG
 using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
+using Oceananigans.Models: p_kernel_parameters
+using Oceananigans.Models.HydrostaticFreeSurfaceModels: update_hydrostatic_pressure!
 using Oceananigans.Solvers: GridWithFFTSolver, GridWithFourierTridiagonalSolver 
 using Oceananigans.Utils: SumOfArrays
 
@@ -102,7 +104,6 @@ prognostic_fields(model::NonhydrostaticModel) = merge(model.velocities, model.tr
 step_lagrangian_particles!(model::NonhydrostaticModel, Δt) = step_lagrangian_particles!(model.particles, model, Δt)
 
 include("solve_for_pressure.jl")
-include("update_hydrostatic_pressure.jl")
 include("update_nonhydrostatic_model_state.jl")
 include("pressure_correction.jl")
 include("nonhydrostatic_tendency_kernel_functions.jl")

src/Models/NonhydrostaticModels/compute_nonhydrostatic_buffer_tendencies.jl

@@ -1,5 +1,10 @@
 import Oceananigans.Models: compute_buffer_tendencies!
 
+using Oceananigans.Models: buffer_tendency_kernel_parameters,
+                           buffer_p_kernel_parameters,
+                           buffer_κ_kernel_parameters,
+                           buffer_parameters
+
 using Oceananigans.TurbulenceClosures: required_halo_size_x, required_halo_size_y
 using Oceananigans.Grids: XFlatGrid, YFlatGrid
 
@@ -24,61 +29,3 @@ function compute_buffer_tendencies!(model::NonhydrostaticModel)
     return nothing
 end
 
-# tendencies need computing in the range 1 : H and N - H + 1 : N 
-function buffer_tendency_kernel_parameters(grid, arch)
-    Nx, Ny, Nz = size(grid)
-    Hx, Hy, _  = halo_size(grid)
-
-    param_west  = (1:Hx,       1:Ny,       1:Nz)
-    param_east  = (Nx-Hx+1:Nx, 1:Ny,       1:Nz)
-    param_south = (1:Nx,       1:Hy,       1:Nz)
-    param_north = (1:Nx,       Ny-Hy+1:Ny, 1:Nz)
-
-    params = (param_west, param_east, param_south, param_north)
-    return buffer_parameters(params, grid, arch)
-end
-
-# p needs computing in the range  0 : 0 and N + 1 : N + 1
-function buffer_p_kernel_parameters(grid, arch)
-    Nx, Ny, _ = size(grid)
-
-    param_west  = (0:0,       1:Ny)
-    param_east  = (Nx+1:Nx+1, 1:Ny)
-    param_south = (1:Nx,      0:0)
-    param_north = (1:Nx,      Ny+1:Ny+1)
-
-    params = (param_west, param_east, param_south, param_north)
-    return buffer_parameters(params, grid, arch)
-end
-
-# diffusivities need recomputing in the range 0 : B and N - B + 1 : N + 1
-function buffer_κ_kernel_parameters(grid, closure, arch)
-    Nx, Ny, Nz = size(grid)
-
-    Bx = required_halo_size_x(closure)
-    By = required_halo_size_y(closure)
-
-    param_west  = (0:Bx,         1:Ny,         1:Nz)
-    param_east  = (Nx-Bx+1:Nx+1, 1:Ny,         1:Nz)
-    param_south = (1:Nx,         0:By,         1:Nz)
-    param_north = (1:Nx,         Ny-By+1:Ny+1, 1:Nz)
-
-    params = (param_west, param_east, param_south, param_north)
-    return buffer_parameters(params, grid, arch)
-end
-
-# Recompute only on communicating sides 
-function buffer_parameters(parameters, grid, arch) 
-    Rx, Ry, _ = arch.ranks
-    Tx, Ty, _ = topology(grid)
-
-    include_west  = !isa(grid, XFlatGrid) && (Rx != 1) && !(Tx == RightConnected)
-    include_east  = !isa(grid, XFlatGrid) && (Rx != 1) && !(Tx == LeftConnected)
-    include_south = !isa(grid, YFlatGrid) && (Ry != 1) && !(Ty == RightConnected)
-    include_north = !isa(grid, YFlatGrid) && (Ry != 1) && !(Ty == LeftConnected)
-
-    include_side = (include_west, include_east, include_south, include_north)
-
-    return Tuple(KernelParameters(parameters[i]...) for i in findall(include_side))
-end
-

src/Models/NonhydrostaticModels/nonhydrostatic_model.jl

@@ -19,6 +19,9 @@ using Oceananigans.TurbulenceClosures.TKEBasedVerticalDiffusivities: FlavorOfCAT
 using Oceananigans.Utils: tupleit
 using Oceananigans.Grids: topology
 
+using Oceananigans.Models.HydrostaticFreeSurfaceModels:
+    materialize_free_surface
+
 import Oceananigans.Architectures: architecture
 import Oceananigans.Models: total_velocities, default_nan_checker, timestepper
 
@@ -29,51 +32,53 @@ const AbstractBGCOrNothing = Union{Nothing, AbstractBiogeochemistry}
 # 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} <: AbstractModel{TS}
 
-         architecture :: A        # Computer `Architecture` on which `Model` is run
-                 grid :: G        # Grid of physical points on which `Model` is solved
-                clock :: Clock{T} # Tracks iteration number and simulation time of `Model`
-            advection :: V        # Advection scheme for velocities _and_ tracers
-             buoyancy :: B        # Set of parameters for buoyancy model
-             coriolis :: R        # Set of parameters for the background rotation rate of `Model`
-         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
-    background_fields :: BG       # Background velocity and tracer fields
-            particles :: P        # Particle set for Lagrangian tracking
-      biogeochemistry :: BGC      # Biogeochemistry for Oceananigans tracers
-           velocities :: U        # Container for velocity fields `u`, `v`, and `w`
-              tracers :: C        # Container for tracer fields
-            pressures :: Φ        # Container for hydrostatic and nonhydrostatic pressure
-   diffusivity_fields :: K        # Container for turbulent diffusivities
-          timestepper :: TS       # Object containing timestepper fields and parameters
-      pressure_solver :: S        # Pressure/Poisson solver
-     auxiliary_fields :: AF       # User-specified auxiliary fields for forcing functions and boundary conditions
+    architecture :: A           # Computer `Architecture` on which `Model` is run
+    grid :: G                   # Grid of physical points on which `Model` is solved
+    clock :: Clock{T}           # Tracks iteration number and simulation time of `Model`
+    advection :: V              # Advection scheme for velocities _and_ tracers
+    buoyancy :: B               # Set of parameters for buoyancy model
+    coriolis :: R               # Set of parameters for the background rotation rate of `Model`
+    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          # Free surface?
+    background_fields :: BG     # Background velocity and tracer fields
+    particles :: P              # Particle set for Lagrangian tracking
+    biogeochemistry :: BGC      # Biogeochemistry for Oceananigans tracers
+    velocities :: U             # Container for velocity fields `u`, `v`, and `w`
+    tracers :: C                # Container for tracer fields
+    pressures :: Φ              # Container for hydrostatic and nonhydrostatic pressure
+    diffusivity_fields :: K     # Container for turbulent diffusivities
+    timestepper :: TS           # Object containing timestepper fields and parameters
+    pressure_solver :: S        # Pressure/Poisson solver
+    auxiliary_fields :: AF      # User-specified auxiliary fields for forcing functions and boundary conditions
 end
 
 """
-    NonhydrostaticModel(;           grid,
-                                    clock = Clock{eltype(grid)}(time = 0),
-                                advection = Centered(),
-                                 buoyancy = nothing,
-                                 coriolis = nothing,
-                             stokes_drift = nothing,
-                      forcing::NamedTuple = NamedTuple(),
-                                  closure = nothing,
-          boundary_conditions::NamedTuple = NamedTuple(),
-                                  tracers = (),
-                              timestepper = :RungeKutta3,
-            background_fields::NamedTuple = NamedTuple(),
-            particles::ParticlesOrNothing = nothing,
-    biogeochemistry::AbstractBGCOrNothing = nothing,
-                               velocities = nothing,
-                  nonhydrostatic_pressure = CenterField(grid),
-             hydrostatic_pressure_anomaly = DefaultHydrostaticPressureAnomaly(),
-                       diffusivity_fields = nothing,
-                          pressure_solver = nothing,
-                         auxiliary_fields = NamedTuple())
+    NonhydrostaticModel(; grid,
+                        clock = Clock{eltype(grid)}(time = 0),
+                        advection = Centered(),
+                        buoyancy = nothing,
+                        coriolis = nothing,
+                        stokes_drift = nothing,
+                        forcing::NamedTuple = NamedTuple(),
+                        closure = nothing,
+                        free_surface = nothing,
+                        boundary_conditions::NamedTuple = NamedTuple(),
+                        tracers = (),
+                        timestepper = :RungeKutta3,
+                        background_fields::NamedTuple = NamedTuple(),
+                        particles::ParticlesOrNothing = nothing,
+                        biogeochemistry::AbstractBGCOrNothing = nothing,
+                        velocities = nothing,
+                        nonhydrostatic_pressure = CenterField(grid),
+                        hydrostatic_pressure_anomaly = DefaultHydrostaticPressureAnomaly(),
+                        diffusivity_fields = nothing,
+                        pressure_solver = nothing,
+                        auxiliary_fields = NamedTuple())
 
 Construct a model for a non-hydrostatic, incompressible fluid on `grid`, using the Boussinesq
 approximation when `buoyancy != nothing`. By default, all Bounded directions are rigid and impenetrable.
@@ -91,6 +96,7 @@ Keyword arguments
   - `stokes_drift`: Parameters for Stokes drift fields associated with surface waves. Default: `nothing`.
   - `forcing`: `NamedTuple` of user-defined forcing functions that contribute to solution tendencies.
   - `closure`: The turbulence closure for `model`. See `Oceananigans.TurbulenceClosures`.
+  - `free_surface`: The free surface formulation to use. Default: `nothing`, which makes the rigid lid approximation.
   - `boundary_conditions`: `NamedTuple` containing field boundary conditions.
   - `tracers`: A tuple of symbols defining the names of the modeled tracers, or a `NamedTuple` of
                preallocated `CenterField`s.
@@ -119,6 +125,7 @@ function NonhydrostaticModel(; grid,
                              stokes_drift = nothing,
                              forcing::NamedTuple = NamedTuple(),
                              closure = nothing,
+                             free_surface = nothing,
                              boundary_conditions::NamedTuple = NamedTuple(),
                              tracers = (),
                              timestepper = :RungeKutta3,
@@ -212,6 +219,11 @@ function NonhydrostaticModel(; grid,
     pressures          = (pNHS=nonhydrostatic_pressure, pHY′=hydrostatic_pressure_anomaly)
     diffusivity_fields = DiffusivityFields(diffusivity_fields, grid, tracernames(tracers), boundary_conditions, 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
+
     if isnothing(pressure_solver)
         pressure_solver = nonhydrostatic_pressure_solver(grid)
     end
@@ -228,7 +240,7 @@ function NonhydrostaticModel(; grid,
     forcing = model_forcing(model_fields; forcing...)
 
     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)
 
     update_state!(model; compute_tendencies = false)

src/Models/NonhydrostaticModels/pressure_correction.jl

@@ -1,19 +1,25 @@
 import Oceananigans.TimeSteppers: calculate_pressure_correction!, pressure_correct_velocities!
 
+using Oceananigans.Models.HydrostaticFreeSurfaceModels: step_free_surface!
+
 """
     calculate_pressure_correction!(model::NonhydrostaticModel, Δt)
 
 Calculate the (nonhydrostatic) pressure correction associated `tendencies`, `velocities`, and step size `Δt`.
 """
 function calculate_pressure_correction!(model::NonhydrostaticModel, Δt)
 
+    if !isnothing(model.free_surface)
+        step_free_surface!(model.free_surface, model, model.timestepper, Δt)
+        # "First" barotropic pressure correction
+        pressure_correct_velocities!(model, model.free_surface, Δt)
+    end
+
     # 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)
 
     return nothing