make oop mutation free

Author: oscardssmith

lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_caches.jl

@@ -697,15 +697,14 @@ end
 
 ### Rodas4 methods
 
-struct Rodas4ConstantCache{TF, UF, Tab, JType, WType, F, AD, rateType} <: RosenbrockConstantCache
+struct Rodas4ConstantCache{TF, UF, Tab, JType, WType, F, AD} <: RosenbrockConstantCache
     tf::TF
     uf::UF
     tab::Tab
     J::JType
     W::WType
     linsolve::F
     autodiff::AD
-    ks::Vector{rateType}
 end
 
 tabtype(::Rodas4) = Rodas4Tableau
@@ -722,11 +721,10 @@ function alg_cache(alg::Union{Rodas4, Rodas42, Rodas4P, Rodas4P2}, u, rate_proto
     J, W = build_J_W(alg, u, uprev, p, t, dt, f, uEltypeNoUnits, Val(false))
     linprob = nothing #LinearProblem(W,copy(u); u0=copy(u))
     linsolve = nothing #init(linprob,alg.linsolve,alias_A=true,alias_b=true)
-    ks = Vector{typeof(rate_prototype)}(undef, 6)
     Rodas4ConstantCache(tf, uf,
         tabtype(alg)(constvalue(uBottomEltypeNoUnits),
             constvalue(tTypeNoUnits)), J, W, linsolve,
-        alg_autodiff(alg), ks)
+        alg_autodiff(alg))
 end
 
 function alg_cache(alg::Union{Rodas4, Rodas42, Rodas4P, Rodas4P2}, u, rate_prototype, ::Type{uEltypeNoUnits},

lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_perform_step.jl

@@ -1208,7 +1208,7 @@ end
 
 @muladd function perform_step!(integrator, cache::Rodas4ConstantCache, repeat_step = false)
     (;t, dt, uprev, u, f, p) = integrator
-    (;tf, uf, ks) = cache
+    (;tf, uf) = cache
     (;A, C, gamma, c, d, H) = cache.tab
 
     # Precalculations
@@ -1228,11 +1228,15 @@ end
         return nothing
     end
 
-    # Initialize k arrays
+    # Initialize ks
     num_stages = size(A,1)
-
+    du = f(u, p, t)
+    linsolve_tmp = @.. du + dtd[1] * dT
+    k1 = _reshape(W \ -_vec(linsolve_tmp), axes(uprev))
+    # constant number for type stability make sure this is greater than num_stages
+    ks = ntuple(Returns(k1), 10)
     # Loop for stages
-    for stage in 1:num_stages
+    for stage in 2:num_stages
         u = uprev
         for i in 1:stage-1
             u = @.. u + A[stage, i] * ks[i]
@@ -1255,7 +1259,7 @@ end
         end
         linsolve_tmp = @.. du + dtd[stage] * dT + linsolve_tmp1
 
-        ks[stage] = _reshape(W \ -_vec(linsolve_tmp), axes(uprev))
+        ks = Base.setindex(ks, _reshape(W \ -_vec(linsolve_tmp), axes(uprev)), stage)
         integrator.stats.nsolve += 1
     end
     #@show ks
@@ -1271,7 +1275,7 @@ end
         for j in eachindex(integrator.k)
             integrator.k[j] = zero(integrator.k[1])
         end
-        for i in eachindex(ks)
+        for i in 1:num_stages
             for j in eachindex(integrator.k)
                 integrator.k[j] = @.. integrator.k[j] + H[j, i] * ks[i]
             end

lib/OrdinaryDiffEqRosenbrock/src/stiff_addsteps.jl

@@ -291,7 +291,7 @@ function _ode_addsteps!(k, t, uprev, u, dt, f, p, cache::Rodas4ConstantCache,
         always_calc_begin = false, allow_calc_end = true,
         force_calc_end = false)
     if length(k) < 2 || always_calc_begin
-        (;tf, uf, ks) = cache
+        (;tf, uf) = cache
         (;A, C, gamma, c, d, H) = cache.tab
 
         # Precalculations
@@ -317,10 +317,16 @@ function _ode_addsteps!(k, t, uprev, u, dt, f, p, cache::Rodas4ConstantCache,
             J = ForwardDiff.derivative(uf, uprev)
             W = 1 / dtgamma - J
         end
-
-        num_stages = size(A, 1)
+        
+        
+        num_stages = size(A,1)
+        du = f(u, p, t)
+        linsolve_tmp = @.. du + dtd[1] * dT
+        k1 = _reshape(W \ -_vec(linsolve_tmp), axes(uprev))
+        # constant number for type stability make sure this is greater than num_stages
+        ks = ntuple(Returns(k1), 10)
         # Last stage doesn't affect ks
-        for stage in 1:num_stages-1
+        for stage in 2:num_stages-1
             u = uprev
             for i in 1:stage-1
                 u = @.. u + A[stage, i] * ks[i]
@@ -341,14 +347,14 @@ function _ode_addsteps!(k, t, uprev, u, dt, f, p, cache::Rodas4ConstantCache,
                 end
             end
 
-            ks[stage] = _reshape(W \ _vec(linsolve_tmp), axes(uprev))
+            ks = Base.setindex(ks, _reshape(W \ _vec(linsolve_tmp), axes(uprev)), stage)
         end
 
         k1 = zero(ks[1])
         k2 = zero(ks[1])
         H = cache.tab.H
         # Last stage doesn't affect ks
-        for i in 1:length(ks)-1
+        for i in 1:num_stages-1
             k1 = @.. k1 + H[1, i] * ks[i]
             k2 = @.. k2 + H[2, i] * ks[i]
         end