Nonlinear Runge-Kutta

src/ODEs/ODETools/IMEXRungeKutta.jl

@@ -154,6 +154,9 @@ mutable struct IMEXRungeKuttaUpdateNonlinearOperator <: RungeKuttaNonlinearOpera
   bₑ::Vector{Float64}
 end
 
+IMEXRungeKuttaNonlinearOperator = Union{IMEXRungeKuttaStageNonlinearOperator,
+  IMEXRungeKuttaUpdateNonlinearOperator}
+
 """
 residual!(b,op::IMEXRungeKuttaStageNonlinearOperator,x)
 
@@ -251,14 +254,29 @@ function jacobian!(J::AbstractMatrix,
   jacobian!(J,op.odeop,op.ti,(uf,vf),2,1.0/(op.dt),op.ode_cache)
 end
 
-function explicit_rhs!(op::RungeKuttaNonlinearOperator, x::AbstractVector)
+function rhs!(op::IMEXRungeKuttaNonlinearOperator, x::AbstractVector)
+  u = x
+  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)
+end
+
+function explicit_rhs!(op::IMEXRungeKuttaNonlinearOperator, x::AbstractVector)
   u = x
   v = op.vi
   @. v = (x-op.u0)/(op.dt)
   g = op.gi
   explicit_rhs!(g[op.i],op.odeop,op.ti,(u,v),op.ode_cache)
 end
 
+function lhs!(b::AbstractVector, op::IMEXRungeKuttaNonlinearOperator, x::AbstractVector)
+  u = x
+  v = op.vi
+  @. v = (x-op.u0)/(op.dt)
+  lhs!(b,op.odeop,op.ti,(u,v),op.ode_cache)
+end
+
 function update!(op::RungeKuttaNonlinearOperator,ti::Float64,fi::AbstractVector,gi::AbstractVector,i::Int)
   op.ti = ti
   op.fi = fi

src/ODEs/ODETools/ODETools.jl

@@ -44,6 +44,7 @@ using Gridap.Algebra: jacobian
 using Gridap.Algebra: symbolic_setup
 using Gridap.Algebra: numerical_setup
 using Gridap.Algebra: numerical_setup!
+using Gridap.Algebra: LinearSolverCache
 import Gridap.Algebra: residual!
 import Gridap.Algebra: jacobian!
 import Gridap.Algebra: allocate_residual

src/ODEs/ODETools/RungeKutta.jl

@@ -31,60 +31,52 @@ function solve_step!(uf::AbstractVector,
   # Create cache if not there
   if cache === nothing
     ode_cache = allocate_cache(op)
-    vi = similar(u0)
-    fi = [similar(u0)]
+    ui = similar(u0)
+    ki = Vector{typeof(u0)}()
+    sizehint!(ki,s)
+    [push!(ki,similar(u0)) for i in 1:s]
+    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, ui, ki, rhs, nls_stage_cache, nls_update_cache = cache
+  end
+
+  # Initialize states to zero
+  for i in 1:s
+    @. ki[i] *= 0.0
   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,ui,ki,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))
-    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,i)
 
-    if(a[i,i]==0)
-      # Skip stage solve if a_ii=0 => u_i=u_0, f_i = f_0
-      @. uf = u0
-    else
-      # solve at stage i
-      nls_stage_cache = solve!(uf,solver.nls_stage,nlop_stage,nls_stage_cache)
-    end
+    # solve at stage i
+    nls_stage_cache = solve!(uf,solver.nls_stage,nlop_stage,nls_stage_cache)
 
-    # Update RHS at stage i using solution at u_i
-    rhs!(nlop_stage, uf)
+    # Update stage unknown
+    @. nlop_stage.ki[i] = uf
 
   end
 
   # Update final time
   tf = t0+dt
 
-  # Skip final update if not necessary
-  if !(c[s]==1.0 && a[s,:] == b)
-
-    # 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)
-
-    # solve at final update
-    nls_update_cache = solve!(uf,solver.nls_update,nlop_update,nls_update_cache)
-
+  # Update final solution
+  @. uf = u0
+  for i in 1:s
+    @. uf = uf + dt * b[i] * nlop_stage.ki[i]
   end
 
   # Update final cache
-  cache = (ode_cache, vi, fi, nls_stage_cache, nls_update_cache)
+  cache = (ode_cache, ui, ki, rhs, nls_stage_cache, nls_update_cache)
 
   return (uf,tf,cache)
 
@@ -94,109 +86,52 @@ abstract type RungeKuttaNonlinearOperator <: NonlinearOperator end
 
 """
 Nonlinear operator that represents the Runge-Kutta nonlinear operator at a
-given time step and stage, i.e., A(t,u_i,(u_i-u_n)/dt)
+given time step and stage, i.e., A(tᵢ,uᵢ,kᵢ)
 """
 mutable struct RungeKuttaStageNonlinearOperator <: RungeKuttaNonlinearOperator
   odeop::ODEOperator
   ti::Float64
   dt::Float64
   u0::AbstractVector
   ode_cache
-  vi::AbstractVector
-  fi::Vector{AbstractVector}
+  ui::AbstractVector
+  ki::Vector{AbstractVector}
+  rhs::AbstractVector
   i::Int
   a::Matrix
 end
 
-"""
-Nonlinear operator that represents the Runge-Kutta nonlinear operator at the
-final updated of a given time step, i.e., A(t,u_f,(u_f-u_n)/dt)
-"""
-mutable struct RungeKuttaUpdateNonlinearOperator <: RungeKuttaNonlinearOperator
-  odeop::ODEOperator
-  ti::Float64
-  dt::Float64
-  u0::AbstractVector
-  ode_cache
-  vi::AbstractVector
-  fi::Vector{AbstractVector}
-  s::Int
-  b::Vector{Number}
-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,ki) = M(ti) ki - f(u₀ + ∑_{j<=i} Δt * a_ij * kj, 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 = M(ti) Ki
+b - f(u₀ + ∑_{j<=i} Δt * a_ij * kj, 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
-end
-
-"""
-Compute the residual of the final update of the Runge-Kutta nonlinear operator.
-```math
-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
-```
-"""
-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
+  @. b = b - op.rhs
   b
 end
 
 function jacobian!(A::AbstractMatrix,op::RungeKuttaStageNonlinearOperator,x::AbstractVector)
-  # @assert (abs(op.a[op.i,op.i]) > 0.0)
-  ui = x
-  vi = op.vi
-  @. vi = (x-op.u0)/(op.dt)
-  z = zero(eltype(A))
-  fillstored!(A,z)
-  jacobians!(A,op.odeop,op.ti,(ui,vi),(op.a[op.i,op.i],1.0/op.dt),op.ode_cache)
-end
-
-function jacobian!(A::AbstractMatrix,op::RungeKuttaUpdateNonlinearOperator,x::AbstractVector)
-  uf = x
-  vf = op.vi
-  @. vf = (x-op.u0)/(op.dt)
-  z = zero(eltype(A))
-  fillstored!(A,z)
-  jacobian!(A,op.odeop,op.ti,(uf,vf),2,1.0/(op.dt),op.ode_cache)
-end
-
-function jacobian!(A::AbstractMatrix,op::RungeKuttaStageNonlinearOperator,x::AbstractVector,
-  i::Integer,γᵢ::Real)
-  ui = x
-  vi = op.vi
-  @. vi = (x-op.u0)/(op.dt)
+  u = op.ui
+  @. u = op.u0
+  for j in 1:op.i-1
+    @. u = u + op.dt * op.a[op.i,j] * op.ki[j]
+  end
+  @. u = u + op.dt * op.a[op.i,op.i] * x
   z = zero(eltype(A))
   fillstored!(A,z)
-  jacobian!(A,op.odeop,op.ti,(ui,vi),i,γᵢ,op.ode_cache)
+  jacobians!(A,op.odeop,op.ti,(u,x),(op.dt*op.a[op.i,op.i],1.0),op.ode_cache)
 end
 
 function allocate_residual(op::RungeKuttaNonlinearOperator,x::AbstractVector)
@@ -213,29 +148,55 @@ function zero_initial_guess(op::RungeKuttaNonlinearOperator)
   x0
 end
 
-function rhs!(op::RungeKuttaNonlinearOperator, x::AbstractVector)
-  u = x
-  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)
+function rhs!(op::RungeKuttaStageNonlinearOperator, x::AbstractVector)
+  u = op.ui
+  @. u = op.u0
+  for j in 1:op.i-1
+    @. u = u + op.dt * op.a[op.i,j] * op.ki[j]
+  end
+  @. u = u + op.dt * op.a[op.i,op.i] * x
+  rhs!(op.rhs,op.odeop,op.ti,(u,x),op.ode_cache)
 end
 
 function lhs!(b::AbstractVector, op::RungeKuttaNonlinearOperator, x::AbstractVector)
-  u = x
-  v = op.vi
-  @. v = (x-op.u0)/(op.dt)
-  lhs!(b,op.odeop,op.ti,(u,v),op.ode_cache)
+  u = op.ui
+  @. u *= 0
+  lhs!(b,op.odeop,op.ti,(u,x),op.ode_cache)
 end
 
-function update!(op::RungeKuttaNonlinearOperator,ti::Float64,fi::AbstractVector,i::Int)
+function update!(op::RungeKuttaNonlinearOperator,ti::Float64,i::Int)
   op.ti = ti
-  op.fi = fi
   op.i = i
 end
 
-function _mass_matrix!(A,odeop,t,ode_cache,u)
-  z = zero(eltype(A))
-  fillstored!(A,z)
-  jacobian!(A,odeop,t,(u,u),2,1.0,ode_cache)
+# Redefining solve! function to enforce computation of the jacobian within
+# each stage of the Runge-Kutta method when the solver is "LinearSolver".
+function solve!(x::AbstractVector,
+  ls::LinearSolver,
+  op::RungeKuttaNonlinearOperator,
+  cache::Nothing)
+  fill!(x,zero(eltype(x)))
+  b = residual(op, x)
+  A = jacobian(op, x)
+  ss = symbolic_setup(ls, A)
+  ns = numerical_setup(ss,A)
+  rmul!(b,-1)
+  solve!(x,ns,b)
+  LinearSolverCache(A,b,ns)
+end
+
+function solve!(x::AbstractVector,
+  ls::LinearSolver,
+  op::RungeKuttaNonlinearOperator,
+  cache)
+  fill!(x,zero(eltype(x)))
+  b = cache.b
+  A = cache.A
+  ns = cache.ns
+  residual!(b, op, x)
+  jacobian!(A, op, x)
+  numerical_setup!(ns,A)
+  rmul!(b,-1)
+  solve!(x,ns,b)
+  cache
 end

test/ODEsTests/ODEsTests/ODESolversTests.jl

@@ -136,6 +136,8 @@ 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
+# uf-u0 = dt*u1 => uf = u1
 odesol = RungeKutta(ls,ls,dt,:BE_1_0_1)
 cache = nothing
 uf, tf, cache = solve_step!(uf,odesol,op,u0,t0,cache)
@@ -144,6 +146,11 @@ uf, tf, cache = solve_step!(uf,odesol,op,u0,t0,cache)
 @test test_ode_solver(odesol,op,u0,t0,tf)
 
 # RK: CN 2nd order
+# k1 = u0
+# k2 = u0 + dt * 0.5 * u0 + dt * 0.5 * k2
+# k2 = u0 * (1+dt*0.5)/(1-dt*0.5)
+# un+1 = u0 + dt * 0.5 * u0 + dt * 0.5 * u0 * (1+dt*0.5)/(1-dt*0.5)
+# un+1 = u0 * (1+ dt * 0.5 + dt * 0.5* (1+dt*0.5)/(1-dt*0.5))
 odesol = RungeKutta(ls,ls,dt,:CN_2_0_2)
 cache = nothing
 uf, tf, cache = solve_step!(uf,odesol,op,u0,t0,cache)
@@ -152,6 +159,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
+# k1 = u0 + dt * 0.25 * k1
+# k1 = u0 * 1/(1-dt*0.25)
+# k2 = u0 + dt * 0.5 * k1 + dt * 0.25 * k2
+# k2 = u0 * 1/(1-dt*0.25) * (1 + dt*0.5/(1-dt*0.25))
+# un+1 = u0 + dt * 0.5 * k1 + dt * 0.5 * k2
+# un+1 = u0 + dt * 0.5 * u0 * 1/(1-dt*0.25) + dt * 0.5 * u0 * 1/(1-dt*0.25) * (1 + dt*0.5/(1-dt*0.25))
+# un+1 = u0 * (1 + dt*0.5/(1-dt*0.25) + dt*0.5/(1-dt*0.25) * (1 + dt*0.5/(1-dt*0.25))
 odesol = RungeKutta(ls,ls,dt,:SDIRK_2_0_2)
 cache = nothing
 uf, tf, cache = solve_step!(uf,odesol,op,u0,t0,cache)

test/ODEsTests/TransientFEsTests/TransientFEOperatorsTests.jl

@@ -31,13 +31,17 @@ V0 = FESpace(
   dirichlet_tags="boundary")
 U = TransientTrialFESpace(V0,u)
 Ω = Triangulation(model)
+# Γ = BoundaryTriangulation(model,tags="boundary")
 degree = 2*order
 dΩ = Measure(Ω,degree)
+# dΓ = Measure(Γ,degree)
+# nΓ = get_normal_vector(Γ)
+# h = 1/partition[1]
 
 # Affine FE operator
-a(u,v) = ∫(∇(v)⊙∇(u))dΩ
+a(u,v) = ∫(∇(v)⊙∇(u))dΩ #- ∫(0.0*v⋅(nΓ⋅∇(u))  + u⋅(nΓ⋅∇(v)) - 10/h*(v⋅u))dΓ
 m(u,v) = ∫(v*u)dΩ
-b(v,t) = ∫(v*f(t))dΩ
+b(v,t) = ∫(v*f(t))dΩ #- ∫(u(t)⋅(nΓ⋅∇(v)) - 10/h*(v⋅u(t)) )dΓ
 res(t,u,v) = a(u,v) + m(∂t(u),v) - b(v,t)
 lhs(t,u,v) = m(∂t(u),v)
 rhs(t,u,v) = b(v,t) - a(u,v)