Skip to content

Commit

Permalink
RK3 for the HydrostaticFreeSurfaceModel version 2 (#3930)
Browse files Browse the repository at this point in the history
* 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

* give it a try

* runs

* small test

* test it in a more serious simulation

* works!

* remove the show

* remove useless terms

* add some validation

* some small changes

* 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

* new commit

* better

* update

* 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?

* ok let's go

* add to docs

* start changing everything

* continue

* bugfix

* continue

* getting there...

* proceding

* 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

* cleaner

* should work nicely

* bugfix in integrated tendencies

* rk3 working

* think about it a bit more

* try putting it back in the ab2step

* revert quickly to see if this works

* compute_free_surface_tendecny! is the only problem

* no need for const c and const f

* merge

* adapted to explicit

* some improvements

* fix errors

* fix rk3

* bugfix

* eliding NaNs + separate compute_slow_tendencies

* whoops wrong name

* correct compute tendencies

* should work?

* ok this works!

* works

* works nicely

* improvements

* warnings

* no need for this extra defnition

* no more `implicit_free_surface_step!`

* fix bugs

* fix test

* CFL changes

* ready PR for merging

* remove duplicate code

* fix AB2

* Update split_hydrostatic_runge_kutta_3.jl

* Update src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl

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

* S⁻ -> Ψ⁻

* small bugfix

---------

Co-authored-by: Gregory L. Wagner <[email protected]>
Co-authored-by: Navid C. Constantinou <[email protected]>
  • Loading branch information
3 people authored Nov 23, 2024
1 parent 9cc7112 commit d289fed
Show file tree
Hide file tree
Showing 20 changed files with 544 additions and 94 deletions.
10 changes: 10 additions & 0 deletions docs/oceananigans.bib
Original file line number Diff line number Diff line change
Expand Up @@ -979,3 +979,13 @@ @article{BouZeid05
year = {2005},
doi = {10.1063/1.1839152},
}

@article{Lan2022,
author = {Lan, Rihui and Ju, Lili and Wanh, Zhu and Gunzburger, Max and Jones, Philip},
title = {High-order multirate explicit time-stepping schemes for the baroclinic-barotropic split dynamics in primitive equations},
journal = {Journal of Computational Physics},
volume = {457},
pages = {111050},
year = {2022},
doi = {https://doi.org/10.1016/j.jcp.2022.111050},
}
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@ import Oceananigans.Advection: cell_advection_timescale
import Oceananigans.TimeSteppers: step_lagrangian_particles!
import Oceananigans.Architectures: on_architecture

using Oceananigans.TimeSteppers: SplitRungeKutta3TimeStepper, QuasiAdamsBashforth2TimeStepper

abstract type AbstractFreeSurface{E, G} end

# This is only used by the cubed sphere for now.
Expand Down Expand Up @@ -132,6 +134,7 @@ include("compute_hydrostatic_free_surface_tendencies.jl")
include("compute_hydrostatic_free_surface_buffers.jl")
include("update_hydrostatic_free_surface_model_state.jl")
include("hydrostatic_free_surface_ab2_step.jl")
include("hydrostatic_free_surface_rk3_step.jl")
include("store_hydrostatic_free_surface_tendencies.jl")
include("prescribed_hydrostatic_velocity_fields.jl")
include("single_column_model_mode.jl")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ using KernelAbstractions.Extras.LoopInfo: @unroll

import Oceananigans.Models.HydrostaticFreeSurfaceModels: initialize_free_surface!,
materialize_free_surface,
ab2_step_free_surface!,
step_free_surface!,
compute_free_surface_tendency!,
explicit_barotropic_pressure_x_gradient,
explicit_barotropic_pressure_y_gradient
Expand Down
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
#####
##### Compute slow tendencies with an AB2 timestepper
#####

# Calculate RHS for the barotropic time step.
@kernel function _compute_integrated_ab2_tendencies!(Gᵁ, Gⱽ, grid, ::Nothing, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, χ)
Expand Down Expand Up @@ -38,32 +41,114 @@ end
return ifelse(immersed, zero(grid), Gⁿ⁺¹)
end

@inline function compute_split_explicit_forcing!(GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ,
timestepper::QuasiAdamsBashforth2TimeStepper, stage)
active_cells_map = retrieve_surface_active_cells_map(grid)

Gu⁻ = timestepper.G⁻.u
Gv⁻ = timestepper.G⁻.v

launch!(architecture(grid), grid, :xy, _compute_integrated_ab2_tendencies!, GUⁿ, GVⁿ, grid,
active_cells_map, Gu⁻, Gv⁻, Guⁿ, Gvⁿ, timestepper.χ; active_cells_map)

return nothing
end

#####
##### Compute slow tendencies with a RK3 timestepper
#####

@inline function G_vertical_integral(i, j, grid, Gⁿ, ℓx, ℓy, ℓz)
immersed = peripheral_node(i, j, 1, grid, ℓx, ℓy, ℓz)

Gⁿ⁺¹ = Δz(i, j, 1, grid, ℓx, ℓy, ℓz) * ifelse(immersed, zero(grid), Gⁿ[i, j, 1])

@inbounds for k in 2:grid.Nz
immersed = peripheral_node(i, j, k, grid, ℓx, ℓy, ℓz)
Gⁿ⁺¹ += Δz(i, j, k, grid, ℓx, ℓy, ℓz) * ifelse(immersed, zero(grid), Gⁿ[i, j, k])
end

return Gⁿ⁺¹
end

@kernel function _compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, grid, active_cells_map, Guⁿ, Gvⁿ, stage)
idx = @index(Global, Linear)
i, j = active_linear_index_to_tuple(idx, active_cells_map)
compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, stage)
end

@kernel function _compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, grid, ::Nothing, Guⁿ, Gvⁿ, stage)
i, j = @index(Global, NTuple)
compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, stage)
end

@inline function compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, ::Val{1})
@inbounds GUⁿ[i, j, 1] = G_vertical_integral(i, j, grid, Guⁿ, Face(), Center(), Center())
@inbounds GVⁿ[i, j, 1] = G_vertical_integral(i, j, grid, Gvⁿ, Center(), Face(), Center())

@inbounds GU⁻[i, j, 1] = GUⁿ[i, j, 1]
@inbounds GV⁻[i, j, 1] = GVⁿ[i, j, 1]
end

@inline function compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, ::Val{2})
FT = eltype(GUⁿ)

@inbounds GUⁿ[i, j, 1] = G_vertical_integral(i, j, grid, Guⁿ, Face(), Center(), Center())
@inbounds GVⁿ[i, j, 1] = G_vertical_integral(i, j, grid, Gvⁿ, Center(), Face(), Center())

@inbounds GU⁻[i, j, 1] = (GUⁿ[i, j, 1] + GU⁻[i, j, 1]) / 6
@inbounds GV⁻[i, j, 1] = (GVⁿ[i, j, 1] + GV⁻[i, j, 1]) / 6
end

@inline function compute_integrated_rk3_tendencies!(GUⁿ, GVⁿ, GU⁻, GV⁻, i, j, grid, Guⁿ, Gvⁿ, ::Val{3})
FT = eltype(GUⁿ)

GUi = G_vertical_integral(i, j, grid, Guⁿ, Face(), Center(), Center())
GVi = G_vertical_integral(i, j, grid, Gvⁿ, Center(), Face(), Center())

@inbounds GUⁿ[i, j, 1] = convert(FT, 2/3) * GUi + GU⁻[i, j, 1]
@inbounds GVⁿ[i, j, 1] = convert(FT, 2/3) * GVi + GV⁻[i, j, 1]
end

@inline function compute_split_explicit_forcing!(GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ,
timestepper::SplitRungeKutta3TimeStepper, stage)

GU⁻ = timestepper.G⁻.U
GV⁻ = timestepper.G⁻.V

active_cells_map = retrieve_surface_active_cells_map(grid)
launch!(architecture(grid), grid, :xy, _compute_integrated_rk3_tendencies!,
GUⁿ, GVⁿ, GU⁻, GV⁻, grid, active_cells_map, Guⁿ, Gvⁿ, stage; active_cells_map)

return nothing
end

#####
##### Free surface setup
#####

# 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)
function compute_free_surface_tendency!(grid, model, free_surface::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.χ)
barotropic_timestepper = free_surface.timestepper
baroclinic_timestepper = model.timestepper

stage = model.clock.stage

@apply_regionally begin
compute_split_explicit_forcing!(GUⁿ, GVⁿ, grid, Guⁿ, Gvⁿ, baroclinic_timestepper, Val(stage))
initialize_free_surface_state!(free_surface, baroclinic_timestepper, barotropic_timestepper, Val(stage))
end

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
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
using Oceananigans.ImmersedBoundaries: retrieve_surface_active_cells_map, peripheral_node
using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, SplitRungeKutta3TimeStepper
using Oceananigans.Operators: Δz

# This file contains two different initializations methods performed at different stages of the simulation.
#
Expand All @@ -19,13 +21,40 @@ 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)
function initialize_free_surface_state!(free_surface, baroclinic_timestepper, timestepper, stage)

initialize_free_surface_timestepper!(timestepper, η, velocities)
η = free_surface.η
U, V = free_surface.barotropic_velocities

fill!(filtered_state.η, 0)
fill!(filtered_state.U, 0)
fill!(filtered_state.V, 0)
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)

return nothing
end

# At the last stage we reset the velocities and perform the complete substepping from n to n+1
function initialize_free_surface_state!(free_surface, baroclinic_ts::SplitRungeKutta3TimeStepper, barotropic_ts, ::Val{3})

η = free_surface.η
U, V = free_surface.barotropic_velocities

Uⁿ⁻¹ = baroclinic_ts.Ψ⁻.U
Vⁿ⁻¹ = baroclinic_ts.Ψ⁻.V
ηⁿ⁻¹ = baroclinic_ts.Ψ⁻.η

# Restart from the state at baroclinic step n
parent(U) .= parent(Uⁿ⁻¹)
parent(V) .= parent(Vⁿ⁻¹)
parent(η) .= parent(ηⁿ⁻¹)

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)

return nothing
end
Original file line number Diff line number Diff line change
Expand Up @@ -120,12 +120,12 @@ end

initialize_free_surface_timestepper!(::ForwardBackwardScheme, args...) = nothing

function initialize_free_surface_timestepper!(timestepper::AdamsBashforth3Scheme, η, velocities)
parent(timestepper.Uᵐ⁻¹) .= parent(velocities.U)
parent(timestepper.Vᵐ⁻¹) .= parent(velocities.V)
function initialize_free_surface_timestepper!(timestepper::AdamsBashforth3Scheme, η, U, V)
parent(timestepper.Uᵐ⁻¹) .= parent(U)
parent(timestepper.Vᵐ⁻¹) .= parent(V)

parent(timestepper.Uᵐ⁻²) .= parent(velocities.U)
parent(timestepper.Vᵐ⁻²) .= parent(velocities.V)
parent(timestepper.Uᵐ⁻²) .= parent(U)
parent(timestepper.Vᵐ⁻²) .= parent(V)

parent(timestepper.ηᵐ) .= parent(η)
parent(timestepper.ηᵐ⁻¹) .= parent(η)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -115,17 +115,13 @@ end
##### SplitExplicitFreeSurface barotropic subcylicing
#####

ab2_step_free_surface!(free_surface::SplitExplicitFreeSurface, model, Δt, χ) =
split_explicit_free_surface_step!(free_surface, model, Δt)

function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurface, model, Δt)
function step_free_surface!(free_surface::SplitExplicitFreeSurface, model, baroclinic_timestepper, Δt)

# Note: free_surface.η.grid != model.grid for DistributedSplitExplicitFreeSurface
# since halo_size(free_surface.η.grid) != halo_size(model.grid)
free_surface_grid = free_surface.η.grid
filtered_state = free_surface.filtered_state
substepping = free_surface.substepping
timestepper = free_surface.timestepper

barotropic_velocities = free_surface.barotropic_velocities

Expand Down Expand Up @@ -155,8 +151,6 @@ function split_explicit_free_surface_step!(free_surface::SplitExplicitFreeSurfac

# reset free surface averages
@apply_regionally begin
initialize_free_surface_state!(filtered_state, free_surface.η, barotropic_velocities, timestepper)

# Solve for the free surface at tⁿ⁺¹
iterate_split_explicit!(free_surface, free_surface_grid, GUⁿ, GVⁿ, Δτᴮ, weights, Val(Nsubsteps))

Expand Down
90 changes: 55 additions & 35 deletions src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl
Original file line number Diff line number Diff line change
Expand Up @@ -50,53 +50,44 @@ end
##### Time stepping
#####

ab2_step_free_surface!(free_surface::ExplicitFreeSurface, model, Δt, χ) =
@apply_regionally explicit_ab2_step_free_surface!(free_surface, model, Δt, χ)
step_free_surface!(free_surface::ExplicitFreeSurface, model, timestepper::QuasiAdamsBashforth2TimeStepper, Δt) =
@apply_regionally explicit_ab2_step_free_surface!(free_surface, model, Δt, timestepper.χ)

step_free_surface!(free_surface::ExplicitFreeSurface, model, timestepper::SplitRungeKutta3TimeStepper, Δt) =
@apply_regionally explicit_rk3_step_free_surface!(free_surface, model, Δt, timestepper)

@inline rk3_coeffs(ts, ::Val{1}) = (1, 0)
@inline rk3_coeffs(ts, ::Val{2}) = (ts.γ², ts.ζ²)
@inline rk3_coeffs(ts, ::Val{3}) = (ts.γ³, ts.ζ³)

function explicit_rk3_step_free_surface!(free_surface, model, Δt, timestepper)

γⁿ, ζⁿ = rk3_coeffs(timestepper, Val(model.clock.stage))

launch!(model.architecture, model.grid, :xy,
_explicit_rk3_step_free_surface!, free_surface.η, Δt, γⁿ, ζⁿ,
model.timestepper.Gⁿ.η, model.timestepper.Ψ⁻.η, size(model.grid, 3))

return nothing
end

explicit_ab2_step_free_surface!(free_surface, model, Δt, χ) =
launch!(model.architecture, model.grid, :xy,
_explicit_ab2_step_free_surface!, free_surface.η, Δt, χ,
model.timestepper.Gⁿ.η, model.timestepper.G⁻.η, size(model.grid, 3))

#####
##### Kernel
##### Kernels
#####

@kernel function _explicit_ab2_step_free_surface!(η, Δt, χ::FT, Gηⁿ, Gη⁻, Nz) where FT
@kernel function _explicit_rk3_step_free_surface!(η, Δt, γⁿ, ζⁿ, Gⁿ, η⁻, Nz)
i, j = @index(Global, NTuple)

@inbounds begin
η[i, j, Nz+1] += Δt * ((FT(1.5) + χ) * Gηⁿ[i, j, Nz+1] - (FT(0.5) + χ) * Gη⁻[i, j, Nz+1])
end
@inbounds η[i, j, Nz+1] += ζⁿ * η⁻[i, j, k] + γⁿ * (η[i, j, k] + convert(FT, Δt) * Gⁿ[i, j, k])
end

compute_free_surface_tendency!(grid, model, ::ExplicitFreeSurface) =
@apply_regionally compute_explicit_free_surface_tendency!(grid, model)

# Compute free surface tendency
function compute_explicit_free_surface_tendency!(grid, model)

arch = architecture(grid)

args = tuple(model.velocities,
model.free_surface,
model.tracers,
model.auxiliary_fields,
model.forcing,
model.clock)

launch!(arch, grid, :xy,
compute_hydrostatic_free_surface_Gη!, model.timestepper.Gⁿ.η,
grid, args)

args = (model.clock,
fields(model),
model.closure,
model.buoyancy)

apply_flux_bcs!(model.timestepper.Gⁿ.η, displacement(model.free_surface), arch, args)

return nothing
@kernel function _explicit_ab2_step_free_surface!(η, Δt, χ::FT, Gηⁿ, Gη⁻, Nz) where FT
i, j = @index(Global, NTuple)
@inbounds η[i, j, Nz+1] += Δt * ((FT(1.5) + χ) * Gηⁿ[i, j, Nz+1] - (FT(0.5) + χ) * Gη⁻[i, j, Nz+1])
end

#####
Expand Down Expand Up @@ -140,3 +131,32 @@ The tendency is called ``G_η`` and defined via
return @inbounds ( velocities.w[i, j, k_top]
+ forcings.η(i, j, k_top, grid, clock, model_fields))
end

compute_free_surface_tendency!(grid, model, ::ExplicitFreeSurface) =
@apply_regionally compute_explicit_free_surface_tendency!(grid, model)

# Compute free surface tendency
function compute_explicit_free_surface_tendency!(grid, model)

arch = architecture(grid)

args = tuple(model.velocities,
model.free_surface,
model.tracers,
model.auxiliary_fields,
model.forcing,
model.clock)

launch!(arch, grid, :xy,
compute_hydrostatic_free_surface_Gη!, model.timestepper.Gⁿ.η,
grid, args)

args = (model.clock,
fields(model),
model.closure,
model.buoyancy)

apply_flux_bcs!(model.timestepper.Gⁿ.η, displacement(model.free_surface), arch, args)

return nothing
end
Loading

0 comments on commit d289fed

Please sign in to comment.