Merge pull request #274 from SciML/u/fixes

Fix Rodas4 and other fixes
SciML · Apr 26, 2023 · cd4b302 · cd4b302
2 parents 6abad67 + 7126d70
commit cd4b302
Show file tree

Hide file tree

Showing 7 changed files with 236 additions and 203 deletions.
diff --git a/src/integrators/integrator_utils.jl b/src/integrators/integrator_utils.jl
@@ -10,18 +10,6 @@ function build_adaptive_controller_cache(alg::A, ::Type{T}) where {A, T}
     return beta1, beta2, qmax, qmin, gamma, qoldinit, qold
 end
 
-function build_adaptive_tsit5_controller_cache(::Type{T}) where {T}
-    beta1 = T(7 / 50)
-    beta2 = T(2 / 25)
-    qmax = T(10.0)
-    qmin = T(1 / 5)
-    gamma = T(9 / 10)
-    qoldinit = T(1e-4)
-    qold = qoldinit
-
-    return beta1, beta2, qmax, qmin, gamma, qoldinit, qold
-end
-
 @inline function savevalues!(integrator::DiffEqBase.AbstractODEIntegrator{AlgType, IIP, S, T
                                                                           }, ts,
                              us,

diff --git a/src/integrators/types.jl b/src/integrators/types.jl
diff --git a/src/perform_step/gpu_rodas4_perform_step.jl b/src/perform_step/gpu_rodas4_perform_step.jl
@@ -7,7 +7,7 @@
     integ.uprev = integ.u
     uprev = integ.u
     @unpack a21, a31, a32, a41, a42, a43, a51, a52, a53, a54, C21, C31, C32, C41, C42, C43,
-    C51, C52, C53, C54, C61, C62, C63, C64, C65, gamma, c2, c3, c4, d1, d2, d3,
+    C51, C52, C53, C54, C61, C62, C63, C64, C65, γ, c2, c3, c4, d1, d2, d3,
     d4 = integ.tab
 
     integ.tprev = t
@@ -57,7 +57,7 @@
     dtd2 = dt * d2
     dtd3 = dt * d3
     dtd4 = dt * d4
-    dtgamma = dt * gamma
+    dtgamma = dt * γ
 
     # Starting
     W = J - I * inv(dtgamma)
@@ -188,9 +188,15 @@ end
     reltol = integ.reltol
 
     @unpack a21, a31, a32, a41, a42, a43, a51, a52, a53, a54, C21, C31, C32, C41, C42, C43,
-    C51, C52, C53, C54, C61, C62, C63, C64, C65, gamma, c2, c3, c4, d1, d2, d3,
+    C51, C52, C53, C54, C61, C62, C63, C64, C65, γ, c2, c3, c4, d1, d2, d3,
     d4 = integ.tab
 
+    @unpack h21, h22, h23, h24, h25, h31, h32, h33, h34, h35 = integ.tab
+
+    # Jacobian
+    J = f.jac(uprev, p, t)
+    dT = f.tgrad(uprev, p, t)
+
     if integ.u_modified
         k1 = f(uprev, p, t)
         integ.u_modified = false
@@ -203,10 +209,6 @@ end
     while EEst > convert(T, 1.0)
         dt < convert(T, 1.0f-14) && error("dt<dtmin")
 
-        # Jacobian
-        J = f.jac(uprev, p, t)
-        dT = f.tgrad(uprev, p, t)
-
         # Precalculations
         dtC21 = C21 / dt
         dtC31 = C31 / dt
@@ -228,7 +230,7 @@ end
         dtd2 = dt * d2
         dtd3 = dt * d3
         dtd4 = dt * d4
-        dtgamma = dt * gamma
+        dtgamma = dt * γ
 
         # Starting
         W = J - I * inv(dtgamma)
@@ -288,7 +290,6 @@ end
             dtnew = min(abs(dtnew), abs(tf - t - dt))
 
             @inbounds begin # Necessary for interpolation
-                @unpack h21, h22, h23, h24, h25, h31, h32, h33, h34, h35 = integ.tab
                 integ.k1 = h21 * k1 + h22 * k2 + h23 * k3 + h24 * k4 + h25 * k5
                 integ.k2 = h31 * k1 + h32 * k2 + h33 * k3 + h34 * k4 + h35 * k5
             end

diff --git a/src/perform_step/gpu_tsit5_perform_step.jl b/src/perform_step/gpu_tsit5_perform_step.jl
@@ -82,7 +82,7 @@ end
 
     saveat = _saveat === nothing ? saveat : _saveat
 
-    integ = gputsit5_init(prob.f, false, prob.u0, prob.tspan[1], dt, prob.p, tstops,
+    integ = gputsit5_init(alg, prob.f, false, prob.u0, prob.tspan[1], dt, prob.p, tstops,
                           callback, save_everystep, saveat)
 
     u0 = prob.u0
@@ -121,7 +121,8 @@ end
 #############################Adaptive Version#####################################
 
 @inline function step!(integ::GPUAT5I{false, S, T}, ts, us) where {S, T}
-    beta1, beta2, qmax, qmin, gamma, qoldinit, _ = build_adaptive_tsit5_controller_cache(eltype(integ.u))
+    beta1, beta2, qmax, qmin, gamma, qoldinit, _ = build_adaptive_controller_cache(integ.alg,
+                                                                                   eltype(integ.u))
     c1, c2, c3, c4, c5, c6 = integ.cs
     dt = integ.dtnew
     t = integ.t
@@ -247,7 +248,8 @@ end
     t = tspan[1]
     tf = prob.tspan[2]
 
-    integ = gpuatsit5_init(prob.f, false, prob.u0, prob.tspan[1], prob.tspan[2], dt, prob.p,
+    integ = gpuatsit5_init(alg, prob.f, false, prob.u0, prob.tspan[1], prob.tspan[2], dt,
+                           prob.p,
                            abstol, reltol, DiffEqBase.ODE_DEFAULT_NORM, tstops, callback,
                            saveat)
 

diff --git a/src/perform_step/gpu_vern7_perform_step.jl b/src/perform_step/gpu_vern7_perform_step.jl
@@ -90,7 +90,7 @@ end
 
     saveat = _saveat === nothing ? saveat : _saveat
 
-    integ = gpuvern7_init(prob.f, false, prob.u0, prob.tspan[1], dt, prob.p, tstops,
+    integ = gpuvern7_init(alg, prob.f, false, prob.u0, prob.tspan[1], dt, prob.p, tstops,
                           callback, save_everystep, saveat)
 
     u0 = prob.u0
@@ -128,7 +128,9 @@ end
 #############################Adaptive Version#####################################
 
 @inline function step!(integ::GPUAV7I{false, S, T}, ts, us) where {S, T}
-    beta1, beta2, qmax, qmin, gamma, qoldinit, _ = build_adaptive_tsit5_controller_cache(eltype(integ.u))
+    beta1, beta2, qmax, qmin, gamma, qoldinit, _ = build_adaptive_controller_cache(integ.alg,
+                                                                                   eltype(integ.u))
+
     dt = integ.dtnew
     t = integ.t
     p = integ.p
@@ -272,7 +274,8 @@ end
     t = tspan[1]
     tf = prob.tspan[2]
 
-    integ = gpuavern7_init(prob.f, false, prob.u0, prob.tspan[1], prob.tspan[2], dt, prob.p,
+    integ = gpuavern7_init(alg, prob.f, false, prob.u0, prob.tspan[1], prob.tspan[2], dt,
+                           prob.p,
                            abstol, reltol, DiffEqBase.ODE_DEFAULT_NORM, tstops, callback,
                            saveat)
 

diff --git a/src/perform_step/gpu_vern9_perform_step.jl b/src/perform_step/gpu_vern9_perform_step.jl
@@ -117,7 +117,7 @@ end
 
     saveat = _saveat === nothing ? saveat : _saveat
 
-    integ = gpuvern9_init(prob.f, false, prob.u0, prob.tspan[1], dt, prob.p, tstops,
+    integ = gpuvern9_init(alg, prob.f, false, prob.u0, prob.tspan[1], dt, prob.p, tstops,
                           callback, save_everystep, saveat)
 
     u0 = prob.u0
@@ -156,7 +156,9 @@ end
 #############################Adaptive Version#####################################
 
 @inline function step!(integ::GPUAV9I{false, S, T}, ts, us) where {S, T}
-    beta1, beta2, qmax, qmin, gamma, qoldinit, _ = build_adaptive_tsit5_controller_cache(eltype(integ.u))
+    beta1, beta2, qmax, qmin, gamma, qoldinit, _ = build_adaptive_controller_cache(integ.alg,
+                                                                                   eltype(integ.u))
+
     dt = integ.dtnew
     t = integ.t
     p = integ.p
@@ -331,7 +333,8 @@ end
     t = tspan[1]
     tf = prob.tspan[2]
 
-    integ = gpuavern9_init(prob.f, false, prob.u0, prob.tspan[1], prob.tspan[2], dt, prob.p,
+    integ = gpuavern9_init(alg, prob.f, false, prob.u0, prob.tspan[1], prob.tspan[2], dt,
+                           prob.p,
                            abstol, reltol, DiffEqBase.ODE_DEFAULT_NORM, tstops, callback,
                            saveat)
 

diff --git a/src/tableaus/rodas_tableaus.jl b/src/tableaus/rodas_tableaus.jl
@@ -24,7 +24,7 @@ struct Rodas4Tableau{T, T2}
     C63::T
     C64::T
     C65::T
-    gamma::T
+    γ::T
     c2::T2
     c3::T2
     c4::T2
@@ -45,7 +45,7 @@ struct Rodas4Tableau{T, T2}
 end
 
 function Rodas4Tableau(T::Type{T1}, T2::Type{T1}) where {T1}
-    gamma = convert(T, 0.25)
+    γ = convert(T, 0.25)
     #BET2P=0.0317D0
     #BET3P=0.0635D0
     #BET4P=0.3438D0
@@ -97,6 +97,6 @@ function Rodas4Tableau(T::Type{T1}, T2::Type{T1}) where {T1}
 
     Rodas4Tableau(a21, a31, a32, a41, a42, a43, a51, a52, a53, a54,
                   C21, C31, C32, C41, C42, C43, C51, C52, C53, C54, C61, C62, C63, C64, C65,
-                  gamma, c2, c3, c4, d1, d2, d3, d4,
+                  γ, c2, c3, c4, d1, d2, d3, d4,
                   h21, h22, h23, h24, h25, h31, h32, h33, h34, h35)
 end