in-place fixed

Author: Shreyas-Ekanathan

lib/OrdinaryDiffEqFIRK/src/firk_addsteps.jl

@@ -108,7 +108,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA3Cache, repeat_step = false)
     @unpack t, dt, uprev, u, f, p, fsallast, fsalfirst = integrator
     @unpack T11, T12, T21, T22, TI11, TI12, TI21, TI22 = cache.tab
     @unpack c1, c2, α, β, e1, e2 = cache.tab
-    @unpack κ, cont1, cont2 = cache
+    @unpack κ = cache
     @unpack z1, z2, w1, w2,
     dw12, cubuff,
     k, k2, fw1, fw2,
@@ -127,14 +127,24 @@ function _ode_addsteps!(integrator, cache::RadauIIA3Cache, repeat_step = false)
         integrator.stats.nw += 1
     end
 
-    #better initial guess
-    uzero = zero(eltype(z1))
-    @. z1 = uzero
-    @. z2 = uzero
-    @. w1 = uzero
-    @. w2 = uzero
-    @. cache.cont1 = uzero
-    @. cache.cont2 = uzero
+    c1m1 = c1 - 1
+    if integrator.iter == 1 || integrator.u_modified || alg.extrapolant == :constant
+        cache.dtprev = one(cache.dtprev)
+        uzero = zero(eltype(u))
+        @.. broadcast=false z1=uzero
+        @.. broadcast=false z2=uzero
+        @.. broadcast=false w1=uzero
+        @.. broadcast=false w2=uzero
+        @.. broadcast=false integrator.k[3]=uzero
+        @.. broadcast=false integrator.k[4]=uzero
+    else
+        c2′ = dt / cache.dtprev
+        c1′ = c1 * c2′
+        @.. broadcast=false z1=c1′ * (k[3] + (c1′ - c1m1) * k[4])
+        @.. broadcast=false z2=c2′ * (k[3] + (c2′ - c1m1) * k[4])
+        @.. broadcast=false w1=TI11 * z1 + TI12 * z2
+        @.. broadcast=false w2=TI21 * z1 + TI22 * z2
+    end
 
     # Newton iteration
     local ndw
@@ -233,6 +243,14 @@ function _ode_addsteps!(integrator, cache::RadauIIA3Cache, repeat_step = false)
     @. u = uprev + z2
     step_limiter!(u, integrator, p, t + dt)
 
+    if integrator.EEst <= oneunit(integrator.EEst)
+        cache.dtprev = dt
+        if alg.extrapolant != :constant
+            integrator.k[3] = (z1 - z2) / c1m1
+            integrator.k[4] = integrator.k[3] - (z1 / c1)
+        end
+    end
+
     f(fsallast, u, p, t + dt)
     OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1)
     return
@@ -384,7 +402,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA5Cache, repeat_step = false)
 @unpack t, dt, uprev, u, f, p, fsallast, fsalfirst, k = integrator
 @unpack T11, T12, T13, T21, T22, T23, T31, TI11, TI12, TI13, TI21, TI22, TI23, TI31, TI32, TI33 = cache.tab
 @unpack c1, c2, γ, α, β, e1, e2, e3 = cache.tab
-@unpack κ, cont1, cont2, cont3 = cache
+@unpack κ = cache
 @unpack z1, z2, z3, w1, w2, w3,
 dw1, ubuff, dw23, cubuff,
 k, k2, k3, fw1, fw2, fw3,
@@ -396,8 +414,8 @@ alg = unwrap_alg(integrator, true)
 mass_matrix = integrator.f.mass_matrix
 
 # precalculations
-c1m1 = (c1 - 1)*dt
-c2m1 = (c2 - 1)*dt
+c1m1 = c1 - 1
+c2m1 = c2 - 1
 c1mc2 = c1 - c2
 γdt, αdt, βdt = γ / dt, α / dt, β / dt
 if (new_W = do_newW(integrator, alg, new_jac, cache.W_γdt))
@@ -422,7 +440,7 @@ if integrator.iter == 1 || integrator.u_modified || alg.extrapolant == :constant
     @.. broadcast=false integrator.k[4] = uzero
     @.. broadcast=false integrator.k[5] = uzero
 else
-    c3′ = dt
+    c3′ = dt / cache.dtprev
     c1′ = c1 * c3′
     c2′ = c2 * c3′
     @.. broadcast=false z1=c1′ * (k[3] + (c1′ - c2m1) * (k[4] + (c1′ - c1m1) * k[5]))
@@ -565,8 +583,8 @@ step_limiter!(u, integrator, p, t + dt)
 if integrator.EEst <= oneunit(integrator.EEst)
     cache.dtprev = dt
     if alg.extrapolant != :constant
-        integrator.k[3] = (z2 - z3) / (dt * c2m1)
-        @.. tmp=(z1 - z2) / (dt * c1mc2)
+        integrator.k[3] = (z2 - z3) / c2m1
+        @.. tmp=(z1 - z2) / c1mc2
         integrator.k[4] = (tmp - integrator.k[3]) / c1m1
         integrator.k[5] = integrator.k[4] - (tmp - z1 / c1) / c2
     end
@@ -807,7 +825,7 @@ function _ode_addsteps!(integrator, cache::RadauIIA9Cache, repeat_step = false)
 @unpack T11, T12, T13, T14, T15, T21, T22, T23, T24, T25, T31, T32, T33, T34, T35, T41, T42, T43, T44, T45, T51 = cache.tab #= T52 = 1, T53 = 0, T54 = 1, T55 = 0=#
 @unpack TI11, TI12, TI13, TI14, TI15, TI21, TI22, TI23, TI24, TI25, TI31, TI32, TI33, TI34, TI35, TI41, TI42, TI43, TI44, TI45, TI51, TI52, TI53, TI54, TI55 = cache.tab
 @unpack c1, c2, c3, c4, γ, α1, β1, α2, β2, e1, e2, e3, e4, e5 = cache.tab
-@unpack κ= cache
+@unpack κ = cache
 @unpack z1, z2, z3, z4, z5, w1, w2, w3, w4, w5 = cache
 @unpack dw1, ubuff, dw23, dw45, cubuff1, cubuff2 = cache
 @unpack k, k2, k3, k4, k5, fw1, fw2, fw3, fw4, fw5 = cache

lib/OrdinaryDiffEqFIRK/src/firk_caches.jl

@@ -41,8 +41,6 @@ mutable struct RadauIIA3Cache{uType, cuType, uNoUnitsType, rateType, JType, W1Ty
     w2::uType
     dw12::cuType
     cubuff::cuType
-    cont1::uType
-    cont2::uType
     du1::rateType
     fsalfirst::rateType
     k::rateType
@@ -86,8 +84,6 @@ function alg_cache(alg::RadauIIA3, u, rate_prototype, ::Type{uEltypeNoUnits},
     recursivefill!(dw12, false)
     cubuff = similar(u, Complex{eltype(u)})
     recursivefill!(cubuff, false)
-    cont1 = zero(u)
-    cont2 = zero(u)
 
     fsalfirst = zero(rate_prototype)
     k = zero(rate_prototype)
@@ -118,7 +114,7 @@ function alg_cache(alg::RadauIIA3, u, rate_prototype, ::Type{uEltypeNoUnits},
 
     RadauIIA3Cache(u, uprev,
         z1, z2, w1, w2,
-        dw12, cubuff, cont1, cont2,
+        dw12, cubuff,
         du1, fsalfirst, k, k2, fw1, fw2,
         J, W1,
         uf, tab, κ, one(uToltype), 10000,
@@ -462,7 +458,7 @@ function alg_cache(alg::RadauIIA9, u, rate_prototype, ::Type{uEltypeNoUnits},
 
     RadauIIA9Cache(u, uprev,
         z1, z2, z3, z4, z5, w1, w2, w3, w4, w5,
-        dw1, ubuff, dw23, dw45, cubuff1, cubuff2, 
+        dw1, ubuff, dw23, dw45, cubuff1, cubuff2,
         du1, fsalfirst, k, k2, k3, k4, k5, fw1, fw2, fw3, fw4, fw5,
         J, W1, W2, W3,
         uf, tab, κ, one(uToltype), 10000,

lib/OrdinaryDiffEqFIRK/src/firk_interpolants.jl

@@ -10,6 +10,15 @@ FIRK_WITH_INTERPOLATIONS = Union{RadauIIA3ConstantCache, RadauIIA3Cache, RadauII
     @..  y₀ + Θ * (cont1 + (Θ - c1m1) * cont2)
 end
 
+@muladd function _ode_interpolant!(
+    out, Θ, dt, y₀, y₁, k, cache::Union{RadauIIA3ConstantCache, RadauIIA3Cache},
+    idxs::Nothing, T::Type{Val{0}}, differential_vars)
+    @unpack c1 = cache.tab
+    c1m1 = c1 - 1
+    Θdt = 1 - Θ
+    @.. out = y₁ - Θdt * (k[3] - (Θdt + c1m1) * k[4])
+end
+
 @muladd function _ode_interpolant(
     Θ, dt, y₀, y₁, k, cache::Union{RadauIIA5ConstantCache, RadauIIA5Cache},
     idxs::Nothing, T::Type{Val{0}}, differential_vars)
@@ -26,9 +35,9 @@ end
     idxs::Nothing, T::Type{Val{0}}, differential_vars)
     @unpack c1, c2 = cache.tab
     @unpack dtprev = cache
-    c1m1 = (c1 - 1) * dt
-    c2m1 = (c2 - 1) * dt
-    Θdt = (1 - Θ)*dt
+    c1m1 = c1 - 1
+    c2m1 = c2 - 1
+    Θdt = 1 - Θ
     @.. out = y₁ - Θdt * (k[3] - (Θdt + c2m1) * (k[4] - (Θdt + c1m1) * k[5]))
 end
 
@@ -41,7 +50,6 @@ end
     c3m1 = c3 - 1
     c4m1 = c4 - 1
     Θdt = 1 - Θ
-    print("yes")
     @.. y₁ - Θdt * (k[3] - (Θdt + c4m1) * (k[4] - (Θdt + c3m1) * (k[5] - (Θdt + c2m1) * (k[6] - (Θdt + c1m1) * k[7]))))
 end
 
@@ -63,7 +71,6 @@ end
     @unpack num_stages, index = cache
     @unpack c = cache.tabs[index]
     Θdt = 1 - Θ
-    @show k
     tmp = k[num_stages + 1] - k[num_stages + 2] * (Θdt + c[1] - 1) 
     j = num_stages - 2
     while j > 0