SSPRK2 and 3 implemented

milankl · milankl · commit 7dba5fd0ff53 · 2020-11-19T18:38:57.000Z
diff --git a/src/Constants.jl b/src/Constants.jl
@@ -1,8 +1,11 @@
 struct Constants{T<:AbstractFloat,Tprog<:AbstractFloat}
 
-    # RUNGE-KUTTA COEFFICIENTS 3rd/4th order including timestep Δt
+    # RUNGE-KUTTA COEFFICIENTS 2nd/3rd/4th order including timestep Δt
     RKaΔt::Array{Tprog,1}
     RKbΔt::Array{Tprog,1}
+    Δ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/Δ
 
     # BOUNDARY CONDITIONS
     one_minus_α::Tprog      # tangential boundary condition for the ghost-point copy
@@ -39,6 +42,13 @@ function Constants{T,Tprog}(P::Parameter,G::Grid) where {T<:AbstractFloat,Tprog<
         RKbΔt = Tprog.([.5,.5,1.]*G.dtint/G.Δ)
     end
 
+    # Δt/(s-1) for SSPRK2
+    Δt_Δs = Tprog(G.dtint/G.Δ/(P.RKs-1))
+
+    # time step and half the time step including the grid spacing as this is not included in the RHS
+    Δ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
 
@@ -65,7 +75,8 @@ function Constants{T,Tprog}(P::Parameter,G::Grid) where {T<:AbstractFloat,Tprog<
     ωFx = 2π*P.ωFx/24/365.25/3600
     ωFy = 2π*P.ωFy/24/365.25/3600
 
-    return Constants{T,Tprog}(  RKaΔt,RKbΔt,one_minus_α,
+    return Constants{T,Tprog}(  RKaΔt,RKbΔt,Δt_Δs,Δt_Δ,Δt_Δ_half,
+                                one_minus_α,
                                 g,cD,rD,γ,cSmag,νB,rSST,
                                 jSST,SSTmin,ωFη,ωFx,ωFy)
 end
diff --git a/src/Continuity.jl b/src/Continuity.jl
@@ -69,16 +69,31 @@ end
 
 
 """Transit function to call the specified continuity function."""
-function continuity!(   η::AbstractMatrix,
+function continuity!(   u::AbstractMatrix,
+                        v::AbstractMatrix,
+                        η::AbstractMatrix,
                         Diag::DiagnosticVars,
                         S::ModelSetup,
                         t::Int)
+    
+    @unpack U,V,dUdx,dVdy = Diag.VolumeFluxes
+    @unpack nstep_advcor = S.grid
+    @unpack time_scheme,surface_relax,surface_forcing = S.parameters
 
-    if S.parameters.surface_relax
+    if nstep_advcor > 0 ||              # then UV wasn't calculated inside rhs!
+        time_scheme != "RK"             # then u,v changed before evaluating semi-implciti continuity
+        UVfluxes!(u,v,η,Diag,S)
+    end
+
+    # divergence of mass flux
+    ∂x!(dUdx,U)
+    ∂y!(dVdy,V)
+
+    if surface_relax
         continuity_surf_relax!(η,Diag,S,t)
-    elseif S.parameters.surface_forcing
+    elseif surface_forcing
         continuity_forcing!(Diag,S,t)
     else
         continuity_itself!(Diag,S,t)
     end
-end
+end
diff --git a/src/DefaultParameters.jl b/src/DefaultParameters.jl
@@ -47,7 +47,9 @@
     wk::Real=10e3                       # width [m] in y of Gaussian used for surface forcing
 
     # TIME STEPPING OPTIONS
-    RKo::Int=4                          # Order of the RK time stepping scheme (3 or 4)
+    time_scheme::String="RK"            # Runge-Kutta ("RK") or strong-stability preserving RK ("SSPRK2","SSPRK3")
+    RKo::Int=4                          # Order of the RK time stepping scheme (2, 3 or 4)
+    RKs::Int=2                          # Number of stages for SSPRK2
     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.
@@ -118,19 +120,21 @@
     @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 RKo in [2,3,4]        "RKo has to be 2,3 or 4, $RKo given."
+    @assert time_scheme in ["RK","SSPRK2","SSPRK3"] "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."
     @assert nstep_diff > 0      "nstep_diff has to be >0, $nstep_diff given."
     @assert nstep_advcor >= 0   "nstep_advcor has to be >=0, $nstep_advcor given."
     @assert bc in ["periodic","nonperiodic"]    "boundary condition '$bc' unsupported."
     @assert α >= 0.0 && α <= 2.0    "Tangential boundary condition α has to be in (0,2), $α given."
     @assert adv_scheme in ["Sadourny","ArakawaHsu"] "Advection scheme '$adv_scheme' unsupported"
-    @assert dynamics in ["linear","nonlinear"]  "Dynamics '$dynamics' unsupoorted."
+    @assert dynamics in ["linear","nonlinear"]  "Dynamics '$dynamics' unsupported."
     @assert bottom_drag in ["quadratic","linear","none"] "Bottom drag '$bottom_drag' unsupported."
     @assert cD >= 0.0    "Bottom drag coefficient cD has to be >=0, $cD given."
     @assert τD >= 0.0    "Bottom drag coefficient τD has to be >=0, $τD given."
-    @assert diffusion in ["Smagorinsky", "constant"] "Diffusion '$diffusion' unsupoorted."
-    @assert νB > 0.0     "Diffusion scaling constant νB has to be > 0, $νB given."
+    @assert diffusion in ["Smagorinsky", "constant"] "Diffusion '$diffusion' unsupported."
+    @assert νB > 0.0     "Diffusion scaling constant νB has to be >0, $νB given."
     @assert cSmag > 0.0  "Smagorinsky coefficient cSmag has to be >0, $cSmag given."
     @assert injection_region in ["west","south"] "Injection region '$injection_region' unsupported."
     @assert Uadv > 0.0   "Advection velocity scale Uadv has to be >0, $Uadv given."
diff --git a/src/GhostPoints.jl b/src/GhostPoints.jl
@@ -176,6 +176,24 @@ function ghost_points!( u::AbstractMatrix,
     end
 end
 
+"""Decide on boundary condition P.bc which ghost point function to execute."""
+function ghost_points!( u::AbstractMatrix,
+                        v::AbstractMatrix,
+                        S::ModelSetup)
+
+    @unpack bc = S.parameters
+    C = S.constants
+
+    if bc == "periodic"
+        @unpack Tcomm = S.parameters
+        ghost_points_u_periodic!(Tcomm,u,C)
+        ghost_points_v_periodic!(Tcomm,v)
+    else
+        ghost_points_u_nonperiodic!(C,u)
+        ghost_points_v_nonperiodic!(C,v)
+    end
+end
+
 """Decide on boundary condition P.bc which ghost point function to execute."""
 function ghost_points_sst!(sst::AbstractMatrix,S::ModelSetup)
 
diff --git a/src/InitialConditions.jl b/src/InitialConditions.jl
@@ -68,12 +68,10 @@ function initial_conditions(::Type{T},S::ModelSetup) where {T<:AbstractFloat}
         v = reshape(v,size(v)[1:2])
         η = reshape(η,size(η)[1:2])
 
-        # Interpolation in case the grids don't match?
-
         nx_old,ny_old = size(η)
 
         if (nx_old,ny_old) != (nx,ny)
-            if init_interpolation
+            if init_interpolation           # interpolation in case the grids don't match
 
                 # old grids
                 x_T = collect(0.5:nx_old-0.5)
diff --git a/src/TimeIntegration.jl b/src/TimeIntegration.jl
@@ -9,16 +9,21 @@ function time_integration(  Prog::PrognosticVars{Tprog},
     @unpack du,dv,dη = Diag.Tendencies
     @unpack um,vm = Diag.SemiLagrange
 
-    @unpack dynamics,RKo,tracer_advection = S.parameters
+    @unpack dynamics,RKo,RKs,tracer_advection = S.parameters
+    @unpack time_scheme = S.parameters
     @unpack RKaΔt,RKbΔt = S.constants
+    @unpack Δt_Δ,Δt_Δs = S.constants
 
     @unpack nt,dtint = S.grid
     @unpack nstep_advcor,nstep_diff,nadvstep,nadvstep_half = S.grid
 
-    if dynamics == "linear"
-        Ix!(Diag.VolumeFluxes.h_u,S.forcing.H)
-        Iy!(Diag.VolumeFluxes.h_v,S.forcing.H)
-    end
+    # some precalculations
+    thickness!(Diag.VolumeFluxes.h,η,S.forcing.H)
+    Ix!(Diag.VolumeFluxes.h_u,Diag.VolumeFluxes.h)
+    Iy!(Diag.VolumeFluxes.h_v,Diag.VolumeFluxes.h)
+    Ixy!(Diag.Vorticity.h_q,Diag.VolumeFluxes.h)
+    advection_coriolis!(u,v,η,Diag,S)
+    PVadvection!(Diag,S)
 
     # propagate initial conditions
     copyto!(u0,u)
@@ -40,29 +45,104 @@ function time_integration(  Prog::PrognosticVars{Tprog},
         copyto!(v1,v)
         copyto!(η1,η)
 
-        # Runge-Kutta 4th order / 3rd order
-        for rki = 1:RKo
-            if rki > 1
-                ghost_points!(u1,v1,η1,S)
+        if time_scheme == "RK"   # classic RK4,3 or 2
+
+            for rki = 1:RKo
+                if rki > 1
+                    ghost_points!(u1,v1,η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)          # momentum only
+                continuity!(u1rhs,v1rhs,η1rhs,Diag,S,t)   # continuity equation 
+
+                if rki < RKo
+                    caxb!(u1,u,RKbΔt[rki],du)   #u1 .= u .+ RKb[rki]*Δt*du
+                    caxb!(v1,v,RKbΔt[rki],dv)   #v1 .= v .+ RKb[rki]*Δt*dv
+                    caxb!(η1,η,RKbΔt[rki],dη)   #η1 .= η .+ RKb[rki]*Δt*dη
+                end
+
+                # sum RK-substeps on the go
+                axb!(u0,RKaΔt[rki],du)          #u0 .+= RKa[rki]*Δt*du
+                axb!(v0,RKaΔt[rki],dv)          #v0 .+= RKa[rki]*Δt*dv
+                axb!(η0,RKaΔt[rki],dη)          #η0 .+= RKa[rki]*Δt*dη
             end
 
-            # type conversion for mixed precision
-            u1rhs = convert(Diag.PrognosticVarsRHS.u,u1)
-            v1rhs = convert(Diag.PrognosticVarsRHS.v,v1)
-            η1rhs = convert(Diag.PrognosticVarsRHS.η,η1)
+        elseif time_scheme == "SSPRK2"  # s-stage 2nd order SSPRK
+
+            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)
+                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)
 
-            if rki < RKo
-                caxb!(u1,u,RKbΔt[rki],du)   #u1 .= u .+ RKb[rki]*Δt*du
-                caxb!(v1,v,RKbΔt[rki],dv)   #v1 .= v .+ RKb[rki]*Δt*dv
-                caxb!(η1,η,RKbΔt[rki],dη)   #η1 .= η .+ RKb[rki]*Δt*dη
+                # 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)
+                u1rhs = convert(Diag.PrognosticVarsRHS.u,u1)
+                v1rhs = convert(Diag.PrognosticVarsRHS.v,v1)
+                continuity!(u1rhs,v1rhs,η1rhs,Diag,S,t)
+                axb!(η1,Δt_Δs,dη)       # η1 = η1 + Δt/(s-1)*RHS(u1)
             end
 
-            # sum RK-substeps on the go
-            axb!(u0,RKaΔt[rki],du)          #u0 .+= RKa[rki]*Δt*du
-            axb!(v0,RKaΔt[rki],dv)          #v0 .+= RKa[rki]*Δt*dv
-            axb!(η0,RKaΔt[rki],dη)          #η0 .+= RKa[rki]*Δt*dη
+            a = 1/RKs
+            b = (RKs-1)/RKs
+            cxayb!(u0,a,u,b,u1)
+            cxayb!(v0,a,v,b,v1)
+            cxayb!(η0,a,η,b,η1)
+            
+        elseif time_scheme == "SSPRK3"   # 4-stage SSPRK3
+        
+            for rki = 1:4
+                if rki > 1
+                    ghost_points!(u1,v1,η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)
+
+                caxb!(u0,u1,Δt_Δ,du)        # store Euler update into u0,v0
+                caxb!(v0,v1,Δt_Δ,dv)
+                cxab!(u1,1/2,u1,u0)         # average u0,u1 and store in u1
+                cxab!(v1,1/2,v1,v0)         # same
+
+                # semi-implicit for continuity equation, use u1,v1 to calcualte dη
+                ghost_points!(u1,v1,S)
+                u1rhs = convert(Diag.PrognosticVarsRHS.u,u1)
+                v1rhs = convert(Diag.PrognosticVarsRHS.v,v1)
+                continuity!(u1rhs,v1rhs,η1rhs,Diag,S,t)
+                
+                caxb!(η0,η1,Δt_Δ,dη)    # store Euler update into η0
+                cxab!(η1,1/2,η1,η0)         # average η0,η1 and store in η1
+
+                if rki == 3
+                    cxayb!(u1,2/3,u,1/3,u1)
+                    cxayb!(v1,2/3,v,1/3,v1)
+                    cxayb!(η1,2/3,η,1/3,η1)
+                elseif rki == 4
+                    copyto!(u0,u1)
+                    copyto!(v0,v1)
+                    copyto!(η0,η1)
+                end
+            end
         end
 
         ghost_points!(u0,v0,η0,S)
@@ -76,7 +156,9 @@ function time_integration(  Prog::PrognosticVars{Tprog},
         # although included in the tendency of every RK substep,
         # only update every nstep_advcor steps if nstep_advcor > 0
         if dynamics == "nonlinear" && nstep_advcor > 0 && (i % nstep_advcor) == 0
+            UVfluxes!(u0rhs,v0rhs,η0rhs,Diag,S)
             advection_coriolis!(u0rhs,v0rhs,η0rhs,Diag,S)
+            PVadvection!(Diag,S)
         end
 
         # DIFFUSIVE TERMS - SEMI-IMPLICIT EULER
@@ -86,23 +168,26 @@ 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)
         end
 
         t += dtint
 
         # TRACER ADVECTION
+        u0rhs = convert(Diag.PrognosticVarsRHS.u,u0)  # copy back as add_drag_diff_tendencies changed u0,v0
+        v0rhs = convert(Diag.PrognosticVarsRHS.v,v0)
         tracer!(i,u0rhs,v0rhs,Prog,Diag,S)
 
         # feedback and output
         feedback.i = i
         feedback!(Prog,feedback,S)
-        output_nc!(i,netCDFfiles,Prog,Diag,S)
+        output_nc!(i,netCDFfiles,Prog,Diag,S)       # uses u0,v0,η0
 
         if feedback.nans_detected
             break
         end
 
-        # RK3/4 copy back from substeps
+        # Copy back from substeps
         copyto!(u,u0)
         copyto!(v,v0)
         copyto!(η,η0)
@@ -134,14 +219,42 @@ function caxb!(c::Array{T,2},a::Array{T,2},x::T,b::Array{T,2}) where {T<:Abstrac
     @boundscheck (m,n) == size(b) || throw(BoundsError())
     @boundscheck (m,n) == size(c) || throw(BoundsError())
 
-    #TODO @simd?
     @inbounds for j ∈ 1:n
         for i ∈ 1:m
            c[i,j] = a[i,j] + x*b[i,j]
         end
     end
 end
 
+"""c equals add x multiplied to a plus b. c = x*(a+b) """
+function cxab!(c::Array{T,2},x::Real,a::Array{T,2},b::Array{T,2}) where {T<:AbstractFloat}
+    m,n = size(a)
+    @boundscheck (m,n) == size(b) || throw(BoundsError())
+    @boundscheck (m,n) == size(c) || throw(BoundsError())
+
+    xT = T(x)
+    @inbounds for j ∈ 1:n
+        for i ∈ 1:m
+           c[i,j] = xT*(a[i,j] + b[i,j])
+        end
+    end
+end
+
+"""c equals add x multiplied to a plus b. c = x*(a+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())
+    @boundscheck (m,n) == size(c) || throw(BoundsError())
+
+    xT = T(x)
+    yT = T(y)
+    @inbounds for j ∈ 1:n
+        for i ∈ 1:m
+           c[i,j] = xT*a[i,j] + yT*b[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}
diff --git a/src/rhs.jl b/src/rhs.jl
