From 324f163c15e1b2cfdb29052b946299e126390d37 Mon Sep 17 00:00:00 2001
From: Oriol Colomes <oriol.colomes@gmail.com>
Date: Fri, 20 Oct 2023 22:07:49 +0200
Subject: [PATCH] Fixes in RK

---
 src/ODEs/ODETools/RungeKutta.jl             | 25 +++++++++------------
 test/ODEsTests/ODEsTests/ODESolversTests.jl |  8 +++++++
 2 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/ODEs/ODETools/RungeKutta.jl b/src/ODEs/ODETools/RungeKutta.jl
index 3aaa0578d..1b8993d2b 100644
--- a/src/ODEs/ODETools/RungeKutta.jl
+++ b/src/ODEs/ODETools/RungeKutta.jl
@@ -32,7 +32,9 @@ function solve_step!(uf::AbstractVector,
   if cache === nothing
     ode_cache = allocate_cache(op)
     vi = similar(u0)
-    ui = [similar(u0)]
+    ui = Vector{typeof(u0)}()
+    sizehint!(ui,s)
+    [push!(ui,similar(u0)) for i in 1:s]
     rhs = similar(u0)
     nls_stage_cache = nothing
     nls_update_cache = nothing
@@ -46,11 +48,6 @@ function solve_step!(uf::AbstractVector,
   # Compute intermediate stages
   for i in 1:s
 
-    # allocate space to store the RHS at i
-    if (length(ui) < i)
-      push!(ui,similar(u0))
-    end
-
     # Update time
     ti = t0 + c[i]*dt
     ode_cache = update_cache!(ode_cache,op,ti)
@@ -214,22 +211,22 @@ function zero_initial_guess(op::RungeKuttaNonlinearOperator)
 end
 
 function rhs!(op::RungeKuttaStageNonlinearOperator, x::AbstractVector)
-  u = zeros(eltype(x),length(x))
-  for j in 1:op.i
-    @. u = u + op.a[op.i,j] * op.ui[j]
-  end
   v = op.vi
   @. v = (x-op.u0)/(op.dt)
+  u = op.a[op.i,op.i] * op.ui[op.i]
+  for j in 1:op.i-1
+      @. u += op.a[op.i,j] * op.ui[j]
+  end
   rhs!(op.rhs,op.odeop,op.ti,(u,v),op.ode_cache)
 end
 
 function rhs!(op::RungeKuttaUpdateNonlinearOperator, x::AbstractVector)
-  u = zeros(eltype(x),length(x))
-  for i in 1:op.s
-    @. u = u + op.b[i] * op.ui[i]
-  end
   v = op.vi
   @. v = (x-op.u0)/(op.dt)
+  u = op.b[op.s] * op.ui[op.s]
+  for i in 1:op.s-1
+    @. u = u + op.b[i] * op.ui[i]
+  end
   rhs!(op.rhs,op.odeop,op.ti,(u,v),op.ode_cache)
 end
 
diff --git a/test/ODEsTests/ODEsTests/ODESolversTests.jl b/test/ODEsTests/ODEsTests/ODESolversTests.jl
index 3794a1168..11a8010d0 100644
--- a/test/ODEsTests/ODEsTests/ODESolversTests.jl
+++ b/test/ODEsTests/ODEsTests/ODESolversTests.jl
@@ -132,6 +132,7 @@ ufθ, tf, cache = solve_step!(ufθ,odesolθ,op,u0,t0,nothing)
 
 # RK tests
 # RK: BE equivalent
+# u1-u0 = dt*u1 => u1 = u0/(1-dt) = 2.2222222222222223
 odesol = RungeKutta(ls,ls,dt,:BE_1_0_1)
 cache = nothing
 uf, tf, cache = solve_step!(uf,odesol,op,u0,t0,cache)
@@ -148,6 +149,13 @@ uf, tf, cache = solve_step!(uf,odesol,op,u0,t0,cache)
 @test test_ode_solver(odesol,op,u0,t0,tf)
 
 # RK: SDIRK 2nd order
+# u1-u0 = 0.25*dt*u1
+# u1(1-0.25*dt) = u0 => u1 = u0/(1-0.25*dt) = 2.0512820512820515
+# u2-u0 = 0.25*dt*u2 + 0.5*dt*u0/(1-0.25*dt)
+# u2(1-0.25*dt) = u0 + 0.5*dt*u0/(1-0.25*dt) = u0*(1+0.5*0.1/(1-0.25*0.1))
+# u2 = u0*(1+0.5*0.1/(1-0.25*0.1))/(1-0.25*dt) = 2.156476002629849
+# uf-u0 = 0.5*dt*u1 + 0.5*dt*u2 = 0.5*dt*u0*( 1/(1-0.25*dt) + (1+0.5*dt/(1-0.25*dt)) )
+# uf = u0*(1 + 0.5*dt/(1-0.25*dt) + 0.5*dt*(1+0.5*0.1/(1-0.25*0.1))/(1-0.25*dt))
 odesol = RungeKutta(ls,ls,dt,:SDIRK_2_0_2)
 cache = nothing
 uf, tf, cache = solve_step!(uf,odesol,op,u0,t0,cache)