s-stage SSPRK3

Author: milankl

src/Constants.jl

@@ -1,3 +1,31 @@
+"""Coefficients for strong stability-preserving Runge-Kutta 3rd order.
+From: KETCHESON, LOĆZI, AND PARSANI, 2014. INTERNAL ERROR PROPAGATION IN EXPLICIT RUNGE–KUTTA METHODS, 
+SIAM J NUMER ANAL 52/5. DOI:10.1137/130936245"""
+struct SSPRK3coeff{T<:AbstractFloat}
+    n::Int
+    s::Int
+    kn::Int
+    mn::Int
+    Δt_Δn::T
+    kna::T
+    knb::T
+    Δt_Δnc::T
+end
+
+"""Generator function for a SSPRK3coeff struct."""
+function SSPRK3coeff{T}(P::Parameter,Δt_Δ::T) where T
+    n = P.RKn
+    s = n^2
+    kn = n*(n+1) ÷ 2 + 1
+    mn = (n-1)*(n-2) ÷ 2 + 1
+    Δt_Δn = T(Δt_Δ/(n^2-n))
+    kna = T((n-1)/(2n-1))
+    knb = T(n/(2n-1))
+    Δt_Δnc = T(Δt_Δ/(n*(2n-1)))
+
+    return SSPRK3coeff{T}(n,s,kn,mn,Δt_Δn,kna,knb,Δt_Δnc)
+end
+
 struct Constants{T<:AbstractFloat,Tprog<:AbstractFloat}
 
     # RUNGE-KUTTA COEFFICIENTS 2nd/3rd/4th order including timestep Δt
@@ -6,6 +34,7 @@ struct Constants{T<:AbstractFloat,Tprog<:AbstractFloat}
     Δt_Δs::Tprog            # Δt/(s-1) wher s the number of stages
     Δt_Δ::Tprog             # Δt/Δ - timestep divided by grid spacing
     Δt_Δ_half::Tprog        # 1/2 * Δt/Δ
+    SSPRK3c::SSPRK3coeff    # struct containing all coefficients for SSPRK3
 
     # BOUNDARY CONDITIONS
     one_minus_α::Tprog      # tangential boundary condition for the ghost-point copy
@@ -49,8 +78,12 @@ function Constants{T,Tprog}(P::Parameter,G::Grid) where {T<:AbstractFloat,Tprog<
     Δt_Δ = Tprog(G.dtint/G.Δ)
     Δt_Δ_half = Tprog(G.dtint/G.Δ/2)
 
-    one_minus_α = Tprog(1-P.α)    # for the ghost point copy/tangential boundary conditions
-    g = T(P.g)                # gravity - for Bernoulli potential
+    # coefficients for SSPRK3
+    SSPRK3c = SSPRK3coeff{Tprog}(P,Δt_Δ)
+
+    # BOUNDARY CONDITIONS AND PHYSICS
+    one_minus_α = Tprog(1-P.α)      # for the ghost point copy/tangential boundary conditions
+    g = T(P.g)                      # gravity - for Bernoulli potential
 
     # BOTTOM FRICTION COEFFICENTS
     # incl grid spacing Δ for non-dimensional gradients
@@ -76,7 +109,7 @@ function Constants{T,Tprog}(P::Parameter,G::Grid) where {T<:AbstractFloat,Tprog<
     ωFy = 2π*P.ωFy/24/365.25/3600
 
     return Constants{T,Tprog}(  RKaΔt,RKbΔt,Δt_Δs,Δt_Δ,Δt_Δ_half,
-                                one_minus_α,
+                                SSPRK3c,one_minus_α,
                                 g,cD,rD,γ,cSmag,νB,rSST,
                                 jSST,SSTmin,ωFη,ωFx,ωFy)
 end

src/DefaultParameters.jl

@@ -47,9 +47,11 @@
     wk::Real=10e3                       # width [m] in y of Gaussian used for surface forcing
 
     # TIME STEPPING OPTIONS
-    time_scheme::String="RK"            # Runge-Kutta ("RK") or strong-stability preserving RK ("SSPRK2","SSPRK3")
+    time_scheme::String="RK"            # Runge-Kutta ("RK") or strong-stability preserving RK
+                                        # "SSPRK2","SSPRK3","4SSPRK3"
     RKo::Int=4                          # Order of the RK time stepping scheme (2, 3 or 4)
     RKs::Int=3                          # Number of stages for SSPRK2
+    RKn::Int=2                          # n^2 = s = Number of stages  for SSPRK3
     cfl::Real=1.0                       # CFL number (1.0 recommended for RK4, 0.6 for RK3)
     Ndays::Real=10.0                    # number of days to integrate for
     nstep_diff::Int=1                   # diffusive part every nstep_diff time steps.
@@ -120,7 +122,7 @@
     @assert topo_width > 0.0    "topo_width has to be >0, $topo_width given."
     @assert t_relax > 0.0       "t_relax has to be >0, $t_relax given."
     @assert η_refw > 0.0        "η_refw has to be >0, $η_refw given."
-    @assert time_scheme in ["RK","SSPRK2","SSPRK3"] "Time scheme $time_scheme unsupported."
+    @assert time_scheme in ["RK","SSPRK2","SSPRK3","4SSPRK3"] "Time scheme $time_scheme unsupported."
     @assert RKo in [2,3,4]        "RKo has to be 2,3 or 4; $RKo given."
     @assert RKs > 1               "RKs has to be >= 2; $RKs given."
     @assert Ndays > 0.0         "Ndays has to be >0, $Ndays given."

src/GhostPoints.jl

@@ -177,7 +177,7 @@ function ghost_points!( u::AbstractMatrix,
 end
 
 """Decide on boundary condition P.bc which ghost point function to execute."""
-function ghost_points!( u::AbstractMatrix,
+function ghost_points_uv!( u::AbstractMatrix,
                         v::AbstractMatrix,
                         S::ModelSetup)
 
@@ -194,6 +194,20 @@ function ghost_points!( u::AbstractMatrix,
     end
 end
 
+"""Decide on boundary condition P.bc which ghost point function to execute."""
+function ghost_points_η!(   η::AbstractMatrix,
+                            S::ModelSetup)
+
+    @unpack bc = S.parameters
+
+    if bc == "periodic"
+        @unpack Tcomm = S.parameters
+        ghost_points_η_periodic!(Tcomm,η)
+    else
+        ghost_points_η_nonperiodic!(η)
+    end
+end
+
 """Decide on boundary condition P.bc which ghost point function to execute."""
 function ghost_points_sst!(sst::AbstractMatrix,S::ModelSetup)
 

src/TimeIntegration.jl

@@ -81,23 +81,22 @@ function time_integration(  Prog::PrognosticVars{Tprog},
 
             for rki = 1:RKs
                 if rki > 1
-                    # TODO technically the ghost point copy for u1,v1 is redundant as done further down
-                    ghost_points!(u1,v1,η1,S)
+                    ghost_points_η!(η1,S)
                 end
 
                 # type conversion for mixed precision
                 u1rhs = convert(Diag.PrognosticVarsRHS.u,u1)
                 v1rhs = convert(Diag.PrognosticVarsRHS.v,v1)
                 η1rhs = convert(Diag.PrognosticVarsRHS.η,η1)
 
-                rhs!(u1rhs,v1rhs,η1rhs,Diag,S,t)
+                rhs!(u1rhs,v1rhs,η1rhs,Diag,S,t)        # momentum only
 
                 # the update step
                 axb!(u1,Δt_Δs,du)       # u1 = u1 + Δt/(s-1)*RHS(u1)
                 axb!(v1,Δt_Δs,dv)
 
-                # semi-implicit for continuity equation, use u1,v1 to calcualte dη
-                ghost_points!(u1,v1,S)
+                # semi-implicit for continuity equation, use new u1,v1 to calcualte dη
+                ghost_points_uv!(u1,v1,S)
                 u1rhs = convert(Diag.PrognosticVarsRHS.u,u1)
                 v1rhs = convert(Diag.PrognosticVarsRHS.v,v1)
                 continuity!(u1rhs,v1rhs,η1rhs,Diag,S,t)
@@ -109,8 +108,54 @@ function time_integration(  Prog::PrognosticVars{Tprog},
             cxayb!(u0,a,u,b,u1)
             cxayb!(v0,a,v,b,v1)
             cxayb!(η0,a,η,b,η1)
+        
+        elseif time_scheme == "SSPRK3"  # s-stage 3rd order SSPRK
+
+            @unpack s,kn,mn,kna,knb,Δt_Δnc,Δt_Δn = S.constants.SSPRK3c
+
+            for rki = 1:s       # number of stages
+                if rki > 1
+                    ghost_points_η!(η1,S)
+                end
+
+                # type conversion for mixed precision
+                u1rhs = convert(Diag.PrognosticVarsRHS.u,u1)
+                v1rhs = convert(Diag.PrognosticVarsRHS.v,v1)
+                η1rhs = convert(Diag.PrognosticVarsRHS.η,η1)
+
+                rhs!(u1rhs,v1rhs,η1rhs,Diag,S,t)
+
+                if rki == kn    # special case combining more previous stages  
+                    dxaybzc!(u1,kna,u1,knb,u0,Δt_Δnc,du)
+                    dxaybzc!(v1,kna,v1,knb,v0,Δt_Δnc,dv)
+                else                                # normal update case
+                    axb!(u1,Δt_Δn,du)   
+                    axb!(v1,Δt_Δn,dv)
+                end
+
+                # semi-implicit for continuity equation, use new u1,v1 to calcualte dη
+                ghost_points_uv!(u1,v1,S)
+                u1rhs = convert(Diag.PrognosticVarsRHS.u,u1)
+                v1rhs = convert(Diag.PrognosticVarsRHS.v,v1)
+                continuity!(u1rhs,v1rhs,η1rhs,Diag,S,t)
+
+                if rki == kn
+                    dxaybzc!(η1,kna,η1,knb,η0,Δt_Δnc,dη)
+                else
+                    axb!(η1,Δt_Δn,dη)
+                end
+
+                # special stage that is needed later for the kn-th stage, store in u0,v0,η0 therefore
+                # or for the last step, as u0,v0,η0 is used as the last step's result of any RK scheme.
+                if rki == mn || rki == s
+                    copyto!(u0,u1)
+                    copyto!(v0,v1)
+                    ghost_points_η!(η1,S)
+                    copyto!(η0,η1)
+                end
+            end
 
-        elseif time_scheme == "SSPRK3"   # 4-stage SSPRK3
+        elseif time_scheme == "4SSPRK3"   # 4-stage SSPRK3
 
             for rki = 1:4
                 if rki > 1
@@ -130,7 +175,7 @@ function time_integration(  Prog::PrognosticVars{Tprog},
                 cxab!(v1,1/2,v1,v0)         # same
 
                 # semi-implicit for continuity equation, use u1,v1 to calcualte dη
-                ghost_points!(u1,v1,S)
+                ghost_points_uv!(u1,v1,S)
                 u1rhs = convert(Diag.PrognosticVarsRHS.u,u1)
                 v1rhs = convert(Diag.PrognosticVarsRHS.v,v1)
                 continuity!(u1rhs,v1rhs,η1rhs,Diag,S,t)
@@ -172,7 +217,7 @@ function time_integration(  Prog::PrognosticVars{Tprog},
             bottom_drag!(u0rhs,v0rhs,η0rhs,Diag,S)
             diffusion!(u0rhs,v0rhs,Diag,S)
             add_drag_diff_tendencies!(u0,v0,Diag,S)
-            ghost_points!(u0,v0,S)
+            ghost_points_uv!(u0,v0,S)
         end
 
         t += dtint
@@ -244,7 +289,7 @@ function cxab!(c::Array{T,2},x::Real,a::Array{T,2},b::Array{T,2}) where {T<:Abst
     end
 end
 
-"""c equals add x multiplied to a plus b. c = x*(a+b) """
+"""c = x*a + y*b"""
 function cxayb!(c::Array{T,2},x::Real,a::Array{T,2},y::Real,b::Array{T,2}) where {T<:AbstractFloat}
     m,n = size(a)
     @boundscheck (m,n) == size(b) || throw(BoundsError())
@@ -259,6 +304,27 @@ function cxayb!(c::Array{T,2},x::Real,a::Array{T,2},y::Real,b::Array{T,2}) where
     end
 end
 
+"""d = x*a + y*b + z*c"""
+function dxaybzc!(  d::Array{T,2},
+                    x::Real,a::Array{T,2},
+                    y::Real,b::Array{T,2},
+                    z::Real,c::Array{T,2}) where {T<:AbstractFloat}
+    m,n = size(a)
+    @boundscheck (m,n) == size(b) || throw(BoundsError())
+    @boundscheck (m,n) == size(c) || throw(BoundsError())
+    @boundscheck (m,n) == size(d) || throw(BoundsError())
+
+    xT = T(x)       # convert to type T
+    yT = T(y)
+    zT = T(z)
+
+    @inbounds for j ∈ 1:n
+        for i ∈ 1:m
+           d[i,j] = xT*a[i,j] + yT*b[i,j] + zT*c[i,j]
+        end
+    end
+end
+
 """Convert function for two arrays, X1, X2, in case their eltypes differ.
 Convert every element from X1 and store it in X2."""
 function Base.convert(X2::Array{T2,N},X1::Array{T1,N}) where {T1,T2,N}

test/runtests.jl

@@ -1,4 +1,4 @@
-using ShallowWaters
+# using ShallowWaters
 using Test
 
 @testset "No Forcing" begin
@@ -39,6 +39,12 @@ end
     @test all(Prog.u .!= 0.0f0)
 end
 
+@testset "RK3 tests" begin
+    Prog = RunModel(time_scheme="RK",RKo=3,cfl=0.6,nstep_advcor=0)
+    @test all(abs.(Prog.η) .< 10)    # sea surface height shouldn't exceed +-10m
+    @test all(Prog.u .!= 0.0f0)
+end
+
 @testset "SSPRK2 tests" begin
     Prog = RunModel(time_scheme="SSPRK2",RKs=2,cfl=0.7,nstep_advcor=1)
     @test all(abs.(Prog.η) .< 10)    # sea surface height shouldn't exceed +-10m
@@ -60,17 +66,17 @@ end
     @test all(abs.(Prog.η) .< 10)    # sea surface height shouldn't exceed +-10m
     @test all(Prog.u .!= 0.0f0)
 
-    Prog = RunModel(time_scheme="SSPRK2",RKs=4,cfl=1.2,nstep_advcor=0)
+    Prog = RunModel(time_scheme="SSPRK2",RKs=4,cfl=1.4,nstep_advcor=0)
     @test all(abs.(Prog.η) .< 10)    # sea surface height shouldn't exceed +-10m
     @test all(Prog.u .!= 0.0f0)
 end
 
 @testset "SSPRK3 tests" begin
-    Prog = RunModel(time_scheme="SSPRK3",cfl=1.2,nstep_advcor=1)
+    Prog = RunModel(time_scheme="SSPRK3",RKn=2,cfl=1.2,nstep_advcor=1)
     @test all(abs.(Prog.η) .< 10)    # sea surface height shouldn't exceed +-10m
     @test all(Prog.u .!= 0.0f0)
 
-    Prog = RunModel(time_scheme="SSPRK3",cfl=1.2,nstep_advcor=0)
+    Prog = RunModel(time_scheme="SSPRK3",RKn=3,cfl=2.5,nstep_advcor=1)
     @test all(abs.(Prog.η) .< 10)    # sea surface height shouldn't exceed +-10m
     @test all(Prog.u .!= 0.0f0)
 end