Add Generic Rosenbrock Algorithm

.buildkite/pipeline.yml

@@ -6,10 +6,18 @@ env:
   CLIMATEMACHINE_SETTINGS_FIX_RNG_SEED: "true"
 
 steps:
-  - label: "init perf env"
-    key: "init_perf_env"
+  - label: "init cpu env"
+    key: "init_cpu_env"
     command:
       - "echo $JULIA_DEPOT_PATH"
+
+      - echo "--- Instantiate test"
+      - "julia --project=test -e 'using Pkg; Pkg.develop(path=\".\")'"
+      - "julia --project=test -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
+      - "julia --project=test -e 'using Pkg; Pkg.precompile(;strict=true)'"
+      - "julia --project=perf -e 'using Pkg; Pkg.status()'"
+
+      - echo "--- Instantiate perf"
       - "julia --project=perf -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
       - "julia --project=perf -e 'using Pkg; Pkg.precompile()'"
       - "julia --project=perf -e 'using Pkg; Pkg.status()'"
@@ -35,6 +43,15 @@ steps:
 
   - wait
 
+  - label: "Unit tests and plots"
+    command:
+      - "julia --project -e 'using Pkg; Pkg.test()'"
+    artifact_paths: "test/output/*"
+    agents:
+      config: cpu
+      queue: central
+      slurm_ntasks: 1
+
   - label: "GPU tests"
     command:
       - "julia --project -e 'using Pkg; Pkg.test(;test_args=[\"CuArray\"])'"

src/ClimaTimeSteppers.jl

@@ -70,21 +70,24 @@ end
 SciMLBase.allowscomplex(alg::DistributedODEAlgorithm) = true
 include("integrators.jl")
 
+include("solvers/matrix_utilities.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/integrators.jl

@@ -29,7 +29,7 @@ function DiffEqBase.__init(
     args...;
     dt,  # required
     stepstop=-1,
-    adjustfinal=false,
+    adjustfinal=true,
     callback=nothing,
     save_func=(u,t,integrator)->copy(u),
     saveat=nothing,

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/matrix_utilities.jl

@@ -0,0 +1,33 @@
+lower_plus_diagonal(matrix::T) where {T} =
+    T(LinearAlgebra.LowerTriangular(matrix))
+diagonal(matrix::T) where {T} = T(LinearAlgebra.Diagonal(matrix))
+lower(matrix) = lower_plus_diagonal(matrix) - diagonal(matrix)
+
+lower_triangular_inv(matrix::T) where {T} =
+    T(inv(LinearAlgebra.LowerTriangular(matrix)))
+
+to_enumerated_rows(x) = x
+function to_enumerated_rows(matrix::AbstractMatrix)
+    rows = tuple(1:size(matrix, 1)...)
+    nonzero_indices = map(i -> findall(matrix[i, :] .!= 0), rows)
+    enumerated_rows = map(
+        i -> tuple(zip(nonzero_indices[i], matrix[i, nonzero_indices[i]])...),
+        rows,
+    )
+    return enumerated_rows
+end
+
+linear_combination_terms(enumerated_row, vectors) =
+    map(((j, val),) -> broadcasted(*, val, vectors[j]), enumerated_row)
+function linear_combination(enumerated_row, vectors)
+    length(enumerated_row) == 0 && return nothing
+    terms = linear_combination_terms(enumerated_row, vectors)
+    length(enumerated_row) == 1 && return terms[1]
+    return broadcasted(+, terms...)
+end
+
+@generated foreachval(f::F, ::Val{N}) where {F, N} =
+    quote
+        Base.@nexprs $N i -> f(Val(i))
+        return nothing
+    end

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