Generalized RK for nonlinear ODEs

src/ODEs/ODETools/RungeKutta.jl

@@ -32,28 +32,29 @@ function solve_step!(uf::AbstractVector,
   if cache === nothing
     ode_cache = allocate_cache(op)
     vi = similar(u0)
-    fi = [similar(u0)]
+    ui = [similar(u0)]
+    rhs = similar(u0)
     nls_stage_cache = nothing
     nls_update_cache = nothing
   else
-    ode_cache, vi, fi, nls_stage_cache, nls_update_cache = cache
+    ode_cache, vi, ui, rhs, nls_stage_cache, nls_update_cache = cache
   end
 
   # Create RKNL stage operator
-  nlop_stage = RungeKuttaStageNonlinearOperator(op,t0,dt,u0,ode_cache,vi,fi,0,a)
+  nlop_stage = RungeKuttaStageNonlinearOperator(op,t0,dt,u0,ode_cache,vi,ui,rhs,0,a)
 
   # Compute intermediate stages
   for i in 1:s
 
     # allocate space to store the RHS at i
-    if (length(fi) < i)
-      push!(fi,similar(u0))
+    if (length(ui) < i)
+      push!(ui,similar(u0))
     end
 
     # Update time
     ti = t0 + c[i]*dt
     ode_cache = update_cache!(ode_cache,op,ti)
-    update!(nlop_stage,ti,fi,i)
+    update!(nlop_stage,ti,ui,i)
 
     if(a[i,i]==0)
       # Skip stage solve if a_ii=0 => u_i=u_0, f_i = f_0
@@ -63,8 +64,14 @@ function solve_step!(uf::AbstractVector,
       nls_stage_cache = solve!(uf,solver.nls_stage,nlop_stage,nls_stage_cache)
     end
 
-    # Update RHS at stage i using solution at u_i
-    rhs!(nlop_stage, uf)
+    # Update stage unknown
+    @. nlop_stage.ui[i] = uf
+    println("uf RK = ",uf, "ti = ", ti)
+    println( "op.ui = ",nlop_stage.ui)
+    println( "ui = ",ui)
+
+    # # Update RHS at stage i using solution at u_i
+    # rhs!(nlop_stage, uf)
 
   end
 
@@ -76,15 +83,20 @@ function solve_step!(uf::AbstractVector,
 
     # Create RKNL final update operator
     ode_cache = update_cache!(ode_cache,op,tf)
-    nlop_update = RungeKuttaUpdateNonlinearOperator(op,tf,dt,u0,ode_cache,vi,fi,s,b)
+    nlop_update = RungeKuttaUpdateNonlinearOperator(op,tf,dt,u0,ode_cache,vi,ui,rhs,s,b)
+
+
+    println( "op_up.ui = ",nlop_update.ui)
+    println( "ui = ",ui)
 
     # solve at final update
     nls_update_cache = solve!(uf,solver.nls_update,nlop_update,nls_update_cache)
+    println("uf RK update = ",uf, "tf =",tf)
 
   end
 
   # Update final cache
-  cache = (ode_cache, vi, fi, nls_stage_cache, nls_update_cache)
+  cache = (ode_cache, vi, ui, rhs, nls_stage_cache, nls_update_cache)
 
   return (uf,tf,cache)
 
@@ -103,7 +115,8 @@ mutable struct RungeKuttaStageNonlinearOperator <: RungeKuttaNonlinearOperator
   u0::AbstractVector
   ode_cache
   vi::AbstractVector
-  fi::Vector{AbstractVector}
+  ui::Vector{AbstractVector}
+  rhs::AbstractVector
   i::Int
   a::Matrix
 end
@@ -119,7 +132,8 @@ mutable struct RungeKuttaUpdateNonlinearOperator <: RungeKuttaNonlinearOperator
   u0::AbstractVector
   ode_cache
   vi::AbstractVector
-  fi::Vector{AbstractVector}
+  ui::Vector{AbstractVector}
+  rhs::AbstractVector
   s::Int
   b::Vector{Number}
 end
@@ -129,23 +143,21 @@ end
 """
 Compute the residual of the Runge-Kutta nonlinear operator at stage i.
 ```math
-A(t,ui,∂ui/∂t) = ∂ui/∂t - a_ii * f(ui,ti) - ∑_{j<i} a_ij * f(uj,tj) = 0
+A(t,ui,∂ui/∂t) = ∂ui/∂t - f(∑_{j<=i} a_ij * uj, tj) = 0
 ```
 
 Uses the vector b as auxiliar variable to store the residual of the left-hand side of
 the i-th stage ODE operator, then adds the corresponding contribution from right-hand side
 at all earlier stages.
 ```math
 b = M(ui,ti)∂u/∂t
-b - ∑_{j<=i} a_ij * f(uj,tj) = 0
+b - f(∑_{j<=i} a_ij * uj,tj) = 0
 ```
 """
 function residual!(b::AbstractVector,op::RungeKuttaStageNonlinearOperator,x::AbstractVector)
   rhs!(op,x)
   lhs!(b,op,x)
-  for j in 1:op.i
-    @. b = b - op.a[op.i,j] * op.fi[j]
-  end
+  @. b = b - op.rhs
   b
 end
 
@@ -158,15 +170,14 @@ A(t,uf,∂uf/∂t) = ∂uf/∂t - ∑_{i<=s} b_i * f(ui,ti) = 0
 Uses the vector b as auxiliar variable to store the residual of the update ODE
 operator (e.g. identity or mass matrix), then adds the corresponding contribution from earlier stages.
 ```math
-b = [∂u/∂t]
-b - ∑_{i<=s} bi * f(ui,ti) = 0
+b = M(uf,tf)[∂u/∂t]
+b - f(∑_{i<=s} bi * ui,ti) = 0
 ```
 """
 function residual!(b::AbstractVector,op::RungeKuttaUpdateNonlinearOperator,x::AbstractVector)
   lhs!(b,op,x)
-  for i in 1:op.s
-    @. b = b - op.b[i] * op.fi[i]
-  end
+  rhs!(op,x)
+  @. b = b - op.rhs
   b
 end
 
@@ -213,12 +224,24 @@ function zero_initial_guess(op::RungeKuttaNonlinearOperator)
   x0
 end
 
-function rhs!(op::RungeKuttaNonlinearOperator, x::AbstractVector)
-  u = x
+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)
+  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)
-  f = op.fi
-  rhs!(f[op.i],op.odeop,op.ti,(u,v),op.ode_cache)
+  rhs!(op.rhs,op.odeop,op.ti,(u,v),op.ode_cache)
 end
 
 function lhs!(b::AbstractVector, op::RungeKuttaNonlinearOperator, x::AbstractVector)
@@ -228,9 +251,9 @@ function lhs!(b::AbstractVector, op::RungeKuttaNonlinearOperator, x::AbstractVec
   lhs!(b,op.odeop,op.ti,(u,v),op.ode_cache)
 end
 
-function update!(op::RungeKuttaNonlinearOperator,ti::Float64,fi::AbstractVector,i::Int)
+function update!(op::RungeKuttaNonlinearOperator,ti::Float64,ui::AbstractVector,i::Int)
   op.ti = ti
-  op.fi = fi
+  op.ui = ui
   op.i = i
 end