Copy internal stages to solution step for all Runge-Kutta integrators.

JuliaGNI · Jan 24, 2024 · 5b84683 · 5b84683
1 parent 14c5da2
commit 5b84683
Show file tree

Hide file tree

Showing 8 changed files with 162 additions and 5 deletions.
diff --git a/src/integrators/rk/integrators_dirk.jl b/src/integrators/rk/integrators_dirk.jl
@@ -95,6 +95,26 @@ reset!(cache::DIRKCache, t, q, λ = missing) = copyto!(cache.q̄, q)
 nlsolution(cache::DIRKCache, i) = cache.x[i]
 
 
+function internal_variables(method::DIRKMethod, problem::AbstractProblemODE{DT,TT}) where {DT,TT}
+    S = nstages(method)
+    D = ndims(problem)
+
+    Q = create_internal_stage_vector(DT, D, S)
+    V = create_internal_stage_vector(DT, D, S)
+    Y = create_internal_stage_vector(DT, D, S)
+
+    # solver = get_solver_status(int.solver)
+
+    (Q=Q, V=V, Y=Y)#, solver=solver)
+end
+
+function copy_internal_variables(solstep::SolutionStep, cache::DIRKCache)
+    haskey(internal(solstep), :Q) && copyto!(internal(solstep).Q, cache.Q)
+    haskey(internal(solstep), :V) && copyto!(internal(solstep).V, cache.V)
+    haskey(internal(solstep), :Y) && copyto!(internal(solstep).Y, cache.Y)
+end
+
+
 function initial_guess!(int::GeometricIntegrator{<:DIRK})
     for i in eachstage(int)
         initialguess!(solstep(int).t̄ + timestep(int) * tableau(int).c[i], cache(int).Q[i], cache(int).V[i], solstep(int), problem(int), iguess(int))
@@ -184,4 +204,7 @@ function integrate_step!(int::GeometricIntegrator{<:DIRK, <:AbstractProblemODE})
 
     # compute final update
     update!(solstep(int), cache(int).V, tableau(int), timestep(int))
+
+    # copy internal stage variables
+    copy_internal_variables(solstep(int), cache(int))
 end
diff --git a/src/integrators/rk/integrators_eprk.jl b/src/integrators/rk/integrators_eprk.jl
@@ -86,6 +86,31 @@ end
 @inline CacheType(ST, problem::EquationProblem, method::EPRK) = EPRKCache{ST, ndims(problem), nstages(tableau(method))}
 
 
+function internal_variables(method::EPRK, problem::AbstractProblemPODE{DT,TT}) where {DT,TT}
+    S = nstages(method)
+    D = ndims(problem)
+
+    Q = create_internal_stage_vector_with_zero(DT, D, S)
+    P = create_internal_stage_vector_with_zero(DT, D, S)
+
+    Y = create_internal_stage_vector(DT, D, S)
+    Z = create_internal_stage_vector(DT, D, S)
+    V = create_internal_stage_vector(DT, D, S)
+    F = create_internal_stage_vector(DT, D, S)
+
+    (Q=Q, P=P, V=V, F=F, Y=Y, Z=Z)
+end
+
+function copy_internal_variables(solstep::SolutionStep, cache::EPRKCache)
+    haskey(internal(solstep), :Q) && copyto!(internal(solstep).Q, cache.Q)
+    haskey(internal(solstep), :P) && copyto!(internal(solstep).P, cache.P)
+    haskey(internal(solstep), :Y) && copyto!(internal(solstep).Y, cache.Y)
+    haskey(internal(solstep), :Z) && copyto!(internal(solstep).Z, cache.Z)
+    haskey(internal(solstep), :V) && copyto!(internal(solstep).V, cache.V)
+    haskey(internal(solstep), :F) && copyto!(internal(solstep).F, cache.F)
+end
+
+
 # Compute Q stages of explicit partitioned Runge-Kutta methods.
 function compute_stage_q!(solstep::SolutionStepPODE, problem::AbstractProblemPODE, method::EPRK, cache, i, jmax, t)
     # obtain cache
@@ -155,4 +180,7 @@ function integrate_step!(int::GeometricIntegrator{<:EPRK, <:AbstractProblemPODE}
 
     # compute final update
     update!(solstep(int), cache(int).V, cache(int).F, tableau(int), timestep(int))
+
+    # copy internal stage variables
+    copy_internal_variables(solstep(int), cache(int))
 end
diff --git a/src/integrators/rk/integrators_erk.jl b/src/integrators/rk/integrators_erk.jl
@@ -58,6 +58,22 @@ end
 @inline CacheType(ST, problem::AbstractProblemODE, method::ERK) = ERKCache{ST, ndims(problem), nstages(tableau(method))}
 
 
+function internal_variables(method::ERK, problem::AbstractProblemODE{DT,TT}) where {DT,TT}
+    S = nstages(method)
+    D = ndims(problem)
+
+    Q = create_internal_stage_vector(DT, D, S)
+    V = create_internal_stage_vector(DT, D, S)
+
+    (Q=Q, V=V)
+end
+
+function copy_internal_variables(solstep::SolutionStep, cache::ERKCache)
+    haskey(internal(solstep), :Q) && copyto!(internal(solstep).Q, cache.Q)
+    haskey(internal(solstep), :V) && copyto!(internal(solstep).V, cache.V)
+end
+
+
 function integrate_step!(int::GeometricIntegrator{<:ERK, <:AbstractProblemODE})
     # obtain cache
     local Q = cache(int).Q
@@ -78,4 +94,7 @@ function integrate_step!(int::GeometricIntegrator{<:ERK, <:AbstractProblemODE})
 
     # compute final update
     update!(solstep(int), V, tableau(int), timestep(int))
+
+    # copy internal stage variables
+    copy_internal_variables(solstep(int), cache(int))
 end
diff --git a/src/integrators/rk/integrators_iprk.jl b/src/integrators/rk/integrators_iprk.jl
@@ -149,6 +149,30 @@ function reset!(cache::IPRKCache, t, q, p)
 end
 
 
+function internal_variables(method::IPRK, problem::AbstractProblemPODE{DT,TT}) where {DT,TT}
+    S = nstages(method)
+    D = ndims(problem)
+
+    Q = create_internal_stage_vector(DT, D, S)
+    P = create_internal_stage_vector(DT, D, S)
+    V = create_internal_stage_vector(DT, D, S)
+    F = create_internal_stage_vector(DT, D, S)
+    Y = create_internal_stage_vector(DT, D, S)
+    Z = create_internal_stage_vector(DT, D, S)
+
+    (Q=Q, P=P, V=V, F=F, Y=Y, Z=Z)
+end
+
+function copy_internal_variables(solstep::SolutionStep, cache::IPRKCache)
+    haskey(internal(solstep), :Q) && copyto!(internal(solstep).Q, cache.Q)
+    haskey(internal(solstep), :P) && copyto!(internal(solstep).P, cache.P)
+    haskey(internal(solstep), :V) && copyto!(internal(solstep).V, cache.V)
+    haskey(internal(solstep), :F) && copyto!(internal(solstep).F, cache.F)
+    haskey(internal(solstep), :Y) && copyto!(internal(solstep).Y, cache.Y)
+    haskey(internal(solstep), :Z) && copyto!(internal(solstep).Z, cache.Z)
+end
+
+
 function initial_guess!(int::GeometricIntegrator{<:IPRK, <:AbstractProblemPODE})
     # get cache for nonlinear solution vector and internal stages
     local x = nlsolution(int)
@@ -249,4 +273,7 @@ function integrate_step!(int::GeometricIntegrator{<:IPRK, <:AbstractProblemPODE}
 
     # compute final update
     update!(solstep(int), cache(int).V, cache(int).F, tableau(int), timestep(int))
+
+    # copy internal stage variables
+    copy_internal_variables(solstep(int), cache(int))
 end
diff --git a/src/integrators/rk/integrators_iprk_implicit.jl b/src/integrators/rk/integrators_iprk_implicit.jl
@@ -140,4 +140,7 @@ function integrate_step!(int::GeometricIntegrator{<:IPRK, <:AbstractProblemIODE}
 
     # compute final update
     update!(nlsolution(int), int)
+
+    # copy internal stage variables
+    copy_internal_variables(solstep(int), cache(int))
 end
diff --git a/src/integrators/rk/integrators_irk.jl b/src/integrators/rk/integrators_irk.jl
@@ -125,6 +125,25 @@ end
 @inline CacheType(ST, problem::AbstractProblem, method::IRKMethod) = IRKCache{ST, ndims(problem), nstages(tableau(method))}
 
 
+function internal_variables(method::IRKMethod, problem::AbstractProblemODE{DT,TT}) where {DT,TT}
+    S = nstages(method)
+    D = ndims(problem)
+
+    Q = create_internal_stage_vector(DT, D, S)
+    V = create_internal_stage_vector(DT, D, S)
+    Y = create_internal_stage_vector(DT, D, S)
+
+    # solver = get_solver_status(int.solver)
+
+    (Q=Q, V=V, Y=Y)#, solver=solver)
+end
+
+function copy_internal_variables(solstep::SolutionStep, cache::IRKCache)
+    haskey(internal(solstep), :Q) && copyto!(internal(solstep).Q, cache.Q)
+    haskey(internal(solstep), :V) && copyto!(internal(solstep).V, cache.V)
+    haskey(internal(solstep), :Y) && copyto!(internal(solstep).Y, cache.Y)
+end
+
 
 function initial_guess!(int::GeometricIntegrator{<:IRK, <:AbstractProblemODE})
     # compute initial guess for internal stages
@@ -227,4 +246,10 @@ function integrate_step!(int::GeometricIntegrator{<:IRK, <:AbstractProblemODE})
 
     # compute final update
     update!(nlsolution(int), int)
+
+    # copy internal stage variables
+    copy_internal_variables(solstep(int), cache(int))
+
+    # copy solver status
+    # get_solver_status!(solver(int), solstep(int).internal[:solver])
 end
diff --git a/src/integrators/rk/integrators_irk_implicit.jl b/src/integrators/rk/integrators_irk_implicit.jl
@@ -77,6 +77,28 @@ function solversize(problem::AbstractProblemIODE, method::IRK)
 end
 
 
+function internal_variables(method::IRK, problem::AbstractProblemIODE{DT,TT}) where {DT,TT}
+    S = nstages(method)
+    D = ndims(problem)
+
+    Q = create_internal_stage_vector(DT, D, S)
+    V = create_internal_stage_vector(DT, D, S)
+    Θ = create_internal_stage_vector(DT, D, S)
+    F = create_internal_stage_vector(DT, D, S)
+
+    # solver = get_solver_status(int.solver)
+
+    (Q=Q, V=V, Θ=Θ, F=F)#, solver=solver)
+end
+
+function copy_internal_variables(solstep::SolutionStep, cache::IRKimplicitCache)
+    haskey(internal(solstep), :Q) && copyto!(internal(solstep).Q, cache.Q)
+    haskey(internal(solstep), :V) && copyto!(internal(solstep).V, cache.V)
+    haskey(internal(solstep), :Θ) && copyto!(internal(solstep).Θ, cache.Θ)
+    haskey(internal(solstep), :F) && copyto!(internal(solstep).F, cache.F)
+end
+
+
 function Base.show(io::IO, int::GeometricIntegrator{<:IRK, <:AbstractProblemIODE})
     print(io, "\nRunge-Kutta Integrator for Implicit Equations with:\n")
     print(io, "   Timestep: $(timestep(int))\n")
@@ -247,4 +269,10 @@ function integrate_step!(int::GeometricIntegrator{<:IRK, <:AbstractProblemIODE})
 
     # compute final update
     update!(nlsolution(int), int)
+
+    # copy internal stage variables
+    copy_internal_variables(solstep(int), cache(int))
+
+    # copy solver status
+    # get_solver_status!(solver(int), solstep(int).internal[:solver])
 end
diff --git a/src/integrators/vi/vprk_integrator.jl b/src/integrators/vi/vprk_integrator.jl
@@ -13,18 +13,22 @@ function internal_variables(method::VPRKMethod, problem::AbstractProblemIODE{DT,
     P = create_internal_stage_vector(DT, D, S)
     V = create_internal_stage_vector(DT, D, S)
     F = create_internal_stage_vector(DT, D, S)
+    Y = create_internal_stage_vector(DT, D, S)
+    Z = create_internal_stage_vector(DT, D, S)
 
     # solver = get_solver_status(int.solver)
 
-    (Q=Q, P=P, V=V, F=F)#, solver=solver)
+    (Q=Q, P=P, V=V, F=F, Y=Y, Z=Z)#, solver=solver)
 end
 
 
 function copy_internal_variables(solstep::SolutionStep, cache::VPRKCache)
-    haskey(internal(solstep), :Q) && coypto!(internal(solstep).Q, cache.Q)
-    haskey(internal(solstep), :P) && coypto!(internal(solstep).P, cache.P)
-    haskey(internal(solstep), :V) && coypto!(internal(solstep).V, cache.V)
-    haskey(internal(solstep), :F) && coypto!(internal(solstep).F, cache.F)
+    haskey(internal(solstep), :Q) && copyto!(internal(solstep).Q, cache.Q)
+    haskey(internal(solstep), :P) && copyto!(internal(solstep).P, cache.P)
+    haskey(internal(solstep), :V) && copyto!(internal(solstep).V, cache.V)
+    haskey(internal(solstep), :F) && copyto!(internal(solstep).F, cache.F)
+    haskey(internal(solstep), :Y) && copyto!(internal(solstep).Y, cache.Y)
+    haskey(internal(solstep), :Z) && copyto!(internal(solstep).Z, cache.Z)
 end