(0.94.1) Refactor SplitExplicitFreeSurface (#3894)

* new bottom height * now it should work * comment * comment * remove circular dependency for now * some bugfixes * change name to column_height * correct column height * whoops * another correction * some more changes * another correction * couple of more bugfixes * more bugfixes * this should make it work * unify the formulations * correct implementation * correct implementation * correct partial cell bottom * use center immersed condition for grid fitted boundary * use the *correct* center node * no h for z-values! * simplify partial cells * make sure we don't go out of bounds * back to immersed condition * name changes * bugfix * another bugfix * fix bugs * some corrections * another bugfix * domain_depth * some remaining `z_bottom`s * back as it was * bugfix * correct correction * static_column_depth * Update src/Grids/grid_utils.jl Co-authored-by: Gregory L. Wagner <[email protected]> * address comments * new comment * another name change * AGFBIBG istead of AGFBIB and z_bottom only in TurbulenceClosures * some bugfixes * better definition of bottom height * fixed partial cell * fixed partial cells * remove OrthogonalSphericalShellGrids while we decide what to do * these files shouldn't go here * new operators * this was removed * start refactoring * start refactoring * some more refactoring * big overhaul * more refactoring * compiles * more refactoring... * correction * use a module for now * make module work * some more organizational stuff * some more changes * it compiles * using free_surface_displacement_field * include g_Earth * now it should run * a little bit of cleanup * tests should run now * Update Project.toml * conceptually this is better * fix checkpointer test * fix split explicit settings * fixing some more tests * import with_halo * back on Enzyme * fix tets * some fixes * these are prognostic fields * comment * fix tests * move functions to correct file * apply regionally * fix tests + unify implementation * remove old code * fix comment * bugfix * correction * another bugfix * switch right and left connected * all inside apply regionally * revert * tests corrected * fixed tests * explicit free surface tendency in its own file * test checkpointer * last bugfix * removed test script * last bugfix? * add to docs * thought I had already fixed this * bugfix * unpacking all the fields * last bugfix * fix typo in docs * remove stray spaces * empty line * minor * add reference * split lines for math rendering * better latex rendering * math rendering * Update docs/src/appendix/library.md Co-authored-by: Gregory L. Wagner <[email protected]> * better comment for the corrector * add comments in the docstring * remove kernel_parameters from the initial constructor * Update src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl Co-authored-by: Gregory L. Wagner <[email protected]> * move errors inside constructors * change summary strings * store instead of advance * change name to setup_free_surface! * simplify ab2_step_G * move `compute_free_surface_tendency!` where it should go * minimize using statements * bugfix * bugfix in integrated tendencies * revert quickly to see if this works * compute_free_surface_tendecny! is the only problem * no need for const c and const f * fix errors * eliding NaNs + separate compute_slow_tendencies * whoops wrong name * fix in the summary * Update src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl Co-authored-by: Gregory L. Wagner <[email protected]> * Update src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl Co-authored-by: Gregory L. Wagner <[email protected]> * remove vestigial `@show` --------- Co-authored-by: Gregory L. Wagner <[email protected]> Co-authored-by: Navid C. Constantinou <[email protected]>
CliMA · Nov 14, 2024 · bb42ddd · bb42ddd
1 parent 0be921d
commit bb42ddd
Show file tree

Hide file tree

Showing 31 changed files with 1,165 additions and 1,151 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.94.0"
+version = "0.94.1"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/docs/src/appendix/library.md b/docs/src/appendix/library.md
@@ -122,6 +122,13 @@ Modules = [Oceananigans.Models.HydrostaticFreeSurfaceModels]
 Private = false
 ```
 
+### Split-explicit free-surface
+
+```@autodocs
+Modules = [Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces]
+Private = false
+```
+
 ### Shallow-water models
 
 ```@autodocs

diff --git a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl
@@ -31,8 +31,13 @@ fill_horizontal_velocity_halos!(args...) = nothing
 free_surface_displacement_field(velocities, free_surface, grid) = ZFaceField(grid, indices = (:, :, size(grid, 3)+1))
 free_surface_displacement_field(velocities, ::Nothing, grid) = nothing
 
+# free surface initialization functions
+initialize_free_surface!(free_surface, grid, velocities) = nothing
+
 include("compute_w_from_continuity.jl")
-include("rigid_lid.jl")
+
+# No free surface
+include("nothing_free_surface.jl")
 
 # Explicit free-surface solver functionality
 include("explicit_free_surface.jl")
@@ -46,9 +51,9 @@ include("matrix_implicit_free_surface_solver.jl")
 include("implicit_free_surface.jl")
 
 # Split-Explicit free-surface solver functionality
-include("split_explicit_free_surface.jl")
-include("distributed_split_explicit_free_surface.jl")
-include("split_explicit_free_surface_kernels.jl")
+include("SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl")
+
+using .SplitExplicitFreeSurfaces
 
 include("hydrostatic_free_surface_field_tuples.jl")
 include("hydrostatic_free_surface_model.jl")
@@ -85,6 +90,13 @@ Return a flattened `NamedTuple` of the prognostic fields associated with `Hydros
                                                                                   η = free_surface.η),
                                                                                   tracers)
 
+@inline hydrostatic_prognostic_fields(velocities, free_surface::SplitExplicitFreeSurface, tracers) = merge((u = velocities.u,
+                                                                                                            v = velocities.v,
+                                                                                                            η = free_surface.η,
+                                                                                                            U = free_surface.barotropic_velocities.U,
+                                                                                                            V = free_surface.barotropic_velocities.V),
+                                                                                                            tracers)
+
 @inline hydrostatic_prognostic_fields(velocities, ::Nothing, tracers) = merge((u = velocities.u,
                                                                                v = velocities.v),
                                                                                tracers)
@@ -95,6 +107,14 @@ Return a flattened `NamedTuple` of the prognostic fields associated with `Hydros
                                                                        tracers,
                                                                        (; η = free_surface.η))
 
+@inline hydrostatic_fields(velocities, free_surface::SplitExplicitFreeSurface, tracers) = merge((u = velocities.u,
+                                                                                                 v = velocities.v,
+                                                                                                 w = velocities.w,
+                                                                                                 η = free_surface.η,
+                                                                                                 U = free_surface.barotropic_velocities.U,
+                                                                                                 V = free_surface.barotropic_velocities.V),
+                                                                                                 tracers)
+
 @inline hydrostatic_fields(velocities, ::Nothing, tracers) = merge((u = velocities.u,
                                                                     v = velocities.v,
                                                                     w = velocities.w),

diff --git a/...odels/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl b/...odels/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl
@@ -0,0 +1,44 @@
+module SplitExplicitFreeSurfaces
+
+export SplitExplicitFreeSurface, ForwardBackwardScheme, AdamsBashforth3Scheme
+export FixedSubstepNumber, FixedTimeStepSize
+
+using Oceananigans
+using Oceananigans.Architectures
+using Oceananigans.Architectures: convert_args
+using Oceananigans.Fields
+using Oceananigans.Utils
+using Oceananigans.Grids
+using Oceananigans.Operators
+using Oceananigans.BoundaryConditions
+using Oceananigans.Grids: AbstractGrid, topology
+using Oceananigans.ImmersedBoundaries: active_linear_index_to_tuple, mask_immersed_field!
+using Oceananigans.Models.HydrostaticFreeSurfaceModels: AbstractFreeSurface,
+                                                        free_surface_displacement_field
+
+using Adapt
+using Base
+using KernelAbstractions: @index, @kernel
+using KernelAbstractions.Extras.LoopInfo: @unroll
+
+import Oceananigans.Models.HydrostaticFreeSurfaceModels: initialize_free_surface!,
+                                                         setup_free_surface!,
+                                                         materialize_free_surface,
+                                                         ab2_step_free_surface!,
+                                                         compute_free_surface_tendency!,
+                                                         explicit_barotropic_pressure_x_gradient,
+                                                         explicit_barotropic_pressure_y_gradient
+
+include("split_explicit_timesteppers.jl")
+include("split_explicit_free_surface.jl")
+include("distributed_split_explicit_free_surface.jl")
+include("initialize_split_explicit_substepping.jl")
+include("compute_slow_tendencies.jl")
+include("step_split_explicit_free_surface.jl")
+include("barotropic_split_explicit_corrector.jl")
+
+# extend
+@inline explicit_barotropic_pressure_x_gradient(i, j, k, grid, ::SplitExplicitFreeSurface) = zero(grid)
+@inline explicit_barotropic_pressure_y_gradient(i, j, k, grid, ::SplitExplicitFreeSurface) = zero(grid)
+
+end
diff --git a/...ostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl b/...ostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl
@@ -0,0 +1,60 @@
+# Kernels to compute the vertical integral of the velocities
+@kernel function _barotropic_mode_kernel!(U, V, grid, ::Nothing, u, v)
+    i, j  = @index(Global, NTuple)
+    barotropic_mode_kernel!(U, V, i, j, grid, u, v)
+end
+
+@kernel function _barotropic_mode_kernel!(U, V, grid, active_cells_map, u, v)
+    idx = @index(Global, Linear)
+    i, j = active_linear_index_to_tuple(idx, active_cells_map)
+    barotropic_mode_kernel!(U, V, i, j, grid, u, v)
+end
+
+@inline function barotropic_mode_kernel!(U, V, i, j, grid, u, v)
+    @inbounds U[i, j, 1] = Δzᶠᶜᶜ(i, j, 1, grid) * u[i, j, 1]
+    @inbounds V[i, j, 1] = Δzᶜᶠᶜ(i, j, 1, grid) * v[i, j, 1]
+
+    for k in 2:grid.Nz
+        @inbounds U[i, j, 1] += Δzᶠᶜᶜ(i, j, k, grid) * u[i, j, k]
+        @inbounds V[i, j, 1] += Δzᶜᶠᶜ(i, j, k, grid) * v[i, j, k]
+    end
+
+    return nothing
+end
+
+@inline function compute_barotropic_mode!(U, V, grid, u, v)
+    active_cells_map = retrieve_surface_active_cells_map(grid)
+    launch!(architecture(grid), grid, :xy, _barotropic_mode_kernel!, U, V, grid, active_cells_map, u, v; active_cells_map)
+    return nothing
+end
+
+@kernel function _barotropic_split_explicit_corrector!(u, v, U, V, U̅, V̅, grid)
+    i, j, k = @index(Global, NTuple)
+
+    @inbounds begin
+        Hᶠᶜ = static_column_depthᶠᶜᵃ(i, j, grid)
+        Hᶜᶠ = static_column_depthᶜᶠᵃ(i, j, grid)
+
+        u[i, j, k] = u[i, j, k] + (U[i, j, 1] - U̅[i, j, 1]) / Hᶠᶜ
+        v[i, j, k] = v[i, j, k] + (V[i, j, 1] - V̅[i, j, 1]) / Hᶜᶠ
+    end
+end
+
+# Correcting `u` and `v` with the barotropic mode computed in `free_surface`
+function barotropic_split_explicit_corrector!(u, v, free_surface, grid)
+    state = free_surface.filtered_state
+    U, V  = free_surface.barotropic_velocities
+    U̅, V̅  = state.U, state.V
+    arch  = architecture(grid)
+
+    # NOTE: the filtered `U̅` and `V̅` have been copied in the instantaneous `U` and `V`,
+    # so we use the filtered velocities as "work arrays" to store the vertical integrals
+    # of the instantaneous velocities `u` and `v`.
+    compute_barotropic_mode!(U̅, V̅, grid, u, v)
+
+    # add in "good" barotropic mode
+    launch!(arch, grid, :xyz, _barotropic_split_explicit_corrector!,
+            u, v, U, V, U̅, V̅, grid)
+
+    return nothing
+end
diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl
@@ -0,0 +1,69 @@
+
+# Calculate RHS for the barotropic time step.
+@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, ::Nothing, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+    i, j  = @index(Global, NTuple)
+    ab2_integrate_tendencies!(Gᵁ, Gⱽ, i, j, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+end
+
+@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+    idx = @index(Global, Linear)
+    i, j = active_linear_index_to_tuple(idx, active_cells_map)
+    ab2_integrate_tendencies!(Gᵁ, Gⱽ, i, j, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+end
+
+@inline function ab2_integrate_tendencies!(Gᵁ, Gⱽ, i, j, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+    locU = (Face(), Center(), Center())
+    locV = (Center(), Face(), Center())
+
+    @inbounds Gᵁ[i, j, 1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_G(i, j, 1, grid, locU..., Gu⁻, Guⁿ, χ)
+    @inbounds Gⱽ[i, j, 1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_G(i, j, 1, grid, locV..., Gv⁻, Gvⁿ, χ)
+
+    for k in 2:grid.Nz
+        @inbounds Gᵁ[i, j, 1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_G(i, j, k, grid, locU..., Gu⁻, Guⁿ, χ)
+        @inbounds Gⱽ[i, j, 1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_G(i, j, k, grid, locV..., Gv⁻, Gvⁿ, χ)
+    end
+end
+
+@inline function ab2_step_G(i, j, k, grid, ℓx, ℓy, ℓz, G⁻, Gⁿ, χ::FT) where FT 
+    C₁ = convert(FT, 3/2) + χ
+    C₂ = convert(FT, 1/2) + χ
+
+    # multiply G⁻ by false if C₂ is zero to 
+    # prevent propagationg possible NaNs
+    not_euler = C₂ != 0
+
+    Gⁿ⁺¹ = @inbounds C₁ * Gⁿ[i, j, k] - C₂ * G⁻[i, j, k] * not_euler
+    immersed = peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz)
+
+    return ifelse(immersed, zero(grid), Gⁿ⁺¹)
+end
+
+# Setting up the RHS for the barotropic step (tendencies of the barotropic velocity components)
+# This function is called after `calculate_tendency` and before `ab2_step_velocities!`
+function compute_free_surface_tendency!(grid, model, ::SplitExplicitFreeSurface)
+
+    # we start the time integration of η from the average ηⁿ
+    Gu⁻ = model.timestepper.G⁻.u
+    Gv⁻ = model.timestepper.G⁻.v
+    Guⁿ = model.timestepper.Gⁿ.u
+    Gvⁿ = model.timestepper.Gⁿ.v
+
+    GUⁿ = model.timestepper.Gⁿ.U
+    GVⁿ = model.timestepper.Gⁿ.V
+
+    @apply_regionally compute_free_surface_forcing!(GUⁿ, GVⁿ, model.grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, model.timestepper.χ)
+
+    fields_to_fill = (GUⁿ, GVⁿ)
+    fill_halo_regions!(fields_to_fill; async = true)
+
+    return nothing
+end
+
+@inline function compute_free_surface_forcing!(GUⁿ, GVⁿ, grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+    active_cells_map = retrieve_surface_active_cells_map(grid)
+
+    launch!(architecture(grid), grid, :xy, _compute_integrated_ab2_tendencies!, GUⁿ, GVⁿ, grid,
+            active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ; active_cells_map)
+
+    return nothing
+end
diff --git a/...ticFreeSurfaceModels/SplitExplicitFreeSurfaces/distributed_split_explicit_free_surface.jl b/...ticFreeSurfaceModels/SplitExplicitFreeSurfaces/distributed_split_explicit_free_surface.jl
@@ -0,0 +1,25 @@
+using Oceananigans.DistributedComputations: DistributedField
+using Oceananigans.DistributedComputations: SynchronizedDistributed, synchronize_communication!
+
+const DistributedSplitExplicit = SplitExplicitFreeSurface{<:DistributedField}
+
+wait_free_surface_communication!(free_surface, model, arch) = nothing
+wait_free_surface_communication!(::DistributedSplitExplicit, model, ::SynchronizedDistributed) = nothing
+
+function wait_free_surface_communication!(free_surface::DistributedSplitExplicit, model, arch)
+
+    barotropic_velocities = free_surface.barotropic_velocities
+
+    for field in (barotropic_velocities.U, barotropic_velocities.V)
+        synchronize_communication!(field)
+    end
+
+    Gᵁ = model.timestepper.Gⁿ.U
+    Gⱽ = model.timestepper.Gⁿ.V
+
+    for field in (Gᵁ, Gⱽ)
+        synchronize_communication!(field)
+    end
+
+    return nothing
+end
diff --git a/...taticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl b/...taticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl
@@ -0,0 +1,31 @@
+using Oceananigans.ImmersedBoundaries: retrieve_surface_active_cells_map, peripheral_node
+
+# This file contains two different initializations methods performed at different stages of the simulation.
+#
+# - `initialize_free_surface!`: the first initialization, performed only once at the beginning of the simulation, 
+#                               calculates the barotropic velocities from the velocity initial conditions.
+#
+# - `initialize_free_surface_state!`: is performed at the beginning of the substepping procedure, resets the filtered state to zero
+#                                     and reinitializes the timestepper auxiliaries from the previous filtered state.           
+
+# `initialize_free_surface!` is called at the beginning of the simulation to initialize the free surface state
+# from the initial velocity conditions.
+function initialize_free_surface!(sefs::SplitExplicitFreeSurface, grid, velocities)
+    barotropic_velocities = sefs.barotropic_velocities
+    @apply_regionally compute_barotropic_mode!(barotropic_velocities.U, barotropic_velocities.V, grid, velocities.u, velocities.v)
+    fill_halo_regions!((barotropic_velocities.U, barotropic_velocities.V))
+    return nothing
+end
+
+# `initialize_free_surface_state!` is called at the beginning of the substepping to 
+# reset the filtered state to zero and reinitialize the state from the filtered state.
+function initialize_free_surface_state!(filtered_state, η, velocities, timestepper)
+
+    initialize_free_surface_timestepper!(timestepper, η, velocities)
+
+    fill!(filtered_state.η, 0)
+    fill!(filtered_state.U, 0)
+    fill!(filtered_state.V, 0)
+
+    return nothing
+end