Add Generic Rosenbrock Algorithm

src/ClimaTimeSteppers.jl

@@ -70,21 +70,23 @@ end
 SciMLBase.allowscomplex(alg::DistributedODEAlgorithm) = true
 include("integrators.jl")
 
+include("solvers/update_signal_handler.jl")
 include("solvers/convergence_condition.jl")
 include("solvers/convergence_checker.jl")
 include("solvers/newtons_method.jl")
 include("solvers/imex_ark.jl")
+include("solvers/rosenbrock.jl")
 
 # Include concrete implementations
 include("solvers/imex_ark_tableaus.jl")
+include("solvers/rosenbrock_tableaus.jl")
 include("solvers/multirate.jl")
 include("solvers/lsrk.jl")
 include("solvers/ssprk.jl")
 include("solvers/ark.jl")
 # include("solvers/ars.jl") # previous implementations of ARS schemes
 include("solvers/mis.jl")
 include("solvers/wickerskamarock.jl")
-include("solvers/rosenbrock.jl")
 
 include("callbacks.jl")
 

src/functions.jl

@@ -9,9 +9,19 @@ un .= u .+ dt * f(u, p, t)
 ```
 
 """
-struct ForwardEulerODEFunction{F} <: DiffEqBase.AbstractODEFunction{true}
+struct ForwardEulerODEFunction{F, J, W, T} <:
+    DiffEqBase.AbstractODEFunction{true}
     f::F
+    jac_prototype::J
+    Wfact::W
+    tgrad::T
 end
+ForwardEulerODEFunction(
+    f;
+    jac_prototype = nothing,
+    Wfact = nothing,
+    tgrad = nothing,
+) = ForwardEulerODEFunction(f, jac_prototype, Wfact, tgrad)
 (f::ForwardEulerODEFunction{F})(un, u, p, t, dt) where {F} = f.f(un, u, p, t, dt)
 
 # Don't wrap a ForwardEulerODEFunction in an ODEFunction.

src/solvers/imex_ark.jl

@@ -15,7 +15,7 @@ The abscissae are often defined as c_χ[i] := ∑_{j=1}^s a_χ[i,j] for the expl
 and implicit methods to be "internally consistent", with c_exp[i] = c_imp[i] for
 the overall IMEX method to be "internally consistent", but this is not required.
 If the weights are defined as b_χ[j] := a_χ[s,j], then u_next = U[s]; i.e., the
-method is FSAL (first same as last).
+method is FSAL (first same as last), which means that it is "stiffly accurate".
 
 To simplify our notation, let
     a_χ[s+1,j] := b_χ[j] ∀ j ∈ 1:s,

src/solvers/old_rosenbrock.jl

@@ -0,0 +1,102 @@
+export SSPKnoth
+
+abstract type RosenbrockAlgorithm <: DistributedODEAlgorithm end
+
+struct RosenbrockTableau{N, RT, N²}
+    A::SMatrix{N, N, RT, N²}
+    C::SMatrix{N, N, RT, N²}
+    Γ::SMatrix{N, N, RT, N²}
+    m::SMatrix{N, 1, RT, N}
+end
+
+struct RosenbrockCache{Nstages, RT, N², A}
+    tableau::RosenbrockTableau{Nstages, RT, N²}
+    U::A
+    fU::A
+    k::NTuple{Nstages, A}
+    W
+    linsolve!
+end
+
+function cache(
+    prob::DiffEqBase.AbstractODEProblem,
+    alg::RosenbrockAlgorithm; kwargs...)
+
+    tab = tableau(alg, eltype(prob.u0))
+    Nstages = length(tab.m)
+    U = zero(prob.u0)
+    fU = 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!)
+end
+
+
+function step_u!(int, cache::RosenbrockCache{Nstages, RT}) where {Nstages, RT}
+    tab = cache.tableau
+
+    f! = int.prob.f
+    Wfact_t! = int.prob.f.Wfact_t
+
+    u = int.u
+    p = int.p
+    t = int.t
+    dt = int.dt
+    W = cache.W
+    U = cache.U
+    fU = cache.fU
+    k = cache.k
+    linsolve! = cache.linsolve!
+
+    # 1) compute jacobian factorization
+    γ = dt * tab.Γ[1,1]
+    Wfact_t!(W, u, p, γ, t)
+    for i in 1:Nstages
+        U .= u
+        for j = 1:i-1
+            U .+= tab.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)
+        for j = 1:i-1
+            fU .+= (tab.C[i,j] / dt) .* k[j]
+        end
+        linsolve!(k[i], W, fU)
+    end
+    for i = 1:Nstages
+        u .+= tab.m[i] .* k[i]
+    end
+end
+
+struct SSPKnoth{L} <: RosenbrockAlgorithm
+    linsolve::L
+end
+SSPKnoth(;linsolve)=SSPKnoth(linsolve)
+
+
+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(Γ)
+    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];
+
+end