Merge #133

133: Refactor, unpack efficiently, unify n_stages nstages r=charleskawczynski a=charleskawczynski

This PR applies a few refactoring changes:
 - Changes syntax `a = x.a` to `(; a) = x` for several variables
 - Adds a `n_stages` method for some caches
 - Renames `nstages` to `n_stages`


Co-authored-by: Charles Kawczynski <[email protected]>

Author: bors[bot]

src/solvers/ark.jl

@@ -50,15 +50,16 @@ function init_cache(
 ) where {uType, tType}
 
     tab = tableau(alg, eltype(prob.u0))
-    Nstages = length(tab.B) # TODO: create function for this
+    (; B, Aimpl) = tab
+    Nstages = length(B) # TODO: create function for this
     U = zero(prob.u0)
     L = ntuple(i -> zero(prob.u0), Nstages)
     R = ntuple(i -> zero(prob.u0), Nstages)
 
     if prob.f isa DiffEqBase.ODEFunction
-        W = EulerOperator(prob.f.jvp, -dt * tab.Aimpl[2, 2], prob.p, prob.tspan[1])
+        W = EulerOperator(prob.f.jvp, -dt * Aimpl[2, 2], prob.p, prob.tspan[1])
     elseif prob.f isa DiffEqBase.SplitFunction
-        W = EulerOperator(prob.f.f1, -dt * tab.Aimpl[2, 2], prob.p, prob.tspan[1])
+        W = EulerOperator(prob.f.f1, -dt * Aimpl[2, 2], prob.p, prob.tspan[1])
     end
     linsolve! = alg.linsolve(Val{:init}, W, prob.u0; kwargs...)
 
@@ -108,95 +109,88 @@ function step_u!(int, cache::AdditiveRungeKuttaFullCache{Nstages}) where {Nstage
 end
 
 function step_u!(int, cache::AdditiveRungeKuttaFullCache{Nstages}, f::DiffEqBase.SplitFunction) where {Nstages}
-    tab = cache.tableau
-    U = cache.U
-    Uhat = cache.R[end] # can be used as work array, as only used in last stage
 
+    (; C, Aimpl, Aexpl, B) = cache.tableau
+    (; U, R, L, W, linsolve) = cache
+    (; u, p, t, dt) = int
 
-    u = int.u
-    p = int.p
-    t = int.t
-    dt = int.dt
+    U = U
+    Uhat = R[end] # can be used as work array, as only used in last stage
 
     fL! = f.f1 # linear part
     fR! = f.f2 # remainder
 
     # first stage is always explicit
-    τ = t + tab.C[1] * dt
+    τ = t + C[1] * dt
 
-    #  cache.L[i] .= fL(cache.U[i-1], p, t + tab.C[i]*dt)
-    #  cache.R[1] .= fR(cache.U[i-1], p, t + tab.C[i]*dt)
+    #  L[i] .= fL(U[i-1], p, t + C[i]*dt)
+    #  R[1] .= fR(U[i-1], p, t + C[i]*dt)
 
-    fL!(cache.L[1], u, p, τ)
-    fR!(cache.R[1], u, p, τ)
+    fL!(L[1], u, p, τ)
+    fR!(R[1], u, p, τ)
 
     for i in 2:Nstages
         # solve for W * U = Uhat
         # set U to initial guess based on fully explicit
         # TODO: we don't need this for direct solves
         U .= Uhat .= u
         for j in 1:(i - 1)
-            Uhat .+= (dt * tab.Aimpl[i, j]) .* cache.L[j] .+ (dt * tab.Aexpl[i, j]) .* cache.R[j]
+            Uhat .+= (dt * Aimpl[i, j]) .* L[j] .+ (dt * Aexpl[i, j]) .* R[j]
             # initial value: we only need to do this if using an iterative method:
-            U .+= (dt * tab.Aexpl[i, j]) .* (cache.L[j] .+ cache.R[j])
+            U .+= (dt * Aexpl[i, j]) .* (L[j] .+ R[j])
         end
 
         #  W = I - dt * Aimpl[i,i] * L
         # currently only use SDIRK methods where
         #    Aimpl[i,i] = i == 1 ? 0 : const
         # TODO: handle changing dt & Aimpl[i,i] coeffs
-        if !(DiffEqBase.isconstant(cache.W))
-            cache.W.t = τ
+        if !(DiffEqBase.isconstant(W))
+            W.t = τ
             W_updated = true
         end
-        γ = -dt * tab.Aimpl[i, i]
-        if cache.W.γ != γ
-            cache.W.γ = γ
+        γ = -dt * Aimpl[i, i]
+        if W.γ != γ
+            W.γ = γ
             W_updated = true
         end
-        cache.linsolve!(U, cache.W, Uhat, W_updated)
+        linsolve!(U, W, Uhat, W_updated)
 
-        τ = t + tab.C[i] * dt
-        fL!(cache.L[i], U, p, τ)
-        # or use cache.L[i] .= (U .- Uhat) ./ (dt * tab.Aimpl[i,i]) ?
-        fR!(cache.R[i], U, p, τ)
+        τ = t + C[i] * dt
+        fL!(L[i], U, p, τ)
+        # or use L[i] .= (U .- Uhat) ./ (dt * Aimpl[i,i]) ?
+        fR!(R[i], U, p, τ)
     end
 
     # compute next step
     for i in 1:Nstages
-        u .+= (dt * tab.B[i]) .* (cache.L[i] .+ cache.R[i])
+        u .+= (dt * B[i]) .* (L[i] .+ R[i])
     end
 end
 
 # WIP
 function step_u!(int, cache::AdditiveRungeKuttaFullCache{Nstages}, f::DiffEqBase.ODEFunction) where {Nstages}
-    tab = cache.tableau
-    U = cache.U
-    Uhat = cache.R[end] # can be used as work array, as only used in last stage
-
-
-    u = int.u
-    p = int.p
-    t = int.t
-    dt = int.dt
+    (; C, Aimpl, Aexpl, B) = cache.tableau
+    (; u, p, t, dt) = int
+    (; U, R, L, W, linsolve) = cache
+    Uhat = R[end] # can be used as work array, as only used in last stage
 
     fL! = f.jvp # linear part
     f! = f.f # remainder
 
     # first stage is always explicit
-    τ = t + tab.C[1] * dt
+    τ = t + C[1] * dt
 
-    f!(cache.R[1], u, p, τ)
-    fL!(cache.L[1], u, p, τ)
+    f!(R[1], u, p, τ)
+    fL!(L[1], u, p, τ)
     for i in 2:Nstages
         # solve for W * U = Uhat
         # set U to initial guess based on fully explicit
         # TODO: we don't need this for direct solves
         U .= Uhat .= u
         for j in 1:(i - 1)
-            Uhat .+= (dt * tab.Aimpl[i, j]) .* cache.L[j] .+ (dt * tab.Aexpl[i, j]) .* (cache.R[j] .- cache.L[j])
+            Uhat .+= (dt * Aimpl[i, j]) .* L[j] .+ (dt * Aexpl[i, j]) .* (R[j] .- L[j])
             # initial value: we only need to do this if using an iterative method:
-            U .+= (dt * tab.Aexpl[i, j]) .* cache.R[j]
+            U .+= (dt * Aexpl[i, j]) .* R[j]
         end
         #=
                 # solve for W * U = Uhat
@@ -205,36 +199,36 @@ function step_u!(int, cache::AdditiveRungeKuttaFullCache{Nstages}, f::DiffEqBase
                 V .= u
                 Ω .= 0
                 for j = 1:i-1
-                    V .+= dt * tab.Aexpl[i,j] .* cache.R[j]
-                    Ω .+= (tab.Aimpl[i,j]-tab.Aexpl[i,j])/tab.Aimpl[i,i] .* cache.U[j]
+                    V .+= dt * Aexpl[i,j] .* R[j]
+                    Ω .+= (Aimpl[i,j]-Aexpl[i,j])/Aimpl[i,i] .* U[j]
                 end
                 Vhat .= V .+ Ω
         =#
         #  W = I - dt * Aimpl[i,i] * L
         # currently only use SDIRK methods where
         #    Aimpl[i,i] = i == 1 ? 0 : const
         # TODO: handle changing dt & Aimpl[i,i] coeffs
-        if !(DiffEqBase.isconstant(cache.W))
-            cache.W.t = τ
+        if !(DiffEqBase.isconstant(W))
+            W.t = τ
             W_updated = true
         end
-        γ = -dt * tab.Aimpl[i, i]
-        if cache.W.γ != γ
-            cache.W.γ = γ
+        γ = -dt * Aimpl[i, i]
+        if W.γ != γ
+            W.γ = γ
             W_updated = true
         end
-        cache.linsolve!(U, cache.W, Uhat, W_updated)
+        linsolve!(U, W, Uhat, W_updated)
         # U = V .- Ω
 
-        τ = t + tab.C[i] * dt
-        fL!(cache.L[i], U, p, τ)
-        # or use cache.L[i] .= (U .- Uhat) ./ (dt * tab.Aimpl[i,i]) ?
-        f!(cache.R[i], U, p, τ)
+        τ = t + C[i] * dt
+        fL!(L[i], U, p, τ)
+        # or use L[i] .= (U .- Uhat) ./ (dt * Aimpl[i,i]) ?
+        f!(R[i], U, p, τ)
     end
 
     # compute next step
     for i in 1:Nstages
-        u .+= (dt * tab.B[i]) .* cache.R[i]
+        u .+= (dt * B[i]) .* R[i]
     end
 end
 

src/solvers/lsrk.jl

@@ -40,22 +40,18 @@ function init_cache(prob::DiffEqBase.ODEProblem, alg::LowStorageRungeKutta2N; kw
     return LowStorageRungeKutta2NIncCache(tableau(alg, eltype(du)), du)
 end
 
-nstages(::LowStorageRungeKutta2NIncCache{N}) where {N} = N
+n_stages(::LowStorageRungeKutta2NIncCache{N}) where {N} = N
 
 function step_u!(int, cache::LowStorageRungeKutta2NIncCache)
-    tab = cache.tableau
+    (; C, A, B) = cache.tableau
     du = cache.du
+    (; u, p, t, dt) = int
 
-    u = int.u
-    p = int.p
-    t = int.t
-    dt = int.dt
-
-    for stage in 1:nstages(cache)
-        #  du .= f(u, p, t + tab.C[stage]*dt) .+ tab.A[stage] .* du
-        stage_time = t + tab.C[stage] * dt
-        int.sol.prob.f(du, u, p, stage_time, 1, tab.A[stage])
-        u .+= (dt * tab.B[stage]) .* du
+    for stage in 1:n_stages(cache)
+        #  du .= f(u, p, t + C[stage]*dt) .+ A[stage] .* du
+        stage_time = t + C[stage] * dt
+        int.sol.prob.f(du, u, p, stage_time, 1, A[stage])
+        u .+= (dt * B[stage]) .* du
     end
 end
 
@@ -65,22 +61,23 @@ function init_inner(prob, outercache::LowStorageRungeKutta2NIncCache, dt)
 end
 function update_inner!(innerinteg, outercache::LowStorageRungeKutta2NIncCache, f_slow, u, p, t, dt, stage)
 
+    (; C, A, B) = cache.tableau
     f_offset = innerinteg.sol.prob.f
     tab = outercache.tableau
-    N = nstages(outercache)
+    N = n_stages(outercache)
 
-    τ0 = t + tab.C[stage] * dt
-    τ1 = stage == N ? t + dt : t + tab.C[stage + 1] * dt
+    τ0 = t + C[stage] * dt
+    τ1 = stage == N ? t + dt : t + C[stage + 1] * dt
     f_offset.α = τ0
     innerinteg.t = zero(τ0)
     innerinteg.tstop = τ1 - τ0
 
-    #  du .= f(u, p, t + tab.C[stage]*dt) .+ tab.A[stage] .* du
-    f_slow(f_offset.x, u, p, τ0, 1, tab.A[stage])
+    #  du .= f(u, p, t + C[stage]*dt) .+ A[stage] .* du
+    f_slow(f_offset.x, u, p, τ0, 1, A[stage])
 
-    C0 = tab.C[stage]
-    C1 = stage == N ? one(tab.C[stage]) : tab.C[stage + 1]
-    f_offset.γ = tab.B[stage] / (C1 - C0)
+    C0 = C[stage]
+    C1 = stage == N ? one(C[stage]) : C[stage + 1]
+    f_offset.γ = B[stage] / (C1 - C0)
 end
 
 

src/solvers/mis.jl

@@ -49,7 +49,7 @@ struct MultirateInfinitesimalStepCache{Nstages, A, T <: MultirateInfinitesimalSt
     tableau::T
 end
 
-nstages(::MultirateInfinitesimalStepCache{Nstages}) where {Nstages} = Nstages
+n_stages(::MultirateInfinitesimalStepCache{Nstages}) where {Nstages} = Nstages
 
 function init_cache(
     prob::DiffEqBase.AbstractODEProblem{uType, tType, true},
@@ -76,15 +76,15 @@ end
 function update_inner!(innerinteg, outercache::MultirateInfinitesimalStepCache, f_slow, u, p, t, dt, i)
 
     f_offset = innerinteg.sol.prob.f
-    tab = outercache.tableau
-    N = nstages(outercache)
+    N = n_stages(outercache)
+    (; c, c̃, d) = outercache.tableau
 
     F = outercache.F
     ΔU = outercache.ΔU
 
     # F[i] = f_slow(U[i-1], p, t + c[i-1]*dt)
     u0 = i == 1 ? u : ΔU[i - 1]
-    t0 = i == 1 ? t : t + tab.c[i - 1] * dt
+    t0 = i == 1 ? t : t + c[i - 1] * dt
     f_slow(F[i], u0, p, t0)
 
     # the (i+1)th stage of the paper
@@ -109,16 +109,17 @@ function update_inner!(innerinteg, outercache::MultirateInfinitesimalStepCache,
 
     # KW2014 (9)
     # evaluate f_fast(z(τ), p, t + c̃[i]*dt + (c[i]-c̃[i])/d[i] * τ)
-    f_offset.α = t + tab.c̃[i] * dt
-    f_offset.β = (tab.c[i] - tab.c̃[i]) / tab.d[i]
+    f_offset.α = t + c̃[i] * dt
+    f_offset.β = (c[i] - c̃[i]) / d[i]
 
     innerinteg.t = zero(t)
-    innerinteg.tstop = tab.d[i] * dt
+    innerinteg.tstop = d[i] * dt
 end
 
 @kernel function mis_update!(u, ΔU, F, innerinteg_u, f_offset_x, tab, i, N, dt)
     e = @index(Global, Linear)
     @inbounds begin
+        (; α, β, d, γ) = tab
         if i > 1
             ΔU[i - 1][e] -= u[e]
         end
@@ -128,13 +129,13 @@ end
             innerinteg_u[e] = u[e]
         end
         for j in 1:(i - 1)
-            innerinteg_u[e] += tab.α[i, j] * ΔU[j][e]
+            innerinteg_u[e] += α[i, j] * ΔU[j][e]
         end
 
         # KW2014 (1b) / (9)
-        f_offset_x[e] = tab.β[i, i] / tab.d[i] .* F[i][e]
+        f_offset_x[e] = β[i, i] / d[i] .* F[i][e]
         for j in 1:(i - 1)
-            f_offset_x[e] += (tab.γ[i, j] / (tab.d[i] * dt)) * ΔU[j][e] + tab.β[i, j] / tab.d[i] * F[j][e]
+            f_offset_x[e] += (γ[i, j] / (d[i] * dt)) * ΔU[j][e] + β[i, j] / d[i] * F[j][e]
         end
     end
 end

src/solvers/multirate.jl

@@ -51,7 +51,7 @@ function step_u!(int, cache::MultirateCache)
     innerinteg = cache.innerinteg
     fast_dt = innerinteg.dt
 
-    N = nstages(outercache)
+    N = n_stages(outercache)
     for stage in 1:N
 
         update_inner!(innerinteg, outercache, int.sol.prob.f.f2, u, p, t, dt, stage)

src/solvers/rosenbrock.jl

@@ -33,39 +33,30 @@ end
 
 
 function step_u!(int, cache::RosenbrockCache{Nstages, RT}) where {Nstages, RT}
-    tab = cache.tableau
-
+    (; m, Γ, A, C) = cache.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
 
-    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]
+    γ = dt * Γ[1, 1]
     Wfact_t!(W, u, p, γ, t)
     for i in 1:Nstages
         U .= u
         for j in 1:(i - 1)
-            U .+= tab.A[i, j] .* k[j]
+            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)
         for j in 1:(i - 1)
-            fU .+= (tab.C[i, j] / dt) .* k[j]
+            fU .+= (C[i, j] / dt) .* k[j]
         end
         linsolve!(k[i], W, fU)
     end
     for i in 1:Nstages
-        u .+= tab.m[i] .* k[i]
+        u .+= m[i] .* k[i]
     end
 end