(0.94.1) Refactor SplitExplicitFreeSurface

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"

docs/src/appendix/library.md

@@ -122,6 +122,13 @@ Modules = [Oceananigans.Models.HydrostaticFreeSurfaceModels]
 Private = false
 ```
 
+### Split-explicit free-surface solver
+
+```@autodocs
+Modules = [Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces]
+Private = false
+```
+
 ### Shallow-water models
 
 ```@autodocs

src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl

@@ -31,8 +31,14 @@ 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
+setup_free_surface!(model, free_surface, χ) = 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 +52,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 +91,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 +108,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),

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl

@@ -0,0 +1,38 @@
+module SplitExplicitFreeSurfaces
+
+export SplitExplicitFreeSurface, ForwardBackwardScheme, AdamsBashforth3Scheme
+export FixedSubstepNumber, FixedTimeStepSize
+
+using Oceananigans
+using Oceananigans.Architectures
+using Oceananigans.Fields
+using Oceananigans.Utils
+using Oceananigans.Grids
+using Oceananigans.Operators
+using Oceananigans.Grids: AbstractGrid
+using Oceananigans.Models.HydrostaticFreeSurfaceModels: AbstractFreeSurface,
+                                                        free_surface_displacement_field
+
+using Adapt
+using Base
+using KernelAbstractions: @index, @kernel
+
+import Oceananigans.Models.HydrostaticFreeSurfaceModels: initialize_free_surface!,
+                                                         setup_free_surface!,
+                                                         materialize_free_surface,
+                                                         ab2_step_free_surface!,
+                                                         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("setup_split_explicit.jl")
+include("step_split_explicit_free_surface.jl")
+include("barotropic_correction.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

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_correction.jl

@@ -0,0 +1,58 @@
+# 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
+
+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)
+
+    # take out "bad" barotropic mode, 
+    # !!!! reusing U̅ and V̅ for U and V have already been restarted
+    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

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/distributed_split_explicit_free_surface.jl

@@ -0,0 +1,27 @@
+using Oceananigans.Grids: with_halo
+using Oceananigans.AbstractOperations: GridMetricOperation, Δz
+using Oceananigans.DistributedComputations: DistributedGrid, 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

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/setup_split_explicit.jl

@@ -0,0 +1,84 @@
+# `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
+
+# 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)
+
+    @inbounds Gᵁ[i, j, 1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_Gu(i, j, 1, grid, Gu⁻, Guⁿ, χ)
+    @inbounds Gⱽ[i, j, 1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_Gv(i, j, 1, grid, Gv⁻, Gvⁿ, χ)
+
+    for k in 2:grid.Nz
+        @inbounds Gᵁ[i, j, 1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_Gu(i, j, k, grid, Gu⁻, Guⁿ, χ)
+        @inbounds Gⱽ[i, j, 1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_Gv(i, j, k, grid, Gv⁻, Gvⁿ, χ)
+    end
+end
+
+# Calculate RHS for the barotropic time step.q
+@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)
+
+    @inbounds Gᵁ[i, j, 1] = Δzᶠᶜᶜ(i, j, 1, grid) * ab2_step_Gu(i, j, 1, grid, Gu⁻, Guⁿ, χ)
+    @inbounds Gⱽ[i, j, 1] = Δzᶜᶠᶜ(i, j, 1, grid) * ab2_step_Gv(i, j, 1, grid, Gv⁻, Gvⁿ, χ)
+
+    for k in 2:grid.Nz
+        @inbounds Gᵁ[i, j, 1] += Δzᶠᶜᶜ(i, j, k, grid) * ab2_step_Gu(i, j, k, grid, Gu⁻, Guⁿ, χ)
+        @inbounds Gⱽ[i, j, 1] += Δzᶜᶠᶜ(i, j, k, grid) * ab2_step_Gv(i, j, k, grid, Gv⁻, Gvⁿ, χ)
+    end
+end
+
+@inline ab2_step_Gu(i, j, k, grid, G⁻, Gⁿ, χ::FT) where FT =
+    @inbounds ifelse(peripheral_node(i, j, k, grid, f, c, c), zero(grid), (convert(FT, 1.5) + χ) *  Gⁿ[i, j, k] - G⁻[i, j, k] * (convert(FT, 0.5) + χ))
+
+@inline ab2_step_Gv(i, j, k, grid, G⁻, Gⁿ, χ::FT) where FT =
+    @inbounds ifelse(peripheral_node(i, j, k, grid, c, f, c), zero(grid), (convert(FT, 1.5) + χ) *  Gⁿ[i, j, k] - G⁻[i, j, k] * (convert(FT, 0.5) + χ))
+
+# 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 setup_free_surface!(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 setup_split_explicit_tendency!(GUⁿ, GVⁿ, model.grid, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
+
+    fields_to_fill = (GUⁿ, GVⁿ)
+    fill_halo_regions!(fields_to_fill; async = true)
+
+    return nothing
+end
+
+@inline function setup_split_explicit_tendency!(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