From d289fed2d06e9530e23b60d8c421243cafb6a996 Mon Sep 17 00:00:00 2001 From: Simone Silvestri Date: Sat, 23 Nov 2024 03:28:46 -0800 Subject: [PATCH] RK3 for the HydrostaticFreeSurfaceModel version 2 (#3930) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 * 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 * 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 * 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 * S⁻ -> Ψ⁻ * small bugfix --------- Co-authored-by: Gregory L. Wagner Co-authored-by: Navid C. Constantinou --- docs/oceananigans.bib | 10 + .../HydrostaticFreeSurfaceModels.jl | 3 + .../SplitExplicitFreeSurfaces.jl | 2 +- .../compute_slow_tendencies.jl | 113 +++++++++-- .../initialize_split_explicit_substepping.jl | 39 +++- .../split_explicit_timesteppers.jl | 10 +- .../step_split_explicit_free_surface.jl | 8 +- .../explicit_free_surface.jl | 90 +++++---- .../hydrostatic_free_surface_ab2_step.jl | 12 +- .../hydrostatic_free_surface_field_tuples.jl | 10 + .../hydrostatic_free_surface_model.jl | 14 +- .../hydrostatic_free_surface_rk3_step.jl | 122 ++++++++++++ .../implicit_free_surface.jl | 5 +- .../prescribed_hydrostatic_velocity_fields.jl | 5 +- src/TimeSteppers/TimeSteppers.jl | 1 + src/TimeSteppers/quasi_adams_bashforth_2.jl | 3 +- .../split_hydrostatic_runge_kutta_3.jl | 175 ++++++++++++++++++ ...face_immersed_boundaries_implicit_solve.jl | 4 +- test/test_implicit_free_surface_solver.jl | 10 +- .../test_split_explicit_vertical_integrals.jl | 2 +- 20 files changed, 544 insertions(+), 94 deletions(-) create mode 100644 src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl create mode 100644 src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl diff --git a/docs/oceananigans.bib b/docs/oceananigans.bib index e822bfb919..1cdc5da818 100644 --- a/docs/oceananigans.bib +++ b/docs/oceananigans.bib @@ -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}, +} \ No newline at end of file diff --git a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl index e3e63fa736..033e304998 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/HydrostaticFreeSurfaceModels.jl @@ -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. @@ -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") diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl index f1399a346b..6ce20fd0b7 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/SplitExplicitFreeSurfaces.jl @@ -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 diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl index 1dbd1b3715..30f327ed20 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/compute_slow_tendencies.jl @@ -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ⁿ, χ) @@ -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 \ No newline at end of file diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl index 521fd02737..4c2976591f 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/initialize_split_explicit_substepping.jl @@ -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. # @@ -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 diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl index 6b2eb693f0..e624054d5e 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_timesteppers.jl @@ -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(η) diff --git a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl index 07896866e1..655e2e9068 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/step_split_explicit_free_surface.jl @@ -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 @@ -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)) diff --git a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl index e0645f42a2..e788a1f555 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/explicit_free_surface.jl @@ -50,8 +50,26 @@ 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, @@ -59,44 +77,17 @@ explicit_ab2_step_free_surface!(free_surface, model, Δ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 ##### @@ -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 \ No newline at end of file diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl index 2ba49d9928..6d53944ebd 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_ab2_step.jl @@ -18,8 +18,7 @@ function ab2_step!(model::HydrostaticFreeSurfaceModel, Δt) # Step locally velocity and tracers @apply_regionally local_ab2_step!(model, Δt, χ) - # blocking step for implicit free surface, non blocking for explicit - ab2_step_free_surface!(model.free_surface, model, Δt, χ) + step_free_surface!(model.free_surface, model, model.timestepper, Δt) return nothing end @@ -27,9 +26,11 @@ end function local_ab2_step!(model, Δt, χ) ab2_step_velocities!(model.velocities, model, Δt, χ) ab2_step_tracers!(model.tracers, model, Δt, χ) - return nothing + + return nothing end + ##### ##### Step velocities ##### @@ -44,9 +45,6 @@ function ab2_step_velocities!(velocities, model, Δt, χ) launch!(model.architecture, model.grid, :xyz, ab2_step_field!, velocity_field, Δt, χ, Gⁿ, G⁻) - # TODO: let next implicit solve depend on previous solve + explicit velocity step - # Need to distinguish between solver events and tendency calculation events. - # Note that BatchedTridiagonalSolver has a hard `wait`; this must be solved first. implicit_step!(velocity_field, model.timestepper.implicit_solver, model.closure, @@ -60,7 +58,7 @@ function ab2_step_velocities!(velocities, model, Δt, χ) end ##### -##### Step velocities +##### Step Tracers ##### const EmptyNamedTuple = NamedTuple{(),Tuple{}} diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl index 465cadd146..5cc82b5f37 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_field_tuples.jl @@ -1,5 +1,6 @@ using Oceananigans.Grids: Center, Face using Oceananigans.Fields: XFaceField, YFaceField, ZFaceField, TracerFields +using Oceananigans.TimeSteppers: QuasiAdamsBashforth2TimeStepper, SplitRungeKutta3TimeStepper function HydrostaticFreeSurfaceVelocityFields(::Nothing, grid, clock, bcs=NamedTuple()) u = XFaceField(grid, boundary_conditions=bcs.u) @@ -31,3 +32,12 @@ function HydrostaticFreeSurfaceTendencyFields(velocities, free_surface::SplitExp tracers = TracerFields(tracer_names, grid) return merge((u=u, v=v, U=U, V=V), tracers) end + +PreviousHydrostaticTendencyFields(::Val{:QuasiAdamsBashforth2}, args...) = HydrostaticFreeSurfaceTendencyFields(args...) +PreviousHydrostaticTendencyFields(::Val{:SplitRungeKutta3}, args...) = nothing + +function PreviousHydrostaticTendencyFields(::Val{:SplitRungeKutta3}, velocities, free_surface::SplitExplicitFreeSurface, args...) + U = deepcopy(free_surface.barotropic_velocities.U) + V = deepcopy(free_surface.barotropic_velocities.V) + return (; U=U, V=V) +end diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl index b8591ccc19..270ca8b42b 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_model.jl @@ -12,7 +12,7 @@ using Oceananigans.Forcings: model_forcing 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 +using Oceananigans.TimeSteppers: Clock, TimeStepper, update_state!, AbstractLagrangianParticles, SplitRungeKutta3TimeStepper using Oceananigans.TurbulenceClosures: validate_closure, with_tracers, DiffusivityFields, add_closure_specific_boundary_conditions using Oceananigans.TurbulenceClosures: time_discretization, implicit_diffusion_solver using Oceananigans.Utils: tupleit @@ -111,6 +111,7 @@ function HydrostaticFreeSurfaceModel(; grid, tracers = nothing, forcing::NamedTuple = NamedTuple(), closure = nothing, + timestepper = :QuasiAdamsBashforth2, boundary_conditions::NamedTuple = NamedTuple(), particles::ParticlesOrNothing = nothing, biogeochemistry::AbstractBGCOrNothing = nothing, @@ -190,14 +191,17 @@ function HydrostaticFreeSurfaceModel(; 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) - timestepper = TimeStepper(:QuasiAdamsBashforth2, grid, tracernames(tracers); + implicit_solver = implicit_diffusion_solver(time_discretization(closure), grid) + prognostic_fields = hydrostatic_prognostic_fields(velocities, free_surface, tracers) + + timestepper = TimeStepper(timestepper, grid, tracernames(tracers); implicit_solver = implicit_solver, Gⁿ = HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, grid, tracernames(tracers)), - G⁻ = HydrostaticFreeSurfaceTendencyFields(velocities, free_surface, grid, tracernames(tracers))) + G⁻ = PreviousHydrostaticTendencyFields(Val(timestepper), velocities, free_surface, grid, tracernames(tracers)), + Ψ⁻ = deepcopy(prognostic_fields)) # Regularize forcing for model tracer and velocity fields. - model_fields = merge(hydrostatic_prognostic_fields(velocities, free_surface, tracers), auxiliary_fields) + model_fields = merge(prognostic_fields, auxiliary_fields) forcing = model_forcing(model_fields; forcing...) model = HydrostaticFreeSurfaceModel(arch, grid, clock, advection, buoyancy, coriolis, diff --git a/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl new file mode 100644 index 0000000000..0a02235f89 --- /dev/null +++ b/src/Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_rk3_step.jl @@ -0,0 +1,122 @@ +using Oceananigans.Fields: location +using Oceananigans.TurbulenceClosures: implicit_step! +using Oceananigans.ImmersedBoundaries: retrieve_interior_active_cells_map, retrieve_surface_active_cells_map + +import Oceananigans.TimeSteppers: split_rk3_substep!, _split_rk3_substep_field! + +function split_rk3_substep!(model::HydrostaticFreeSurfaceModel, Δt, γⁿ, ζⁿ) + + grid = model.grid + timestepper = model.timestepper + free_surface = model.free_surface + + compute_free_surface_tendency!(grid, model, free_surface) + + rk3_substep_velocities!(model.velocities, model, Δt, γⁿ, ζⁿ) + rk3_substep_tracers!(model.tracers, model, Δt, γⁿ, ζⁿ) + + # Full step for Implicit and Split-Explicit, substep for Explicit + step_free_surface!(free_surface, model, timestepper, Δt) + + # Average free surface variables + # in the second stage + if model.clock.stage == 2 + rk3_average_free_surface!(free_surface, grid, timestepper, γⁿ, ζⁿ) + end + + return nothing +end + +rk3_average_free_surface!(free_surface, args...) = nothing + +function rk3_average_free_surface!(free_surface::ImplicitFreeSurface, grid, timestepper, γⁿ, ζⁿ) + arch = architecture(grid) + + ηⁿ⁻¹ = timestepper.Ψ⁻.η + ηⁿ = free_surface.η + + launch!(arch, grid, :xy, _rk3_average_free_surface!, ηⁿ, grid, ηⁿ⁻¹, γⁿ, ζⁿ) + + return nothing +end + +function rk3_average_free_surface!(free_surface::SplitExplicitFreeSurface, grid, timestepper, γⁿ, ζⁿ) + + arch = architecture(grid) + + Uⁿ⁻¹ = timestepper.Ψ⁻.U + Vⁿ⁻¹ = timestepper.Ψ⁻.V + Uⁿ = free_surface.barotropic_velocities.U + Vⁿ = free_surface.barotropic_velocities.V + + launch!(arch, grid, :xy, _rk3_average_free_surface!, Uⁿ, grid, Uⁿ⁻¹, γⁿ, ζⁿ) + launch!(arch, grid, :xy, _rk3_average_free_surface!, Vⁿ, grid, Vⁿ⁻¹, γⁿ, ζⁿ) + + return nothing +end + +@kernel function _rk3_average_free_surface!(η, grid, η⁻, γⁿ, ζⁿ) + i, j = @index(Global, NTuple) + k = grid.Nz + 1 + @inbounds η[i, j, k] = ζⁿ * η⁻[i, j, k] + γⁿ * η[i, j, k] +end + +##### +##### Time stepping in each substep +##### + +function rk3_substep_velocities!(velocities, model, Δt, γⁿ, ζⁿ) + + for name in (:u, :v) + Gⁿ = model.timestepper.Gⁿ[name] + Ψ⁻ = model.timestepper.Ψ⁻[name] + velocity_field = velocities[name] + + launch!(model.architecture, model.grid, :xyz, + _split_rk3_substep_field!, velocity_field, Δt, γⁿ, ζⁿ, Gⁿ, Ψ⁻) + + implicit_step!(velocity_field, + model.timestepper.implicit_solver, + model.closure, + model.diffusivity_fields, + nothing, + model.clock, + Δt) + end + + return nothing +end + +##### +##### Step Tracers +##### + +rk3_substep_tracers!(::EmptyNamedTuple, model, Δt, γⁿ, ζⁿ) = nothing + +function rk3_substep_tracers!(tracers, model, Δt, γⁿ, ζⁿ) + + closure = model.closure + grid = model.grid + + # Tracer update kernels + for (tracer_index, tracer_name) in enumerate(propertynames(tracers)) + + Gⁿ = model.timestepper.Gⁿ[tracer_name] + Ψ⁻ = model.timestepper.Ψ⁻[tracer_name] + tracer_field = tracers[tracer_name] + closure = model.closure + + launch!(architecture(grid), grid, :xyz, + _split_rk3_substep_field!, tracer_field, Δt, γⁿ, ζⁿ, Gⁿ, Ψ⁻) + + implicit_step!(tracer_field, + model.timestepper.implicit_solver, + closure, + model.diffusivity_fields, + Val(tracer_index), + model.clock, + Δt) + end + + return nothing +end \ No newline at end of file diff --git a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl index e54ad5c49c..7fab830011 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/implicit_free_surface.jl @@ -123,10 +123,7 @@ build_implicit_step_solver(::Val{:Default}, grid, settings, gravitational_accele """ Implicitly step forward η. """ -ab2_step_free_surface!(free_surface::ImplicitFreeSurface, model, Δt, χ) = - implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt, χ) - -function implicit_free_surface_step!(free_surface::ImplicitFreeSurface, model, Δt, χ) +function step_free_surface!(free_surface::ImplicitFreeSurface, model, timestepper, Δt) η = free_surface.η g = free_surface.gravitational_acceleration rhs = free_surface.implicit_step_solver.right_hand_side diff --git a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl index 94403dc62c..a78b837017 100644 --- a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl +++ b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl @@ -74,7 +74,7 @@ end function HydrostaticFreeSurfaceTendencyFields(::PrescribedVelocityFields, free_surface, grid, tracer_names) tracer_tendencies = TracerFields(tracer_names, grid) - momentum_tendencies = (u = nothing, v = nothing, η = nothing) + momentum_tendencies = (u = nothing, v = nothing) return merge(momentum_tendencies, tracer_tendencies) end @@ -89,7 +89,8 @@ end @inline datatuple(obj::PrescribedVelocityFields) = (; u = datatuple(obj.u), v = datatuple(obj.v), w = datatuple(obj.w)) ab2_step_velocities!(::PrescribedVelocityFields, args...) = nothing -ab2_step_free_surface!(::Nothing, model, Δt, χ) = nothing +rk3_substep_velocities!(::PrescribedVelocityFields, args...) = nothing +step_free_surface!(::Nothing, model, timestepper, Δt) = nothing compute_w_from_continuity!(::PrescribedVelocityFields, args...; kwargs...) = nothing validate_velocity_boundary_conditions(grid, ::PrescribedVelocityFields) = nothing diff --git a/src/TimeSteppers/TimeSteppers.jl b/src/TimeSteppers/TimeSteppers.jl index 2955ca6999..e0e50c8970 100644 --- a/src/TimeSteppers/TimeSteppers.jl +++ b/src/TimeSteppers/TimeSteppers.jl @@ -59,5 +59,6 @@ include("clock.jl") include("store_tendencies.jl") include("quasi_adams_bashforth_2.jl") include("runge_kutta_3.jl") +include("split_hydrostatic_runge_kutta_3.jl") end # module diff --git a/src/TimeSteppers/quasi_adams_bashforth_2.jl b/src/TimeSteppers/quasi_adams_bashforth_2.jl index ed9f55750b..6d8d7ba24a 100644 --- a/src/TimeSteppers/quasi_adams_bashforth_2.jl +++ b/src/TimeSteppers/quasi_adams_bashforth_2.jl @@ -41,7 +41,8 @@ function QuasiAdamsBashforth2TimeStepper(grid, tracers, χ = 0.1; implicit_solver::IT = nothing, Gⁿ = TendencyFields(grid, tracers), - G⁻ = TendencyFields(grid, tracers)) where IT + G⁻ = TendencyFields(grid, tracers), + kw...) where IT FT = eltype(grid) GT = typeof(Gⁿ) diff --git a/src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl b/src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl new file mode 100644 index 0000000000..2dbf8d2d1f --- /dev/null +++ b/src/TimeSteppers/split_hydrostatic_runge_kutta_3.jl @@ -0,0 +1,175 @@ +using Oceananigans.Architectures: architecture +using Oceananigans: fields + +""" + SplitRungeKutta3TimeStepper{FT, TG} <: AbstractTimeStepper + +Holds parameters and tendency fields for a low storage, third-order Runge-Kutta-Wray +time-stepping scheme described by [Lan2022](@citet). +""" +struct SplitRungeKutta3TimeStepper{FT, TG, TE, PF, TI} <: AbstractTimeStepper + γ² :: FT + γ³ :: FT + ζ² :: FT + ζ³ :: FT + Gⁿ :: TG + G⁻ :: TE # only used as storage when needed + Ψ⁻ :: PF # prognostic state at the previous timestep + implicit_solver :: TI +end + +""" + SplitRungeKutta3TimeStepper(grid, tracers; + implicit_solver = nothing, + Gⁿ = TendencyFields(grid, tracers), + G⁻ = TendencyFields(grid, tracers), + Ψ⁻ = TendencyFields(grid, tracers)) + +Return a 3rd-order `SplitRungeKutta3TimeStepper` on `grid` and with `tracers`. +The tendency fields `Gⁿ` and `G⁻`, as well as the previous prognostic state Ψ⁻ can be specified via optional `kwargs`. + +The scheme described by [Lan2022](@citet). In a nutshel, the 3rd-order +Runge Kutta timestepper steps forward the state `Uⁿ` by `Δt` via 3 substeps. A baroptropic velocity correction +step is applied after at each substep. + +The state `U` after each substep `m` is + +```julia +Uᵐ⁺¹ = ζᵐ * Uⁿ + γᵐ * (Uᵐ + Δt * Gᵐ) +``` + +where `Uᵐ` is the state at the ``m``-th substep, `Uⁿ` is the state at the ``n``-th timestep, `Gᵐ` is the tendency +at the ``m``-th substep, and constants ``γ¹ = 1`, ``γ² = 1/4``, ``γ³ = 1/3``, +``ζ¹ = 0``, ``ζ² = 3/4``, ``ζ³ = 1/3``. + +The state at the first substep is taken to be the one that corresponds to the ``n``-th timestep, +`U¹ = Uⁿ`, and the state after the third substep is then the state at the `Uⁿ⁺¹ = U³`. +""" +function SplitRungeKutta3TimeStepper(grid, tracers; + implicit_solver::TI = nothing, + Gⁿ::TG = TendencyFields(grid, tracers), + G⁻::TE = TendencyFields(grid, tracers), + Ψ⁻::PF = TendencyFields(grid, tracers)) where {TI, TG, TE, PF} + + + @warn("Split barotropic-baroclinic time stepping with SplitRungeKutta3TimeStepper is not tested and experimental.\n" * + "Use at own risk, and report any issues encountered.") + + !isnothing(implicit_solver) && + @warn("Implicit-explicit time-stepping with SplitRungeKutta3TimeStepper is not tested. " * + "\n implicit_solver: $(typeof(implicit_solver))") + + γ² = 1 // 4 + γ³ = 2 // 3 + + ζ² = 3 // 4 + ζ³ = 1 // 3 + + FT = eltype(grid) + + return SplitRungeKutta3TimeStepper{FT, TG, TE, PF, TI}(γ², γ³, ζ², ζ³, Gⁿ, G⁻, Ψ⁻, implicit_solver) +end + + +function time_step!(model::AbstractModel{<:SplitRungeKutta3TimeStepper}, Δt; callbacks=[]) + Δt == 0 && @warn "Δt == 0 may cause model blowup!" + + # Be paranoid and update state at iteration 0, in case run! is not used: + model.clock.iteration == 0 && update_state!(model, callbacks; compute_tendencies = true) + + γ² = model.timestepper.γ² + γ³ = model.timestepper.γ³ + + ζ² = model.timestepper.ζ² + ζ³ = model.timestepper.ζ³ + + store_fields!(model) + + #### + #### First stage + #### + + model.clock.stage = 1 + + split_rk3_substep!(model, Δt, nothing, nothing) + calculate_pressure_correction!(model, Δt) + pressure_correct_velocities!(model, Δt) + update_state!(model, callbacks; compute_tendencies = true) + + #### + #### Second stage + #### + + model.clock.stage = 2 + + split_rk3_substep!(model, Δt, γ², ζ²) + calculate_pressure_correction!(model, Δt) + pressure_correct_velocities!(model, Δt) + update_state!(model, callbacks; compute_tendencies = true) + + #### + #### Third stage + #### + + model.clock.stage = 3 + + split_rk3_substep!(model, Δt, γ³, ζ³) + calculate_pressure_correction!(model, Δt) + pressure_correct_velocities!(model, Δt) + update_state!(model, callbacks; compute_tendencies = true) + + step_lagrangian_particles!(model, Δt) + + tick!(model.clock, Δt) + model.clock.last_Δt = Δt + + return nothing +end + +@kernel function _split_rk3_substep_field!(cⁿ, Δt, γⁿ::FT, ζⁿ, Gⁿ, cⁿ⁻¹) where FT + i, j, k = @index(Global, NTuple) + cⁿ[i, j, k] = ζⁿ * cⁿ⁻¹[i, j, k] + γⁿ * (cⁿ[i, j, k] + convert(FT, Δt) * Gⁿ[i, j, k]) +end + +@kernel function _split_rk3_substep_field!(cⁿ, Δt, ::Nothing, ::Nothing, Gⁿ, cⁿ⁻¹) + i, j, k = @index(Global, NTuple) + cⁿ[i, j, k] = cⁿ[i, j, k] + Δt * Gⁿ[i, j, k] +end + +function split_rk3_substep!(model, Δt, γⁿ, ζⁿ) + + grid = model.grid + arch = architecture(grid) + model_fields = prognostic_fields(model) + + for (i, field) in enumerate(model_fields) + kernel_args = (field, Δt, γⁿ, ζⁿ, model.timestepper.Gⁿ[i], model.timestepper.Ψ⁻[i]) + launch!(arch, grid, :xyz, rk3_substep_field!, kernel_args...; exclude_periphery=true) + + # TODO: function tracer_index(model, field_index) = field_index - 3, etc... + tracer_index = Val(i - 3) # assumption + + implicit_step!(field, + model.timestepper.implicit_solver, + model.closure, + model.diffusivity_fields, + tracer_index, + model.clock, + stage_Δt(Δt, γⁿ, ζⁿ)) + end +end + +function store_fields!(model) + + previous_fields = model.timestepper.Ψ⁻ + model_fields = prognostic_fields(model) + + for name in keys(previous_fields) + if !isnothing(previous_fields[name]) + parent(previous_fields[name]) .= parent(model_fields[name]) # Storing also the halos + end + end + + return nothing +end + diff --git a/test/test_hydrostatic_free_surface_immersed_boundaries_implicit_solve.jl b/test/test_hydrostatic_free_surface_immersed_boundaries_implicit_solve.jl index 9cf8116bbb..f7e5c4142f 100644 --- a/test/test_hydrostatic_free_surface_immersed_boundaries_implicit_solve.jl +++ b/test/test_hydrostatic_free_surface_immersed_boundaries_implicit_solve.jl @@ -5,7 +5,7 @@ using Oceananigans.Architectures: on_architecture using Oceananigans.TurbulenceClosures using Oceananigans.Models.HydrostaticFreeSurfaceModels: compute_vertically_integrated_volume_flux!, compute_implicit_free_surface_right_hand_side!, - implicit_free_surface_step!, + step_free_surface!, pressure_correct_velocities! @testset "Immersed boundaries test divergent flow solve with hydrostatic free surface models" begin @@ -55,7 +55,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: compute_vertically_integ v[imm1, jmm1, 1:Nz] .= 1 v[imm1, jmp1, 1:Nz] .= -1 - implicit_free_surface_step!(model.free_surface, model, 1.0, 1.5) + step_free_surface!(model.free_surface, model, model.timestepper, 1.0) sol = (sol..., model.free_surface.η) f = (f..., model.free_surface) diff --git a/test/test_implicit_free_surface_solver.jl b/test/test_implicit_free_surface_solver.jl index d56c47562a..c4ef0f2a19 100644 --- a/test/test_implicit_free_surface_solver.jl +++ b/test/test_implicit_free_surface_solver.jl @@ -10,7 +10,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels: PCGImplicitFreeSurfaceSolver, MatrixImplicitFreeSurfaceSolver, compute_vertically_integrated_lateral_areas!, - implicit_free_surface_step!, + step_free_surface!, implicit_free_surface_linear_operation! using Oceananigans.Grids: with_halo @@ -52,7 +52,7 @@ function run_implicit_free_surface_solver_tests(arch, grid, free_surface) free_surface) set_simple_divergent_velocity!(model) - implicit_free_surface_step!(model.free_surface, model, Δt, 1.5) + step_free_surface!(model.free_surface, model, model.timestepper, Δt) acronym = free_surface.solver_method == :HeptadiagonalIterativeSolver ? "Matrix" : "PCG" @@ -167,9 +167,9 @@ end for m in (mat_model, pcg_model, fft_model) set_simple_divergent_velocity!(m) - implicit_free_surface_step!(m.free_surface, m, Δt₁, 1.5) - implicit_free_surface_step!(m.free_surface, m, Δt₁, 1.5) - implicit_free_surface_step!(m.free_surface, m, Δt₂, 1.5) + step_free_surface!(m.free_surface, m, m.timestepper, Δt₁) + step_free_surface!(m.free_surface, m, m.timestepper, Δt₁) + step_free_surface!(m.free_surface, m, m.timestepper, Δt₂) end mat_η = mat_model.free_surface.η diff --git a/test/test_split_explicit_vertical_integrals.jl b/test/test_split_explicit_vertical_integrals.jl index bedb2162ab..406f31c057 100644 --- a/test/test_split_explicit_vertical_integrals.jl +++ b/test/test_split_explicit_vertical_integrals.jl @@ -38,7 +38,7 @@ using Oceananigans.Models.HydrostaticFreeSurfaceModels.SplitExplicitFreeSurfaces η̅ .= U̅ .= V̅ .= 1.0 # now set equal to zero - initialize_free_surface_state!(state, sefs.η, barotropic_velocities, sefs.timestepper) + initialize_free_surface_state!(sefs, sefs.timestepper, sefs.timestepper, Val(1)) # don't forget the halo points fill_halo_regions!(η̅)