Revive SSPKnoth

docs/make.jl

@@ -12,6 +12,7 @@ pages = [
     "Algorithm Formulations" => [
         "ODE Solvers" => "algorithm_formulations/ode_solvers.md",
         "Newtons Method" => "algorithm_formulations/newtons_method.md",
+        "Rosenbrock Method" => "algorithm_formulations/rosenbrock.md",
         "Old LSRK Formulations" => "algorithm_formulations/lsrk.md",
         "Old MRRK Formulations" => "algorithm_formulations/mrrk.md",
     ],

docs/src/algorithm_formulations/rosenbrock.md

@@ -0,0 +1,122 @@
+# Rosenbrock methods
+
+In this page, we introduce Rosenbrock-type methods to solve ordinary
+differential equations. In doing do, we roughly follow Chapter IV.7 of "Solving
+Ordinary Differential Equations II" by Hairer and Wanner. (Beware, notation is
+not the same.). In this spirit, let us introduce the need for Rosenbrock-type
+methods in the same way as Hairer and Wanner, by quoting Rosenbrock himself:
+
+> When the functions are non-linear, implicit equations can in general be solved
+> only by iteration. This is a severe drawback, as it adds to the problem of
+> stability, that of convergence of the iterative process. An alternative, which
+> avoids this difficulty, is ...
+
+Rosenbrock method!
+
+Before reading this page, we recommend reading the page on ODE solvers first
+[TODO: Add link]
+
+## Introduction to the formalism
+
+Let us consider an ordinary differential equation of the form 
+
+$$ \frac{d}{dt}u(t) = T(u(t), t)\,,$$ 
+
+where $u$ is the state, $T$ is the tendency,
+and $t$ the time. For the sake of simplicity, let us ignore the explicit time
+dependency in the tendency (we will get back to that at the end).
+
+The simplest way to introduce the Rosenbrock method is to start from a
+diagonally implicit Runge-Kutta scheme (see page on DIRK). In an implicit
+Runge-Kutta method with $s$ stages and tableau $a, b, c$, the updated value
+$u_1$ for a step of size $\Delta t$ is obtained starting from the known value $u_0$
+with
+
+$$ u_1 = u_0 + \sum_{i=1}^s b_i k_i\,, $$
+with
+$$ k_i = \Delta t  T ( u_0 + \sum_{j=1}^{i-1}\alpha_{ij}k_j + \alpha_{ii} k_i)\,. $$
+$\alpha_{ij}$, $b_i$ are carefully chosen coefficients.
+
+Rosenbrock's idea consists in linearizing $T$. This operation can be interpreted
+as taking one iteration for Netwon solver and yields
+
+$$ k_i = \Delta t T(g_i) + \Delta t T'(g_i) \alpha_{ii} k_i $$
+
+with $J(g_i)$ Jacobian of the tendency and with 
+
+$$ g_i = u_0 + \sum_{j=1}^{i-1}\alpha_{ij} k_j  $$
+
+In Rosenbrock-type methods, the Jacobian $T'$ is replaced with linear combinations
+of the Jacobian evaluated at $u_0$, $J$:  
+$$ T'(g_i) \alpha_{ii} k_i \approx J \sum_{j=1}^{i-1}\gamma_{ij} k_j + J \gamma_{ii} k_i\,,$$
+with $\gamma_{ij}$ other coefficients that can be chosen. 
+
+Now, each stage consists of solving a system of linear equations in $k_i$ with
+matrix $I - \Delta t \gamma_{ii}$: $$ (I - J \Delta t \gamma_{ii}) k_i = \Delta
+t T(g_i) + J \sum_{j=1}^{i-1}\gamma_{ij} k_j$$ for each $i$ in $1, \dots, s$.
+Once $k_i$ are known, $u_1$ can is easily computed and the process repeated for
+the next step.
+
+In practice, there are computational advantages at implementing a slight
+variation of what presented here.
+
+Let us define $\tilde{k}_{i} = \sum_{j=1}^{i} \gamma_{ij} k_j$. If the matrix
+$\gamma_{ij}$ is invertible, we can move freely from $k_i$ to $\tilde{k}_i$. A
+convenient step is to also define $C$, as 
+
+$$ C = diag(\gamma_{11}^{-1}, \gamma_{22}^{-1}, \dots, \gamma_{ss}^{-1}) - \Gamma^{-1} $$
+
+Which establishes that 
+
+$$ k_i = \frac{\tilde{k}_i}{\gamma_{ii}} - \sum_{j=1}^{i -1} c_{ij} \tilde{k}_j$$
+Substituting this, we obtain the following equations
+
+$$  (J \Delta t \gamma_{ii} - 1) \tilde{k}_i = - \Delta
+t \gamma_{ii} T(g_i) - \gamma_{ii} J \sum_{j=1}^{i-1}c_{ij} \tilde{k}_j \,, $$
+
+$$ g_i =  u_0 + \sum_{j=1}^{i-1}a_{ij} \tilde{k}_j \,,$$
+
+$$ u_1 = u_0 + \sum_{j=1}^{s} m_j \tilde{k}_j\,,$$
+with 
+$$ (a_{ij}) = (\alpha_{ij}) \Gamma^{-1}\,, $$ and $$ (m_i) = (b_i) \Gamma^{-1} $$
+
+Finally, small changes are required to add support for explicit time derivative: 
+
+$$ (J \Delta t \gamma_{ii} - 1) \tilde{k}_i = - \Delta
+t \gamma_{ii} T( t_0 + \alpha_i \Delta t,  g_i) - \gamma_{ii} J \sum_{j=1}^{i-1}c_{ij} \tilde{k}_j - \Delta
+t \gamma_{ii} \gamma_i \Delta
+t \frac{\partial T}{\partial t}(t_0, u_0) $$
+
+where we defined
+
+$\alpha_{i} = \sum_{j=1}^{i-1}\alpha_{ij} $
+
+$ \gamma _{i} = \sum_{j=1}^{i}\gamma_{ij} $
+
+> :note: You may wonder, what is up with all these ugly minus signs? Why don't
+> we cancel them out? The reason for this is because we want to have the
+> prefactor $(J \Delta t \gamma_{ii} - 1)$. This is called `Wfact` in the
+> language of `DifferentialEquations.jl`. Most other algorithms specify this
+> quantity, so it is convenient to be consistent.
+
+## Implementation
+
+In `ClimaTimeSteppers.jl`, we implement the last equation presented in the
+previous section. Currently, the only Rosenbrock-type algorithm implemented is
+`SSPKnoth`. `SSPKnoth` is 3-stage, second-order accurate Rosenbrock with
+
+$$ \alpha = \begin{bmatrix}
+    0 & 0 & 0 \\
+    1 & 0 & 0 \\
+    \frac{1}{4} & \frac{1}{4} & 0 \\
+    \end{bmatrix} $$
+
+$$ \Gamma = \begin{bmatrix}
+    1 & 0 & 0 \\
+    1 & 1 & 0 \\
+    -\frac{3}{4} & \frac{3}{4} & 1 \\
+    \end{bmatrix} $$
+
+and $$ b = (\frac{1}{6}, \frac{1}{6}, \frac{2}{3}) $$.
+
+At each stage, the state is updated to $g_i$ and precomputed quantities are updated. Tendencies are computed and, unless the stage is the last, DSS is applied to the sum of all the tendencies. After the last stage, DSS is applied to the incremented state.

docs/src/plotting_utils.jl

@@ -59,6 +59,7 @@ imex_convergence_orders(::ARK548L2SA2) = (5, 5, 5)
 imex_convergence_orders(::SSP22Heuns) = (2, 2, 2)
 imex_convergence_orders(::SSP33ShuOsher) = (3, 3, 3)
 imex_convergence_orders(::RK4) = (4, 4, 4)
+imex_convergence_orders(::SSPKnoth) = (2, 2, 2)
 
 # Compute a confidence interval for the convergence order, returning the
 # estimated convergence order and its uncertainty.

ext/benchmark_utils.jl

@@ -78,6 +78,7 @@ function get_trial(f, args, name; device, with_cu_prof = :bprofile, trace = fals
 end
 
 get_W(i::CTS.DistributedODEIntegrator) = i.cache.newtons_method_cache.j
+get_W(i::CTS.RosenbrockAlgorithm) = i.cache.newtons_method_cache.j
 get_W(i) = i.cache.W
 f_args(i, f::CTS.ForwardEulerODEFunction) = (copy(i.u), i.u, i.p, i.t, i.dt)
 f_args(i, f) = (similar(i.u), i.u, i.p, i.t)

src/solvers/rosenbrock.jl

@@ -1,93 +1,208 @@
 export SSPKnoth
+using StaticArrays
+import DiffEqBase
+import LinearAlgebra: ldiv!, diagm
 
-abstract type RosenbrockAlgorithm <: DistributedODEAlgorithm end
+abstract type RosenbrockAlgorithmName <: AbstractAlgorithmName end
 
+"""
+    RosenbrockTableau{N, RT, N²}
+
+Contains everything that defines a Rosenbrock-type method.
+
+- N: number of stages,
+- N²: number of stages squared,
+- RT: real type (Float32, Float64, ...)
+
+Refer to the documentation for the precise meaning of the symbols below.
+"""
 struct RosenbrockTableau{N, RT, N²}
+    """A = α Γ⁻¹"""
     A::SMatrix{N, N, RT, N²}
+    """Tableau used for the time-dependent part"""
+    α::SMatrix{N, N, RT, N²}
+    """Stepping matrix. C = 1/diag(Γ) - Γ⁻¹"""
     C::SMatrix{N, N, RT, N²}
+    """Substage contribution matrix"""
     Γ::SMatrix{N, N, RT, N²}
+    """m = b Γ⁻¹, used to compute the increments k"""
     m::SMatrix{N, 1, RT, N}
 end
 
-struct RosenbrockCache{Nstages, RT, N², A}
-    tableau::RosenbrockTableau{Nstages, RT, N²}
+"""
+    RosenbrockAlgorithm(tableau)
+
+Constructs a Rosenbrock algorithm for solving ODEs.
+"""
+struct RosenbrockAlgorithm{T <: RosenbrockTableau} <: ClimaTimeSteppers.DistributedODEAlgorithm
+    tableau::T
+end
+
+"""
+    RosenbrockCache{N, A, WT}
+
+Contains everything that is needed to run a Rosenbrock-type method.
+
+- Nstages: number of stages,
+- A: type of the evolved state (e.g., a ClimaCore.FieldVector),
+- WT: type of the Jacobian (Wfact)
+"""
+struct RosenbrockCache{Nstages, A, WT}
+    """Preallocated space for the state"""
     U::A
+
+    """Preallocated space for the tendency"""
     fU::A
+
+    """Preallocated space for the explicit contribution to the tendency"""
+    fU_exp::A
+
+    """Preallocated space for the limited contribution to the tendency"""
+    fU_lim::A
+
+    """Contributions to the state for each stage"""
     k::NTuple{Nstages, A}
-    W::Any
-    linsolve!::Any
+
+    """Preallocated space for the Wfact, dtγJ - 𝕀, or Wfact_t, 𝕀/dtγ - J, with J the Jacobian of the implicit tendency"""
+    W::WT
+
+    """Preallocated space for the explicit time derivative of the tendency"""
+    ∂Y∂t::A
 end
 
 function init_cache(prob::DiffEqBase.AbstractODEProblem, alg::RosenbrockAlgorithm; kwargs...)
-
-    tab = tableau(alg, eltype(prob.u0))
-    Nstages = length(tab.m)
+    Nstages = length(alg.tableau.m)
     U = zero(prob.u0)
     fU = zero(prob.u0)
+    fU_exp = zero(prob.u0)
+    fU_lim = zero(prob.u0)
+    ∂Y∂t = zero(prob.u0)
     k = ntuple(n -> similar(prob.u0), Nstages)
-    W = prob.f.jac_prototype
-    linsolve! = alg.linsolve(Val{:init}, W, prob.u0; kwargs...)
-
-    return RosenbrockCache(tab, U, fU, k, W, linsolve!)
+    W = prob.f.T_imp!.jac_prototype
+    return RosenbrockCache{Nstages, typeof(U), typeof(W)}(U, fU, fU_exp, fU_lim, k, W, ∂Y∂t)
 end
 
+"""
+    step_u!(int, cache::RosenbrockCache{Nstages})
+
+Take one step with the Rosenbrock-method with the given `cache`.
 
-function step_u!(int, cache::RosenbrockCache{Nstages, RT}) where {Nstages, RT}
-    (; m, Γ, A, C) = cache.tableau
+Some choices are being made here. Most of these are empirically motivated and should be
+revisited on different problems.
+- We do not update dtγ across stages
+- We apply DSS to the sum of the explicit and implicit tendency at all the stages but the last
+- We apply DSS to incremented state (ie, after the final stage is applied)
+"""
+function step_u!(int, cache::RosenbrockCache{Nstages}) where {Nstages}
+    (; m, Γ, A, α, C) = int.alg.tableau
     (; u, p, t, dt) = int
-    (; W, U, fU, k, linsolve!) = cache
-    f! = int.sol.prob.f
-    Wfact_t! = int.sol.prob.f.Wfact_t
+    (; W, U, fU, fU_exp, fU_lim, k, ∂Y∂t) = cache
+    T_imp! = int.sol.prob.f.T_imp!
+    T_exp_lim! = int.sol.prob.f.T_exp_T_lim!
+    tgrad! = int.sol.prob.f.T_imp!.tgrad
+
+    Wfact! = int.sol.prob.f.T_imp!.Wfact
+    Wfact_t! = int.sol.prob.f.T_imp!.Wfact_t
+    if !isnothing(Wfact!) && !isnothing(Wfact_t!)
+        error("Only one between Wfact and Wfact_t can be non-nothing")
+    end
+
+    (; post_explicit!, post_implicit!, dss!) = int.sol.prob.f
+
+    dtγ = dt * Γ[1, 1]
+
+    if isnothing(Wfact_t!)
+        # We have Wfact
+        Wfact!(W, u, p, dtγ, t)
+    else
+        # We have Wfact_t
+        Wfact_t!(W, u, p, dtγ, t)
+    end
+
 
-    # 1) compute jacobian factorization
-    γ = dt * Γ[1, 1]
-    Wfact_t!(W, u, p, γ, t)
     for i in 1:Nstages
+        αi = sum(α[i, 1:(i - 1)])
+
         U .= u
         for j in 1:(i - 1)
             U .+= A[i, j] .* k[j]
         end
-        # TODO: there should be a time modification here (t + c * dt)
-        # if f does depend on time, would need to add tgrad term as well
-        f!(fU, U, p, t)
+
+        # NOTE: post_implicit! is a misnomer
+        post_implicit!(U, p, t)
+
+        if !isnothing(T_imp!)
+            T_imp!(fU, U, p, t)
+        end
+
+        if !isnothing(T_exp_lim!)
+            T_exp_lim!(fU_exp, fU_lim, U, p, t)
+            fU .+= fU_exp
+            fU .+= fU_lim
+        end
+
+        # We dss the tendency at every stage but the last. At the last stage, we
+        # dss the incremented state
+        (i != Nstages) && dss!(fU, p, t)
+
         for j in 1:(i - 1)
             fU .+= (C[i, j] / dt) .* k[j]
         end
-        linsolve!(k[i], W, fU)
+
+        if !isnothing(tgrad!)
+            tgrad!(∂Y∂t, u, p, t)
+            fU .+= αi .* dt .* ∂Y∂t
+        end
+
+        if isnothing(Wfact_t!)
+            fU .*= -dtγ
+        end
+
+        if W isa Matrix
+            ldiv!(k[i], lu(W), fU)
+        else
+            ldiv!(k[i], W, fU)
+        end
     end
+
     for i in 1:Nstages
         u .+= m[i] .* k[i]
     end
-end
 
-struct SSPKnoth{L} <: RosenbrockAlgorithm
-    linsolve::L
+    dss!(u, p, t)
+    return nothing
 end
-SSPKnoth(; linsolve) = SSPKnoth(linsolve)
 
+"""
+    SSPKnoth
+
+`SSPKnoth` is a second-order Rosenbrock method developed by Osward Knoth.
+
+The coefficients are the same as in `CGDycore.jl`, except that for C we add the
+diagonal elements too. Note, however, that the elements on the diagonal of C do
+not really matter because C is only used in its lower triangular part. We add them
+mostly to match literature on the subject
+"""
+struct SSPKnoth end
 
 function tableau(::SSPKnoth, RT)
-    # ROS.transformed=true;
     N = 3
     N² = N * N
     α = @SMatrix RT[
         0 0 0
         1 0 0
         1/4 1/4 0
     ]
-    # ROS.d=ROS.alpha*ones(ROS.nStage,1);
     b = @SMatrix RT[1 / 6 1 / 6 2 / 3]
     Γ = @SMatrix RT[
         1 0 0
         0 1 0
         -3/4 -3/4 1
     ]
     A = α / Γ
-    C = -inv(Γ)
+    invΓ = inv(Γ)
+    diag_invΓ = SMatrix{3, 3}{RT}(diagm([invΓ[i, i] for i in 1:3]))
+    C = diag_invΓ .- inv(Γ)
     m = b / Γ
-    return RosenbrockTableau{N, RT, N²}(A, C, Γ, m)
-    #   ROS.SSP.alpha=[1 0 0
-    #                  3/4 1/4 0
-    #                  1/3 0 2/3];
-
+    return RosenbrockTableau{N, RT, N²}(A, α, C, Γ, m)
 end

test/Project.toml

@@ -3,6 +3,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
+ClimaCorePlots = "cf7c7e5a-b407-4c48-9047-11a94a308626"
 ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"