Reformulate hydrostatic model timestepping

Project.toml

@@ -117,4 +117,4 @@ TimesDates = "bdfc003b-8df8-5c39-adcd-3a9087f5df4a"
 oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 
 [targets]
-test = ["AMDGPU", "CUDA", "oneAPI", "DataDeps", "SafeTestsets", "Test", "Enzyme", "Reactant", "Metal", "CUDA_Runtime_jll", "MPIPreferences", "TimesDates", "NCDatasets"]
+test = ["AMDGPU", "CUDA", "oneAPI", "DataDeps", "SafeTestsets", "Test", "Enzyme", "Reactant", "Metal", "CUDA_Runtime_jll", "MPIPreferences", "TimesDates", "NCDatasets"]

src/DistributedComputations/distributed_grids.jl

@@ -320,12 +320,14 @@ end
 function scatter_local_grids(global_grid::RectilinearGrid, arch::Distributed, local_size)
     x, y, z, topo, halo = scatter_grid_properties(global_grid)
     global_sz = global_size(arch, local_size)
+    global_sz = pop_flat_elements(global_sz, topo)
     return RectilinearGrid(arch, eltype(global_grid); size=global_sz, x=x, y=y, z=z, halo=halo, topology=topo)
 end
 
 function scatter_local_grids(global_grid::LatitudeLongitudeGrid, arch::Distributed, local_size)
     x, y, z, topo, halo = scatter_grid_properties(global_grid)
     global_sz = global_size(arch, local_size)
+    global_sz = pop_flat_elements(global_sz, topo)
     return LatitudeLongitudeGrid(arch, eltype(global_grid); size=global_sz, longitude=x,
                                  latitude=y, z=z, halo=halo, topology=topo, radius=global_grid.radius)
 end

src/ImmersedBoundaries/mutable_immersed_grid.jl

@@ -28,6 +28,16 @@ const MutableGridOfSomeKind = Union{MutableImmersedGrid, AbstractMutableGrid}
 @inline column_depthᶠᶜᵃ(i, j, grid::MutableGridOfSomeKind) = column_depthᶠᶜᵃ(i, j, 1, grid, grid.z.ηⁿ)
 @inline column_depthᶠᶠᵃ(i, j, grid::MutableGridOfSomeKind) = column_depthᶠᶠᵃ(i, j, 1, grid, grid.z.ηⁿ)
 
+@inline column_depthᶜᶜᵃ(i, j, k, grid) = static_column_depthᶜᶜᵃ(i, j, grid)
+@inline column_depthᶜᶠᵃ(i, j, k, grid) = static_column_depthᶜᶠᵃ(i, j, grid)
+@inline column_depthᶠᶜᵃ(i, j, k, grid) = static_column_depthᶠᶜᵃ(i, j, grid)
+@inline column_depthᶠᶠᵃ(i, j, k, grid) = static_column_depthᶠᶠᵃ(i, j, grid)
+
+@inline column_depthᶜᶜᵃ(i, j, k, grid::MutableGridOfSomeKind) = column_depthᶜᶜᵃ(i, j, 1, grid, grid.z.ηⁿ)
+@inline column_depthᶜᶠᵃ(i, j, k, grid::MutableGridOfSomeKind) = column_depthᶜᶠᵃ(i, j, 1, grid, grid.z.ηⁿ)
+@inline column_depthᶠᶜᵃ(i, j, k, grid::MutableGridOfSomeKind) = column_depthᶠᶜᵃ(i, j, 1, grid, grid.z.ηⁿ)
+@inline column_depthᶠᶠᵃ(i, j, k, grid::MutableGridOfSomeKind) = column_depthᶠᶠᵃ(i, j, 1, grid, grid.z.ηⁿ)
+
 # Fallbacks
 @inline σⁿ(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = σⁿ(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz)
 @inline σ⁻(i, j, k, ibg::IBG, ℓx, ℓy, ℓz) = σ⁻(i, j, k, ibg.underlying_grid, ℓx, ℓy, ℓz)

src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl

@@ -19,18 +19,50 @@ import Oceananigans: fields, prognostic_fields, initialize!
 import Oceananigans.Advection: cell_advection_timescale
 import Oceananigans.TimeSteppers: step_lagrangian_particles!
 import Oceananigans.Architectures: on_architecture
+import Oceananigans.BoundaryConditions: fill_halo_regions!
 
 using Oceananigans.TimeSteppers: SplitRungeKutta3TimeStepper, QuasiAdamsBashforth2TimeStepper
 
 abstract type AbstractFreeSurface{E, G} end
 
 struct ZCoordinate end
 
-struct ZStarCoordinate{CC}
+fill_halo_regions!(ztype::ZCoordinate; kwargs...) = nothing
+
+struct ZStarCoordinate{CC} #, FC, CF, FF, H}
     storage :: CC # Storage space used in different ways by different timestepping schemes.
+    # σᶜᶜⁿ :: CC
+    # σᶠᶜⁿ :: FC
+    # σᶜᶠⁿ :: CF
+    # σᶠᶠⁿ :: FF
+    # σᶜᶜ⁻ :: CC
+    # ∂t_σ :: CC
+    # ηⁿ   :: H
+end
+
+function ZStarCoordinate(grid::AbstractGrid)
+    storage = Field{Center, Center, Nothing}(grid)
+    # σᶜᶜⁿ    = Field{Center, Center, Nothing}(grid; data=grid.z.σᶜᶜⁿ)
+    # σᶠᶜⁿ    = Field{Face,   Center, Nothing}(grid; data=grid.z.σᶠᶜⁿ)
+    # σᶜᶠⁿ    = Field{Center, Face,   Nothing}(grid; data=grid.z.σᶜᶠⁿ)
+    # σᶠᶠⁿ    = Field{Face,   Face,   Nothing}(grid; data=grid.z.σᶠᶠⁿ)
+    # σᶜᶜ⁻    = Field{Center, Center, Nothing}(grid; data=grid.z.σᶜᶜ⁻)
+    # ∂t_σ    = Field{Center, Center, Nothing}(grid; data=grid.z.∂t_σ)
+    # ηⁿ      = Field{Center, Center, Nothing}(grid; data=grid.z.ηⁿ  )
+    return ZStarCoordinate(storage) #, σᶜᶜⁿ, σᶠᶜⁿ, σᶜᶠⁿ, σᶠᶠⁿ, σᶜᶜ⁻, ∂t_σ, ηⁿ)
 end
 
-ZStarCoordinate(grid::AbstractGrid) = ZStarCoordinate(Field{Center, Center, Nothing}(grid))
+function fill_halo_regions!(ztype::ZStarCoordinate; kwargs...)
+    fill_halo_regions!(ztype.storage; kwargs...)
+    fill_halo_regions!(ztype.σᶜᶜⁿ; kwargs...)
+    fill_halo_regions!(ztype.σᶠᶜⁿ; kwargs...)
+    fill_halo_regions!(ztype.σᶜᶠⁿ; kwargs...)
+    fill_halo_regions!(ztype.σᶠᶠⁿ; kwargs...)
+    fill_halo_regions!(ztype.σᶜᶜ⁻; kwargs...)
+    fill_halo_regions!(ztype.∂t_σ; kwargs...)
+    fill_halo_regions!(ztype.ηⁿ; kwargs...)
+    return nothing
+end
 
 Base.summary(::ZStarCoordinate) = "ZStarCoordinate"
 Base.show(io::IO, c::ZStarCoordinate) = print(io, summary(c))
@@ -50,8 +82,31 @@ free_surface_displacement_field(velocities, ::Nothing, grid) = nothing
 
 # free surface initialization functions
 initialize_free_surface!(free_surface, grid, velocities) = nothing
+compute_transport_velocities!(model, free_surface) = nothing    
+
+# If U and V exist, we use them
+@inline barotropic_U(i, j, k, grid, U, u) = @inbounds U[i, j, k]
+@inline barotropic_V(i, j, k, grid, V, v) = @inbounds V[i, j, k]
+
+# If either U or V are not available, we compute them
+@inline function barotropic_U(i, j, k, grid, ::Nothing, u)
+    U = 0
+    for k′ in 1:size(grid, 3)
+        @inbounds U += u[i, j, k′] * Δzᶠᶜᶜ(i, j, k′, grid)
+    end
+    return U
+end
+
+@inline function barotropic_V(i, j, k, grid, ::Nothing, v)
+    V = 0
+    for k′ in 1:size(grid, 3)
+        @inbounds V += v[i, j, k′] * Δzᶜᶠᶜ(i, j, k′, grid)
+    end
+    return V
+end
 
 include("compute_w_from_continuity.jl")
+include("hydrostatic_free_surface_field_tuples.jl")
 
 # No free surface
 include("nothing_free_surface.jl")
@@ -71,6 +126,12 @@ include("hydrostatic_free_surface_field_tuples.jl")
 include("SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl")
 using .SplitExplicitFreeSurfaces
 
+# Implicit free-surface solver functionality
+include("implicit_free_surface_utils.jl")
+include("fft_based_implicit_free_surface_solver.jl")
+include("pcg_implicit_free_surface_solver.jl")
+include("implicit_free_surface.jl")
+
 # ZStarCoordinate implementation
 include("z_star_vertical_spacing.jl")
 

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl

@@ -15,7 +15,8 @@ using Oceananigans.ImmersedBoundaries
 using Oceananigans.Grids: AbstractGrid, topology
 using Oceananigans.ImmersedBoundaries: linear_index_to_tuple, mask_immersed_field!
 using Oceananigans.Models.HydrostaticFreeSurfaceModels: AbstractFreeSurface,
-                                                        free_surface_displacement_field
+                                                        free_surface_displacement_field,
+                                                        update_vertical_velocities!
 
 using Adapt
 using Base
@@ -31,6 +32,7 @@ import Oceananigans.Models.HydrostaticFreeSurfaceModels: initialize_free_surface
                                                          materialize_free_surface,
                                                          step_free_surface!,
                                                          compute_free_surface_tendency!,
+                                                         compute_transport_velocities!,
                                                          explicit_barotropic_pressure_x_gradient,
                                                          explicit_barotropic_pressure_y_gradient
 

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/barotropic_split_explicit_corrector.jl

@@ -27,7 +27,7 @@ function barotropic_split_explicit_corrector!(u, v, free_surface, grid)
     state = free_surface.filtered_state
     η     = free_surface.η
     U, V  = free_surface.barotropic_velocities
-    U̅, V̅  = state.U, state.V
+    U̅, V̅  = state.U̅, state.V̅
     arch  = architecture(grid)
 
     # NOTE: the filtered `U̅` and `V̅` have been copied in the instantaneous `U` and `V`,
@@ -54,4 +54,37 @@ end
     end
 end
 
+@kernel function _compute_transport_velocities!(ũ, ṽ, grid, Ũ, Ṽ, u, v, U̅, V̅)
+    i, j = @index(Global, NTuple)
+
+    for k in 1:size(grid, 3)
+        @inline ũ[i, j, k] = u[i, j, k] + (Ũ[i, j, 1] - U̅[i, j, 1]) / column_depthᶠᶜᵃ(i, j, grid)
+        @inline ṽ[i, j, k] = v[i, j, k] + (Ṽ[i, j, 1] - V̅[i, j, 1]) / column_depthᶜᶠᵃ(i, j, grid)
+    end
+end
+
+function compute_transport_velocities!(model, free_surface::SplitExplicitFreeSurface)
+    grid = model.grid
+    u, v, _ = model.velocities
+    ũ, ṽ, _ = model.transport_velocities
+    Ũ = free_surface.filtered_state.Ũ
+    Ṽ = free_surface.filtered_state.Ṽ
+    U̅ = free_surface.filtered_state.U̅
+    V̅ = free_surface.filtered_state.V̅
+
+    compute_barotropic_mode!(U̅, V̅, grid, u, v)
 
+    launch!(architecture(grid), grid, :xy,
+            _compute_transport_velocities!, ũ, ṽ, grid, Ũ, Ṽ, u, v, U̅, V̅)
+
+    mask_immersed_field!(ũ)
+    mask_immersed_field!(ṽ)
+
+    # Fill barotropic stuff...
+    fill_halo_regions!((ũ, ṽ); async=true)
+
+    # Update grid velocity and vertical transport velocity
+    update_vertical_velocities!(model.transport_velocities, model.grid, model)
+
+    return nothing
+end

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl

@@ -67,14 +67,9 @@ end
     @inbounds GVⁿ[i, j, 1] = G_vertical_integral(i, j, grid, Gvⁿ, Center(), Face(), Center())
 end
 
-@inline function compute_split_explicit_forcing!(GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ, timestepper::SplitRungeKutta3TimeStepper)
-
-    active_cells_map = get_active_column_map(grid)    
+@inline compute_split_explicit_forcing!(GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ, ::SplitRungeKutta3TimeStepper) = 
     launch!(architecture(grid), grid, :xy, _compute_integrated_rk3_tendencies!, 
-            GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ; active_cells_map)
-
-    return nothing
-end
+            GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ; active_cells_map = get_active_column_map(grid))
 
 #####
 ##### Free surface setup
@@ -93,9 +88,7 @@ function compute_free_surface_tendency!(grid, model, ::SplitExplicitFreeSurface)
     baroclinic_timestepper = model.timestepper
 
     @apply_regionally compute_split_explicit_forcing!(GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ, baroclinic_timestepper)
-
-    fields_to_fill = (GUⁿ, GVⁿ)
-    fill_halo_regions!(fields_to_fill; async=true)
+    fill_halo_regions!((GUⁿ, GVⁿ); async=true) 
 
     return nothing
 end

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl

@@ -34,9 +34,11 @@ function initialize_free_surface_state!(free_surface, baroclinic_timestepper, ti
 
     initialize_free_surface_timestepper!(timestepper, η, U, V)
 
-    fill!(free_surface.filtered_state.η, 0)
-    fill!(free_surface.filtered_state.U, 0)
-    fill!(free_surface.filtered_state.V, 0)
+    fill!(free_surface.filtered_state.η̅, 0)
+    fill!(free_surface.filtered_state.U̅, 0)
+    fill!(free_surface.filtered_state.V̅, 0)
+    fill!(free_surface.filtered_state.Ũ, 0)
+    fill!(free_surface.filtered_state.Ṽ, 0)
 
     return nothing
 end
@@ -58,9 +60,11 @@ function initialize_free_surface_state!(free_surface, baroclinic_ts::SplitRungeK
 
     initialize_free_surface_timestepper!(barotropic_ts, η, U, V)
 
-    fill!(free_surface.filtered_state.η, 0)
-    fill!(free_surface.filtered_state.U, 0)
-    fill!(free_surface.filtered_state.V, 0)
+    fill!(free_surface.filtered_state.η̅, 0)
+    fill!(free_surface.filtered_state.U̅, 0)
+    fill!(free_surface.filtered_state.V̅, 0)
+    fill!(free_surface.filtered_state.Ũ, 0)
+    fill!(free_surface.filtered_state.Ṽ, 0)
 
     return nothing
 end

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl

@@ -88,7 +88,7 @@ function SplitExplicitFreeSurface(grid = nothing;
                                   substeps = nothing,
                                   cfl = nothing,
                                   fixed_Δt = nothing,
-                                  averaging_kernel = averaging_shape_function,
+                                  averaging_kernel = cosine_averaging_kernel, # averaging_shape_function, # 
                                   timestepper = ForwardBackwardScheme())
 
     if !isnothing(grid)
@@ -121,8 +121,8 @@ end
 # Simplest case: we have the substeps and the averaging kernel
 function split_explicit_substepping(::Nothing, substeps, fixed_Δt, grid, averaging_kernel, gravitational_acceleration)
     FT = eltype(gravitational_acceleration)
-    fractional_step_size, averaging_weights = weights_from_substeps(FT, substeps, averaging_kernel)
-    return FixedSubstepNumber(fractional_step_size, averaging_weights)
+    fractional_step_size, averaging_weights, transport_weights = weights_from_substeps(FT, substeps, averaging_kernel)
+    return FixedSubstepNumber(fractional_step_size, averaging_weights, transport_weights)
 end
 
 # The substeps are calculated dynamically when a cfl without a fixed_Δt is provided
@@ -206,8 +206,10 @@ function materialize_free_surface(free_surface::SplitExplicitFreeSurface, veloci
 
     U̅ = Field{Face, Center, Nothing}(maybe_extended_grid, boundary_conditions = u_bcs)
     V̅ = Field{Center, Face, Nothing}(maybe_extended_grid, boundary_conditions = v_bcs)
+    Ũ = Field{Face, Center, Nothing}(maybe_extended_grid, boundary_conditions = u_bcs)
+    Ṽ = Field{Center, Face, Nothing}(maybe_extended_grid, boundary_conditions = v_bcs)
 
-    filtered_state = (η = η̅, U = U̅, V = V̅)
+    filtered_state = (η̅ = η̅, U̅ = U̅, V̅ = V̅, Ũ = Ũ, Ṽ = Ṽ)
     barotropic_velocities = (U = U, V = V)
 
     kernel_parameters = maybe_augmented_kernel_parameters(TX, TY, substepping, maybe_extended_grid)
@@ -225,9 +227,8 @@ function materialize_free_surface(free_surface::SplitExplicitFreeSurface, veloci
 end
 
 # (p = 2, q = 4, r = 0.18927) minimize dispersion error from Shchepetkin and McWilliams (2005): https://doi.org/10.1016/j.ocemod.2004.08.002
-@inline function averaging_shape_function(τ::FT; p = 2, q = 4, r = FT(0.18927)) where FT
+@inline function averaging_shape_function(τ::FT; p = 2, q = 4, r = FT(0.18927)) where FT #0.12)) where FT
     τ₀ = (p + 2) * (p + q + 2) / (p + 1) / (p + q + 1)
-
     return (τ / τ₀)^p * (1 - (τ / τ₀)^q) - r * (τ / τ₀)
 end
 
@@ -246,6 +247,7 @@ a fixed number of substeps with time step size of `fractional_step_size * Δt_ba
 struct FixedSubstepNumber{B, F}
     fractional_step_size :: B
     averaging_weights    :: F
+    transport_weights    :: F
 end
 
 function FixedTimeStepSize(grid;
@@ -267,18 +269,26 @@ function FixedTimeStepSize(grid;
 end
 
 @inline function weights_from_substeps(FT, substeps, averaging_kernel)
-
+    M  = substeps ÷ 2
     τᶠ = range(FT(0), FT(2), length = substeps+1)
     Δτ = τᶠ[2] - τᶠ[1]
 
     averaging_weights = map(averaging_kernel, τᶠ[2:end])
-    idx = searchsortedlast(averaging_weights, 0, rev=true)
-    substeps = idx
+    M★ = substeps
+
+    # Find the latest allowable weight
+    for i in substeps:-1:1
+        if averaging_weights[i] > 0
+            M★ = i
+            break
+        end
+    end
 
-    averaging_weights = averaging_weights[1:idx]
+    averaging_weights = averaging_weights[1:M★]
     averaging_weights ./= sum(averaging_weights)
+    transport_weights = [sum(averaging_weights[i:M★]) for i in 1:M★] ./ M
 
-    return Δτ, map(FT, tuple(averaging_weights...))
+    return Δτ, map(FT, tuple(averaging_weights...)), map(FT, tuple(transport_weights...))
 end
 
 Base.summary(s::FixedTimeStepSize)  = string("FixedTimeStepSize($(prettytime(s.Δt_barotropic)))")
@@ -299,7 +309,7 @@ function maybe_extend_halos(TX, TY, grid, substepping::FixedSubstepNumber)
 
     old_halos = halo_size(grid)
     Nsubsteps = length(substepping.averaging_weights)
-    step_halo = Nsubsteps + 1
+    step_halo = Nsubsteps+2
 
     Hx = TX() isa ConnectedTopology ? max(step_halo, old_halos[1]) : old_halos[1]
     Hy = TY() isa ConnectedTopology ? max(step_halo, old_halos[2]) : old_halos[2]

src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl

@@ -141,12 +141,14 @@ end
 @inline η★(i, j, k, grid,  ::ForwardBackwardScheme, ηᵐ⁺¹) = @inbounds ηᵐ⁺¹[i, j, k]
 @inline η★(i, j, k, grid, t::AdamsBashforth3Scheme, ηᵐ⁺¹) = @inbounds t.δ * ηᵐ⁺¹[i, j, k] + t.μ * t.ηᵐ[i, j, k] + t.γ * t.ηᵐ⁻¹[i, j, k] + t.ϵ * t.ηᵐ⁻²[i, j, k]
 
-@inline cache_previous_velocities!(::ForwardBackwardScheme,   i, j, k, U) = nothing
-@inline cache_previous_free_surface!(::ForwardBackwardScheme, i, j, k, η) = nothing
+@inline   cache_previous_velocities!(::ForwardBackwardScheme, i, j, k, U, V) = nothing
+@inline cache_previous_free_surface!(::ForwardBackwardScheme, i, j, k, η)    = nothing
 
-@inline function cache_previous_velocities!(t::AdamsBashforth3Scheme, i, j, k, U)
+@inline function cache_previous_velocities!(t::AdamsBashforth3Scheme, i, j, k, U, V)
     @inbounds t.Uᵐ⁻²[i, j, k] = t.Uᵐ⁻¹[i, j, k]
     @inbounds t.Uᵐ⁻¹[i, j, k] =      U[i, j, k]
+    @inbounds t.Vᵐ⁻²[i, j, k] = t.Vᵐ⁻¹[i, j, k]
+    @inbounds t.Vᵐ⁻¹[i, j, k] =      V[i, j, k]
 
     return nothing
 end