From 5b84683e079de6c893b38c2135f49528aec36e0c Mon Sep 17 00:00:00 2001 From: Michael Kraus Date: Wed, 24 Jan 2024 17:29:41 +0100 Subject: [PATCH] Copy internal stages to solution step for all Runge-Kutta integrators. --- src/integrators/rk/integrators_dirk.jl | 23 +++++++++++++++ src/integrators/rk/integrators_eprk.jl | 28 +++++++++++++++++++ src/integrators/rk/integrators_erk.jl | 19 +++++++++++++ src/integrators/rk/integrators_iprk.jl | 27 ++++++++++++++++++ .../rk/integrators_iprk_implicit.jl | 3 ++ src/integrators/rk/integrators_irk.jl | 25 +++++++++++++++++ .../rk/integrators_irk_implicit.jl | 28 +++++++++++++++++++ src/integrators/vi/vprk_integrator.jl | 14 ++++++---- 8 files changed, 162 insertions(+), 5 deletions(-) diff --git a/src/integrators/rk/integrators_dirk.jl b/src/integrators/rk/integrators_dirk.jl index 114733014..07391fe0e 100644 --- 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 index 0a25a5219..55ff5f362 100644 --- 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 index 344824264..4a674469b 100644 --- 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 index d200e7676..f938caa6e 100644 --- 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 index 2b1f39c7b..e4321a2ad 100644 --- 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 index ddf9e8369..d196642ad 100644 --- 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 index 604cd3a09..c57097404 100644 --- 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 index 87c3de088..ca4147cf6 100644 --- 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