Skip to content

Commit

Permalink
Revisit SplitExplicitFreeSurface user-interface and Docs + use `Spl…
Browse files Browse the repository at this point in the history
…itExplicitFreeSurface` as the default when appropriate (#3503)

* make SplitExplicitFreeSurface the default

* promote grid to positional arg in SplitExplicitFreeSurface

* code cleanup

* bit more code cleanup

* minor code beautifications

* simpler split-explicit constructors

* use default_free_surface

* default g that of Earth's

* fix show + code beautifications

* fix validation script

* code cleanup

* code align

* add summary in show method

* grid is arg

* split-explicit is the default

* fix Barotropic Kernels test

* fix Barotropic Correction tests

* clean up

* FFTImplicitFreeSurface only for regular xy rectilinear grids

* few fixes

* SplitExplicitSettings become reproducible by carrying their own settings kwargs

* fix split-explicit tests

* comment on what is the default free surface

* bit of cleanup

* cleaner default_free_surface

* a bit more on required args/kwargs

* add more info regarding averaging_kernel

* bit of cleanup

* small bugfix remaining from on_architecture

* forgot to give grid

* code cleanup

* Update src/Models/HydrostaticFreeSurfaceModels/split_explicit_free_surface.jl

Co-authored-by: Gregory L. Wagner <[email protected]>

* couple of bugfixes

* small bugfix

* default free surface for MultiRegionGrid -> Implicit

* code cleanup

---------

Co-authored-by: Simone Silvestri <[email protected]>
Co-authored-by: Gregory L. Wagner <[email protected]>
  • Loading branch information
3 people committed Mar 15, 2024
1 parent dccba69 commit 9f3f8e6
Show file tree
Hide file tree
Showing 34 changed files with 326 additions and 268 deletions.
7 changes: 2 additions & 5 deletions examples/internal_tide.jl
Original file line number Diff line number Diff line change
Expand Up @@ -114,12 +114,9 @@ u_forcing = Forcing(tidal_forcing, parameters=(; tidal_forcing_amplitude, ω₂)

# ## Model

# We built a `HydrostaticFreeSurfaceModel` with a `SplitExplicitFreeSurface` solver.
# We limit our maximum timestep to 3 minutes.
# We built a `HydrostaticFreeSurfaceModel`:

free_surface = SplitExplicitFreeSurface(; grid, cfl = 0.7)

model = HydrostaticFreeSurfaceModel(; grid, coriolis, free_surface,
model = HydrostaticFreeSurfaceModel(; grid, coriolis,
buoyancy = BuoyancyTracer(),
tracers = :b,
momentum_advection = WENO(),
Expand Down
1 change: 0 additions & 1 deletion src/Grids/rectilinear_grid.jl
Original file line number Diff line number Diff line change
Expand Up @@ -492,4 +492,3 @@ const C = Center
@inline zspacings(grid::RG, ℓx, ℓy, ℓz; kwargs...) = zspacings(grid, ℓz; kwargs...)

@inline isrectilinear(::RG) = true

Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,8 @@ fill_horizontal_velocity_halos!(args...) = nothing
##### HydrostaticFreeSurfaceModel definition
#####

FreeSurfaceDisplacementField(velocities, free_surface, grid) = ZFaceField(grid, indices = (:, :, size(grid, 3)+1))
FreeSurfaceDisplacementField(velocities, ::Nothing, grid) = nothing
free_surface_displacement_field(velocities, free_surface, grid) = ZFaceField(grid, indices = (:, :, size(grid, 3)+1))
free_surface_displacement_field(velocities, ::Nothing, grid) = nothing

include("compute_w_from_continuity.jl")
include("rigid_lid.jl")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ using Oceananigans.AbstractOperations: Ax, Ay, GridMetricOperation
function compute_vertically_integrated_lateral_areas!(∫ᶻ_A)

# we have to account for halos when calculating Integrated areas, in case
# a periodic domain, where it is not guaranteed that ηₙ == ηₙ₊₁
# a periodic domain, where it is not guaranteed that ηₙ == ηₙ₊₁
# 2 halos (instead of only 1) are necessary to accomodate the preconditioner

field_grid = ∫ᶻ_A.xᶠᶜᶜ.grid
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,8 @@ Compute the vertical velocity ``w`` by integrating the continuity equation from
w^{n+1} = -∫ [∂/∂x (u^{n+1}) + ∂/∂y (v^{n+1})] dz
```
"""
compute_w_from_continuity!(model; kwargs...) = compute_w_from_continuity!(model.velocities, model.architecture, model.grid; kwargs...)
compute_w_from_continuity!(model; kwargs...) =
compute_w_from_continuity!(model.velocities, model.architecture, model.grid; kwargs...)

compute_w_from_continuity!(velocities, arch, grid; parameters = w_kernel_parameters(grid)) =
launch!(arch, grid, parameters, _compute_w_from_continuity!, velocities, grid)
Expand All @@ -37,11 +38,11 @@ end
Hx, Hy, _ = halo_size(grid)
Tx, Ty, _ = topology(grid)

Sx = Tx == Flat ? Nx : Nx + 2Hx - 2
Sy = Ty == Flat ? Ny : Ny + 2Hy - 2
Sx = Tx == Flat ? Nx : Nx + 2Hx - 2
Sy = Ty == Flat ? Ny : Ny + 2Hy - 2

Ox = Tx == Flat ? 0 : - Hx + 1
Oy = Ty == Flat ? 0 : - Hy + 1
Ox = Tx == Flat ? 0 : - Hx + 1
Oy = Ty == Flat ? 0 : - Hy + 1

return KernelParameters((Sx, Sy), (Ox, Oy))
end
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ using Oceananigans.DistributedComputations: DistributedGrid, DistributedField
using Oceananigans.DistributedComputations: SynchronizedDistributed, synchronize_communication!
using Oceananigans.Models.HydrostaticFreeSurfaceModels: SplitExplicitState, SplitExplicitFreeSurface

import Oceananigans.Models.HydrostaticFreeSurfaceModels: FreeSurface, SplitExplicitAuxiliaryFields
import Oceananigans.Models.HydrostaticFreeSurfaceModels: materialize_free_surface, SplitExplicitAuxiliaryFields

function SplitExplicitAuxiliaryFields(grid::DistributedGrid)

Expand Down Expand Up @@ -59,7 +59,8 @@ end
return (Ax, Ay)
end

function FreeSurface(free_surface::SplitExplicitFreeSurface, velocities, grid::DistributedGrid)
# Internal function for HydrostaticFreeSurfaceModel
function materialize_free_surface(free_surface::SplitExplicitFreeSurface, velocities, grid::DistributedGrid)

settings = free_surface.settings

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,15 +25,15 @@ Adapt.adapt_structure(to, free_surface::ExplicitFreeSurface) =
ExplicitFreeSurface(Adapt.adapt(to, free_surface.η), free_surface.gravitational_acceleration)

on_architecture(to, free_surface::ExplicitFreeSurface) =
ExplicitFreeSurface(on_architecture(to, free_surface.η),
ExplicitFreeSurface(on_architecture(to, free_surface.η),
on_architecture(to, free_surface.gravitational_acceleration))

#####
##### Interface to HydrostaticFreeSurfaceModel
#####

function FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, velocities, grid)
η = FreeSurfaceDisplacementField(velocities, free_surface, grid)
function materialize_free_surface(free_surface::ExplicitFreeSurface{Nothing}, velocities, grid)
η = free_surface_displacement_field(velocities, free_surface, grid)
g = convert(eltype(grid), free_surface.gravitational_acceleration)

return ExplicitFreeSurface(η, g)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -113,4 +113,3 @@ end
δ_Q = flux_div_xyᶜᶜᶠ(i, j, k_top, grid, ∫ᶻQ.u, ∫ᶻQ.v)
@inbounds rhs[i, j, 1] = (δ_Q - Az * η[i, j, k_top] / Δt) / (g * Lz * Δt * Az)
end

Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
using Oceananigans.Fields: location
using Oceananigans.TimeSteppers: ab2_step_field!
using Oceananigans.TurbulenceClosures: implicit_step!

using Oceananigans.ImmersedBoundaries: active_interior_map, active_surface_map

import Oceananigans.TimeSteppers: ab2_step!
Expand Down Expand Up @@ -91,4 +90,3 @@ function ab2_step_tracers!(tracers, model, Δt, χ)

return nothing
end

Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ end
function HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, grid, tracer_names)
u = XFaceField(grid)
v = YFaceField(grid)
η = FreeSurfaceDisplacementField(velocities, free_surface, grid)
η = free_surface_displacement_field(velocities, free_surface, grid)
tracers = TracerFields(tracer_names, grid)
return merge((u=u, v=v, η=η), tracers)
end
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,7 @@ using Oceananigans.BoundaryConditions: regularize_field_boundary_conditions
using Oceananigans.Biogeochemistry: validate_biogeochemistry, AbstractBiogeochemistry, biogeochemical_auxiliary_fields
using Oceananigans.Fields: Field, CenterField, tracernames, VelocityFields, TracerFields
using Oceananigans.Forcings: model_forcing
using Oceananigans.Grids: halo_size
using Oceananigans.Grids: AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid, architecture
using Oceananigans.Grids: AbstractCurvilinearGrid, AbstractHorizontallyCurvilinearGrid, architecture, halo_size
using Oceananigans.ImmersedBoundaries: ImmersedBoundaryGrid
using Oceananigans.Models: AbstractModel, validate_model_halo, NaNChecker, validate_tracer_advection, extract_boundary_conditions
using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state!, AbstractLagrangianParticles
Expand Down Expand Up @@ -48,14 +47,20 @@ mutable struct HydrostaticFreeSurfaceModel{TS, E, A<:AbstractArchitecture, S,
auxiliary_fields :: AF # User-specified auxiliary fields for forcing functions and boundary conditions
end

default_free_surface(grid::XYRegularRG; gravitational_acceleration=g_Earth) =
ImplicitFreeSurface(; gravitational_acceleration)

default_free_surface(grid; gravitational_acceleration=g_Earth) =
SplitExplicitFreeSurface(grid; cfl = 0.7, gravitational_acceleration)

"""
HydrostaticFreeSurfaceModel(; grid,
clock = Clock{eltype(grid)}(0, 0, 1),
momentum_advection = CenteredSecondOrder(),
tracer_advection = CenteredSecondOrder(),
buoyancy = SeawaterBuoyancy(eltype(grid)),
coriolis = nothing,
free_surface = ImplicitFreeSurface(gravitational_acceleration = g_Earth),
free_surface = default_free_surface(grid, gravitational_acceleration=g_Earth),
forcing::NamedTuple = NamedTuple(),
closure = nothing,
boundary_conditions::NamedTuple = NamedTuple(),
Expand All @@ -74,18 +79,22 @@ Keyword arguments
=================
- `grid`: (required) The resolution and discrete geometry on which `model` is solved. The
architecture (CPU/GPU) that the model is solve is inferred from the architecture
of the grid.
architecture (CPU/GPU) that the model is solved is inferred from the architecture
of the `grid`.
- `momentum_advection`: The scheme that advects velocities. See `Oceananigans.Advection`.
- `tracer_advection`: The scheme that advects tracers. See `Oceananigans.Advection`.
- `buoyancy`: The buoyancy model. See `Oceananigans.BuoyancyModels`.
- `coriolis`: Parameters for the background rotation rate of the model.
- `free_surface`: The free surface model. The default free-surface solver depends on the
geometry of the `grid`. If the `grid` is a `RectilinearGrid` that is
regularly spaced in the horizontal the default is an `ImplicitFreeSurface`
solver with `solver_method = :FFTBasedPoissonSolver`. In all other cases,
the default is a `SplitExplicitFreeSurface`.
- `tracers`: A tuple of symbols defining the names of the modeled tracers, or a `NamedTuple` of
preallocated `CenterField`s.
- `forcing`: `NamedTuple` of user-defined forcing functions that contribute to solution tendencies.
- `free_surface`: The free surface model.
- `closure`: The turbulence closure for `model`. See `Oceananigans.TurbulenceClosures`.
- `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.
- `particles`: Lagrangian particles to be advected with the flow. Default: `nothing`.
- `biogeochemistry`: Biogeochemical model for `tracers`.
- `velocities`: The model velocities. Default: `nothing`.
Expand All @@ -99,11 +108,11 @@ function HydrostaticFreeSurfaceModel(; grid,
tracer_advection = CenteredSecondOrder(),
buoyancy = SeawaterBuoyancy(eltype(grid)),
coriolis = nothing,
free_surface = ImplicitFreeSurface(gravitational_acceleration = g_Earth),
forcing::NamedTuple = NamedTuple(),
free_surface = default_free_surface(grid, gravitational_acceleration=g_Earth),
tracers = (:T, :S),
forcing::NamedTuple = NamedTuple(),
closure = nothing,
boundary_conditions::NamedTuple = NamedTuple(),
tracers = (:T, :S),
particles::ParticlesOrNothing = nothing,
biogeochemistry::AbstractBGCOrNothing = nothing,
velocities = nothing,
Expand Down Expand Up @@ -165,7 +174,7 @@ function HydrostaticFreeSurfaceModel(; grid,
@apply_regionally validate_velocity_boundary_conditions(grid, velocities)

free_surface = validate_free_surface(arch, free_surface)
free_surface = FreeSurface(free_surface, velocities, grid)
free_surface = materialize_free_surface(free_surface, velocities, grid)

# Instantiate timestepper if not already instantiated
implicit_solver = implicit_diffusion_solver(time_discretization(closure), grid)
Expand Down
18 changes: 8 additions & 10 deletions src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl
Original file line number Diff line number Diff line change
Expand Up @@ -90,8 +90,8 @@ on_architecture(to, free_surface::ImplicitFreeSurface) =
on_architecture(to, free_surface.solver_settings))

# Internal function for HydrostaticFreeSurfaceModel
function FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, velocities, grid)
η = FreeSurfaceDisplacementField(velocities, free_surface, grid)
function materialize_free_surface(free_surface::ImplicitFreeSurface{Nothing}, velocities, grid)
η = free_surface_displacement_field(velocities, free_surface, grid)
gravitational_acceleration = convert(eltype(grid), free_surface.gravitational_acceleration)

# Initialize barotropic volume fluxes
Expand All @@ -111,13 +111,11 @@ function FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, velocities, gri
free_surface.solver_settings)
end

is_horizontally_regular(grid) = false
is_horizontally_regular(::RectilinearGrid{<:Any, <:Any, <:Any, <:Any, <:Number, <:Number}) = true
build_implicit_step_solver(::Val{:Default}, grid::XYRegularRG, settings, gravitational_acceleration) =
build_implicit_step_solver(Val(:FastFourierTransform), grid, settings, gravitational_acceleration)

function build_implicit_step_solver(::Val{:Default}, grid, settings, gravitational_acceleration)
default_method = is_horizontally_regular(grid) ? :FastFourierTransform : :HeptadiagonalIterativeSolver
return build_implicit_step_solver(Val(default_method), grid, settings, gravitational_acceleration)
end
build_implicit_step_solver(::Val{:Default}, grid, settings, gravitational_acceleration) =
build_implicit_step_solver(Val(:HeptadiagonalIterativeSolver), grid, settings, gravitational_acceleration)

@inline explicit_barotropic_pressure_x_gradient(i, j, k, grid, ::ImplicitFreeSurface) = 0
@inline explicit_barotropic_pressure_y_gradient(i, j, k, grid, ::ImplicitFreeSurface) = 0
Expand All @@ -141,7 +139,7 @@ function implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model,
# Compute right hand side of implicit free surface equation
@apply_regionally local_compute_integrated_volume_flux!(∫ᶻQ, model.velocities, arch)
fill_halo_regions!(∫ᶻQ)

compute_implicit_free_surface_right_hand_side!(rhs, solver, g, Δt, ∫ᶻQ, η)

# Solve for the free surface at tⁿ⁺¹
Expand All @@ -157,7 +155,7 @@ function implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model,
end

function local_compute_integrated_volume_flux!(∫ᶻQ, velocities, arch)

foreach(mask_immersed_field!, velocities)

# Compute barotropic volume flux. Blocking.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ function solve!(η, implicit_free_surface_solver::MatrixImplicitFreeSurfaceSolve
storage = implicit_free_surface_solver.storage

solve!(storage, solver, rhs, Δt)

set!(η, reshape(storage, solver.problem_size...))

return nothing
Expand All @@ -92,7 +92,7 @@ function compute_implicit_free_surface_right_hand_side!(rhs,
launch!(arch, grid, :xy,
implicit_linearized_free_surface_right_hand_side!,
rhs, grid, g, Δt, ∫ᶻQ, η)

return nothing
end

Expand Down Expand Up @@ -130,8 +130,8 @@ end
@kernel function _compute_coefficients!(diag, Ax, Ay, ∫Ax, ∫Ay, grid, g)
i, j = @index(Global, NTuple)
@inbounds begin
Ay[i, j, 1] = ∫Ay[i, j, 1] / Δyᶜᶠᶠ(i, j, grid.Nz+1, grid)
Ax[i, j, 1] = ∫Ax[i, j, 1] / Δxᶠᶜᶠ(i, j, grid.Nz+1, grid)
diag[i, j, 1] = - Azᶜᶜᶠ(i, j, grid.Nz+1, grid) / g
Ay[i, j, 1] = ∫Ay[i, j, 1] / Δyᶜᶠᶠ(i, j, grid.Nz+1, grid)
Ax[i, j, 1] = ∫Ax[i, j, 1] / Δxᶠᶜᶠ(i, j, grid.Nz+1, grid)
diag[i, j, 1] = - Azᶜᶜᶠ(i, j, grid.Nz+1, grid) / g
end
end
Original file line number Diff line number Diff line change
Expand Up @@ -95,12 +95,12 @@ compute_w_from_continuity!(::PrescribedVelocityFields, args...; kwargs...) = not
validate_velocity_boundary_conditions(grid, ::PrescribedVelocityFields) = nothing
extract_boundary_conditions(::PrescribedVelocityFields) = NamedTuple()

FreeSurfaceDisplacementField(::PrescribedVelocityFields, ::Nothing, grid) = nothing
free_surface_displacement_field(::PrescribedVelocityFields, ::Nothing, grid) = nothing
HorizontalVelocityFields(::PrescribedVelocityFields, grid) = nothing, nothing

FreeSurface(::ExplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, grid) = nothing
FreeSurface(::ImplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, grid) = nothing
FreeSurface(::SplitExplicitFreeSurface, ::PrescribedVelocityFields, grid) = nothing
materialize_free_surface(::ExplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, grid) = nothing
materialize_free_surface(::ImplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, grid) = nothing
materialize_free_surface(::SplitExplicitFreeSurface, ::PrescribedVelocityFields, grid) = nothing

hydrostatic_prognostic_fields(::PrescribedVelocityFields, ::Nothing, tracers) = tracers
compute_hydrostatic_momentum_tendencies!(model, ::PrescribedVelocityFields, kernel_parameters; kwargs...) = nothing
Expand All @@ -121,14 +121,14 @@ on_architecture(to, velocities::PrescribedVelocityFields) =

# If the model only tracks particles... do nothing but that!!!
const OnlyParticleTrackingModel = HydrostaticFreeSurfaceModel{TS, E, A, S, G, T, V, B, R, F, P, U, C} where
{TS, E, A, S, G, T, V, B, R, F, P<:AbstractLagrangianParticles, U<:PrescribedVelocityFields, C<:NamedTuple{(), Tuple{}}}
{TS, E, A, S, G, T, V, B, R, F, P<:AbstractLagrangianParticles, U<:PrescribedVelocityFields, C<:NamedTuple{(), Tuple{}}}

function time_step!(model::OnlyParticleTrackingModel, Δt; callbacks = [], kwargs...)
function time_step!(model::OnlyParticleTrackingModel, Δt; callbacks = [], kwargs...)
model.timestepper.previous_Δt = Δt
tick!(model.clock, Δt)
step_lagrangian_particles!(model, Δt)
update_state!(model, callbacks)
end

update_state!(model::OnlyParticleTrackingModel, callbacks) =
update_state!(model::OnlyParticleTrackingModel, callbacks) =
[callback(model) for callback in callbacks if callback.callsite isa UpdateStateCallsite]
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ function Base.show(io::IO, model::HydrostaticFreeSurfaceModel)
end

if typeof(model.free_surface).name.wrapper == SplitExplicitFreeSurface
print(io, "│ └── substepping: $(model.free_surface.settings.substepping)", "\n")
print(io, "│ └── substepping: $(summary(model.free_surface.settings.substepping))", "\n")
end
end

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,12 @@ const SingleColumnGrid = AbstractGrid{<:AbstractFloat, <:Flat, <:Flat, <:Bounded
#####

PressureField(arch, ::SingleColumnGrid) = (pHY′ = nothing,)
FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, velocities, ::SingleColumnGrid) = nothing
FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, velocities, ::SingleColumnGrid) = nothing
FreeSurface(free_surface::ExplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, ::SingleColumnGrid) = nothing
FreeSurface(free_surface::ImplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, ::SingleColumnGrid) = nothing
materialize_free_surface(free_surface::ExplicitFreeSurface{Nothing}, velocities, ::SingleColumnGrid) = nothing
materialize_free_surface(free_surface::ImplicitFreeSurface{Nothing}, velocities, ::SingleColumnGrid) = nothing
materialize_free_surface(free_surface::SplitExplicitFreeSurface, velocities, ::SingleColumnGrid) = nothing
materialize_free_surface(free_surface::ExplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, ::SingleColumnGrid) = nothing
materialize_free_surface(free_surface::ImplicitFreeSurface{Nothing}, ::PrescribedVelocityFields, ::SingleColumnGrid) = nothing
materialize_free_surface(free_surface::SplitExplicitFreeSurface, ::PrescribedVelocityFields, ::SingleColumnGrid) = nothing

function HydrostaticFreeSurfaceVelocityFields(::Nothing, grid::SingleColumnGrid, clock, bcs=NamedTuple())
u = XFaceField(grid, boundary_conditions=bcs.u)
Expand Down Expand Up @@ -144,4 +146,3 @@ end
@inbounds coriolis = coriolis_array[i, j]
return y_f_cross_U(i, j, k, grid, coriolis, U)
end

Loading

0 comments on commit 9f3f8e6

Please sign in to comment.