diff --git a/lib/OrdinaryDiffEqLowStorageRK/src/OrdinaryDiffEqLowStorageRK.jl b/lib/OrdinaryDiffEqLowStorageRK/src/OrdinaryDiffEqLowStorageRK.jl index 82d338344b2..02b56384365 100644 --- a/lib/OrdinaryDiffEqLowStorageRK/src/OrdinaryDiffEqLowStorageRK.jl +++ b/lib/OrdinaryDiffEqLowStorageRK/src/OrdinaryDiffEqLowStorageRK.jl @@ -25,8 +25,10 @@ using Reexport include("arrayfuse.jl") include("algorithms.jl") include("alg_utils.jl") +include("low_storage_tableaus.jl") include("low_storage_rk_caches.jl") include("low_storage_rk_perform_step.jl") +include("generic_low_storage_perform_step.jl") import PrecompileTools import Preferences diff --git a/lib/OrdinaryDiffEqLowStorageRK/src/generic_low_storage_perform_step.jl b/lib/OrdinaryDiffEqLowStorageRK/src/generic_low_storage_perform_step.jl new file mode 100644 index 00000000000..e9627dcbe9c --- /dev/null +++ b/lib/OrdinaryDiffEqLowStorageRK/src/generic_low_storage_perform_step.jl @@ -0,0 +1,294 @@ +@muladd function _perform_step_iip!( + integrator, cache, tab::LowStorageRKTableau{TwoN} + ) + (; t, dt, u, f, p) = integrator + (; k, tmp, williamson_condition, stage_limiter!, step_limiter!, thread) = cache + (; A2end, B1, B2end, c2end) = tab + + f(k, u, p, t) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + @.. broadcast = false thread = thread tmp = dt * k + @.. broadcast = false thread = thread u = u + B1 * tmp + + for i in eachindex(A2end) + if williamson_condition + f(ArrayFuse(tmp, u, (A2end[i], dt, B2end[i])), u, p, t + c2end[i] * dt) + else + @.. broadcast = false thread = thread tmp = A2end[i] * tmp + stage_limiter!(u, integrator, p, t + c2end[i] * dt) + f(k, u, p, t + c2end[i] * dt) + @.. broadcast = false thread = thread tmp = tmp + dt * k + @.. broadcast = false thread = thread u = u + B2end[i] * tmp + end + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + end + stage_limiter!(u, integrator, p, t + dt) + step_limiter!(u, integrator, p, t + dt) + return nothing +end + +@muladd function _perform_step_oop!( + integrator, tab::LowStorageRKTableau{TwoN} + ) + (; t, dt, u, f, p) = integrator + (; A2end, B1, B2end, c2end) = tab + + tmp = dt * integrator.fsalfirst + u = u + B1 * tmp + + for i in eachindex(A2end) + k = f(u, p, t + c2end[i] * dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + tmp = A2end[i] * tmp + dt * k + u = u + B2end[i] * tmp + end + + integrator.k[1] = integrator.fsalfirst + integrator.fsalfirst = f(u, p, t + dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + integrator.u = u + return nothing +end + +@muladd function _perform_step_iip!( + integrator, cache, tab::LowStorageRKTableau{TwoC} + ) + (; t, dt, u, f, p) = integrator + (; k, fsalfirst, tmp, stage_limiter!, step_limiter!, thread) = cache + (; A2end, B1, B2end, c2end) = tab + + @.. broadcast = false thread = thread k = integrator.fsalfirst + @.. broadcast = false thread = thread u = u + B1 * dt * k + + for i in eachindex(A2end) + @.. broadcast = false thread = thread tmp = u + A2end[i] * dt * k + f(k, tmp, p, t + c2end[i] * dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + @.. broadcast = false thread = thread u = u + B2end[i] * dt * k + end + step_limiter!(u, integrator, p, t + dt) + f(k, u, p, t + dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + return nothing +end + +@muladd function _perform_step_oop!( + integrator, tab::LowStorageRKTableau{TwoC} + ) + (; t, dt, u, f, p) = integrator + (; A2end, B1, B2end, c2end) = tab + + k = integrator.fsalfirst = integrator.f(u, p, t) + integrator.k[1] = integrator.fsalfirst + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + u = u + B1 * dt * k + + for i in eachindex(A2end) + tmp = u + A2end[i] * dt * k + k = integrator.f(tmp, p, t + c2end[i] * dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + u = u + B2end[i] * dt * k + end + + integrator.u = u + return nothing +end + +@muladd function _perform_step_oop!(integrator, tab::LowStorageRK3SConstantCache) + (; t, dt, uprev, u, f, p) = integrator + (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end) = tab + + tmp = u + u = tmp + β1 * dt * integrator.fsalfirst + + for i in eachindex(γ12end) + k = f(u, p, t + c2end[i] * dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + tmp = tmp + δ2end[i] * u + u = γ12end[i] * u + γ22end[i] * tmp + γ32end[i] * uprev + β2end[i] * dt * k + end + + integrator.fsallast = f(u, p, t + dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + integrator.k[1] = integrator.fsalfirst + integrator.u = u + return nothing +end + +@muladd function _perform_step_iip!(integrator, cache, tab::LowStorageRK3SConstantCache) + (; t, dt, uprev, u, f, p) = integrator + (; k, fsalfirst, tmp, stage_limiter!, step_limiter!, thread) = cache + (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end) = tab + + @.. broadcast = false thread = thread tmp = u + @.. broadcast = false thread = thread u = tmp + β1 * dt * integrator.fsalfirst + + for i in eachindex(γ12end) + f(k, u, p, t + c2end[i] * dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + @.. broadcast = false thread = thread tmp = tmp + δ2end[i] * u + @.. broadcast = false thread = thread u = γ12end[i] * u + γ22end[i] * tmp + + γ32end[i] * uprev + + β2end[i] * dt * k + end + + step_limiter!(u, integrator, p, t + dt) + f(k, u, p, t + dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + return nothing +end + +@muladd function _perform_step_oop!(integrator, tab::LowStorageRK3SpConstantCache) + (; t, dt, uprev, u, f, p) = integrator + (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end) = tab + + integrator.fsalfirst = f(uprev, p, t) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + integrator.k[1] = integrator.fsalfirst + tmp = uprev + u = tmp + β1 * dt * integrator.fsalfirst + utilde = u + if integrator.opts.adaptive + utilde = bhat1 * dt * integrator.fsalfirst + end + + for i in eachindex(γ12end) + k = f(u, p, t + c2end[i] * dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + tmp = tmp + δ2end[i] * u + u = γ12end[i] * u + γ22end[i] * tmp + γ32end[i] * uprev + β2end[i] * dt * k + if integrator.opts.adaptive + utilde = utilde + bhat2end[i] * dt * k + end + end + + if integrator.opts.adaptive + atmp = calculate_residuals( + utilde, uprev, u, integrator.opts.abstol, + integrator.opts.reltol, integrator.opts.internalnorm, t + ) + OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t)) + end + + integrator.u = u + return nothing +end + +@muladd function _perform_step_iip!(integrator, cache, tab::LowStorageRK3SpConstantCache) + (; t, dt, uprev, u, f, p) = integrator + (; k, tmp, utilde, atmp, stage_limiter!, step_limiter!, thread) = cache + (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end) = tab + + f(integrator.fsalfirst, uprev, p, t) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + @.. broadcast = false thread = thread tmp = uprev + @.. broadcast = false thread = thread u = tmp + β1 * dt * integrator.fsalfirst + if integrator.opts.adaptive + @.. broadcast = false thread = thread utilde = bhat1 * dt * integrator.fsalfirst + end + + for i in eachindex(γ12end) + stage_limiter!(u, integrator, p, t + c2end[i] * dt) + f(k, u, p, t + c2end[i] * dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + @.. broadcast = false thread = thread tmp = tmp + δ2end[i] * u + @.. broadcast = false thread = thread u = γ12end[i] * u + γ22end[i] * tmp + + γ32end[i] * uprev + β2end[i] * dt * k + if integrator.opts.adaptive + @.. broadcast = false thread = thread utilde = utilde + bhat2end[i] * dt * k + end + end + + stage_limiter!(u, integrator, p, t + dt) + step_limiter!(u, integrator, p, t + dt) + + if integrator.opts.adaptive + calculate_residuals!( + atmp, utilde, uprev, u, integrator.opts.abstol, + integrator.opts.reltol, integrator.opts.internalnorm, t, + thread + ) + OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t)) + end + return nothing +end + +@muladd function _perform_step_oop!(integrator, tab::LowStorageRK3SpFSALConstantCache) + (; t, dt, uprev, u, f, p) = integrator + (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end, bhatfsal) = tab + + tmp = uprev + u = tmp + β1 * dt * integrator.fsalfirst + utilde = u + if integrator.opts.adaptive + utilde = bhat1 * dt * integrator.fsalfirst + end + + for i in eachindex(γ12end) + k = f(u, p, t + c2end[i] * dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + tmp = tmp + δ2end[i] * u + u = γ12end[i] * u + γ22end[i] * tmp + γ32end[i] * uprev + β2end[i] * dt * k + if integrator.opts.adaptive + utilde = utilde + bhat2end[i] * dt * k + end + end + + integrator.fsallast = f(u, p, t + dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + + if integrator.opts.adaptive + utilde = utilde + bhatfsal * dt * integrator.fsallast + atmp = calculate_residuals( + utilde, uprev, u, integrator.opts.abstol, + integrator.opts.reltol, integrator.opts.internalnorm, t + ) + OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t)) + end + + integrator.k[1] = integrator.fsalfirst + integrator.k[2] = integrator.fsallast + integrator.u = u + return nothing +end + +@muladd function _perform_step_iip!(integrator, cache, tab::LowStorageRK3SpFSALConstantCache) + (; t, dt, uprev, u, f, p) = integrator + (; k, tmp, utilde, atmp, stage_limiter!, step_limiter!, thread) = cache + (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end, bhatfsal) = tab + + @.. broadcast = false thread = thread tmp = uprev + @.. broadcast = false thread = thread u = tmp + β1 * dt * integrator.fsalfirst + if integrator.opts.adaptive + @.. broadcast = false thread = thread utilde = bhat1 * dt * integrator.fsalfirst + end + + for i in eachindex(γ12end) + stage_limiter!(u, integrator, p, t + c2end[i] * dt) + f(k, u, p, t + c2end[i] * dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + @.. broadcast = false thread = thread tmp = tmp + δ2end[i] * u + @.. broadcast = false thread = thread u = γ12end[i] * u + γ22end[i] * tmp + + γ32end[i] * uprev + β2end[i] * dt * k + if integrator.opts.adaptive + @.. broadcast = false thread = thread utilde = utilde + bhat2end[i] * dt * k + end + end + + stage_limiter!(u, integrator, p, t + dt) + step_limiter!(u, integrator, p, t + dt) + + f(k, u, p, t + dt) + OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) + + if integrator.opts.adaptive + @.. broadcast = false thread = thread utilde = utilde + bhatfsal * dt * k + calculate_residuals!( + atmp, utilde, uprev, u, integrator.opts.abstol, + integrator.opts.reltol, integrator.opts.internalnorm, t, + thread + ) + OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t)) + end + return nothing +end diff --git a/lib/OrdinaryDiffEqLowStorageRK/src/low_storage_rk_caches.jl b/lib/OrdinaryDiffEqLowStorageRK/src/low_storage_rk_caches.jl index 84f06177900..f6f0cc70cea 100644 --- a/lib/OrdinaryDiffEqLowStorageRK/src/low_storage_rk_caches.jl +++ b/lib/OrdinaryDiffEqLowStorageRK/src/low_storage_rk_caches.jl @@ -17,13 +17,6 @@ get_fsalfirstlast(cache::LowStorageRKMutableCache, u) = (cache.fsalfirst, cache. thread::Thread end -struct LowStorageRK2NConstantCache{N, T, T2} <: OrdinaryDiffEqConstantCache - A2end::NTuple{N, T} # A1 is always zero - B1::T - B2end::NTuple{N, T} - c2end::NTuple{N, T2} # c1 is always zero -end - function ORK256ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} A2 = convert(T, -1.0) A3 = convert(T, -1.55798) @@ -44,7 +37,7 @@ function ORK256ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} c5 = convert(T2, 0.8) c2end = (c2, c3, c4, c5) - return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end) + return LowStorageRKTableau{TwoN}(A2end, B1, B2end, c2end) end function alg_cache( @@ -144,7 +137,7 @@ function CarpenterKennedy2N54ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} c5 = convert(T2, 2802321613138 // 2924317926251) c2end = (c2, c3, c4, c5) - return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end) + return LowStorageRKTableau{TwoN}(A2end, B1, B2end, c2end) end @cache mutable struct SHLDDRK_2NCache{ @@ -359,7 +352,7 @@ function SHLDDRK64ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} c6 = convert(T2, 0.9271047) c2end = (c2, c3, c4, c5, c6) - return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end) + return LowStorageRKTableau{TwoN}(A2end, B1, B2end, c2end) end function DGLDDRK73_CConstantCache(::Type{T}, ::Type{T2}) where {T, T2} @@ -388,7 +381,7 @@ function DGLDDRK73_CConstantCache(::Type{T}, ::Type{T2}) where {T, T2} c7 = convert(T2, 0.998046608462379) c2end = (c2, c3, c4, c5, c6, c7) - return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end) + return LowStorageRKTableau{TwoN}(A2end, B1, B2end, c2end) end function DGLDDRK84_CConstantCache(::Type{T}, ::Type{T2}) where {T, T2} @@ -420,7 +413,7 @@ function DGLDDRK84_CConstantCache(::Type{T}, ::Type{T2}) where {T, T2} c8 = convert(T2, 0.919902896453866) c2end = (c2, c3, c4, c5, c6, c7, c8) - return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end) + return LowStorageRKTableau{TwoN}(A2end, B1, B2end, c2end) end function DGLDDRK84_FConstantCache(::Type{T}, ::Type{T2}) where {T, T2} @@ -452,7 +445,7 @@ function DGLDDRK84_FConstantCache(::Type{T}, ::Type{T2}) where {T, T2} c8 = convert(T2, 0.9484087623348481) c2end = (c2, c3, c4, c5, c6, c7, c8) - return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end) + return LowStorageRKTableau{TwoN}(A2end, B1, B2end, c2end) end function NDBLSRK124ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} @@ -496,7 +489,7 @@ function NDBLSRK124ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} c12 = convert(T2, 0.9032588871651854) c2end = (c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12) - return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end) + return LowStorageRKTableau{TwoN}(A2end, B1, B2end, c2end) end function NDBLSRK134ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} @@ -543,7 +536,7 @@ function NDBLSRK134ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} c13 = convert(T2, 0.9126827615920843) c2end = (c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13) - return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end) + return LowStorageRKTableau{TwoN}(A2end, B1, B2end, c2end) end function NDBLSRK144ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} @@ -593,7 +586,7 @@ function NDBLSRK144ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} c14 = convert(T2, 0.8734213127600976) c2end = (c2, c3, c4, c5, c6, c7, c8, c9, c10, c11, c12, c13, c14) - return LowStorageRK2NConstantCache(A2end, B1, B2end, c2end) + return LowStorageRKTableau{TwoN}(A2end, B1, B2end, c2end) end # 2C low storage methods introduced by Calvo, Franco, Rández (2004) @@ -612,13 +605,6 @@ end thread::Thread end -struct LowStorageRK2CConstantCache{N, T, T2} <: OrdinaryDiffEqConstantCache - A2end::NTuple{N, T} # A1 is always zero - B1::T - B2end::NTuple{N, T} - c2end::NTuple{N, T2} # c1 is always zero -end - function CFRLDDRK64ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} A2 = convert(T, 0.17985400977138) A3 = convert(T, 0.14081893152111) @@ -642,7 +628,7 @@ function CFRLDDRK64ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} c6 = convert(T2, 0.83050587987157) c2end = (c2, c3, c4, c5, c6) - return LowStorageRK2CConstantCache(A2end, B1, B2end, c2end) + return LowStorageRKTableau{TwoC}(A2end, B1, B2end, c2end) end function TSLDDRK74ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} @@ -671,7 +657,7 @@ function TSLDDRK74ConstantCache(::Type{T}, ::Type{T2}) where {T, T2} c7 = convert(T2, 0.91124223849547205) c2end = (c2, c3, c4, c5, c6, c7) - return LowStorageRK2CConstantCache(A2end, B1, B2end, c2end) + return LowStorageRKTableau{TwoC}(A2end, B1, B2end, c2end) end # 3S low storage methods introduced by Ketcheson diff --git a/lib/OrdinaryDiffEqLowStorageRK/src/low_storage_rk_perform_step.jl b/lib/OrdinaryDiffEqLowStorageRK/src/low_storage_rk_perform_step.jl index dfd62619ed8..3edb6d719a9 100644 --- a/lib/OrdinaryDiffEqLowStorageRK/src/low_storage_rk_perform_step.jl +++ b/lib/OrdinaryDiffEqLowStorageRK/src/low_storage_rk_perform_step.jl @@ -1,38 +1,14 @@ -# 2N low storage methods -function initialize!(integrator, cache::LowStorageRK2NConstantCache) - integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal +function initialize!(integrator, cache::LowStorageRKTableau{TwoN}) + integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) integrator.kshortsize = 1 integrator.k = typeof(integrator.k)(undef, integrator.kshortsize) - - # Avoid undefined entries if k is an array of arrays integrator.fsallast = zero(integrator.fsalfirst) return integrator.k[1] = integrator.fsalfirst end -@muladd function perform_step!( - integrator, cache::LowStorageRK2NConstantCache, - repeat_step = false - ) - (; t, dt, u, f, p) = integrator - (; A2end, B1, B2end, c2end) = cache - - # u1 - tmp = dt * integrator.fsalfirst - u = u + B1 * tmp - - # other stages - for i in eachindex(A2end) - k = f(u, p, t + c2end[i] * dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - tmp = A2end[i] * tmp + dt * k - u = u + B2end[i] * tmp - end - - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - integrator.k[1] = integrator.fsalfirst - integrator.fsalfirst = f(u, p, t + dt) # For interpolation, then FSAL'd - integrator.u = u +function perform_step!(integrator, cache::LowStorageRKTableau{TwoN}, repeat_step = false) + return _perform_step_oop!(integrator, cache) end get_fsalfirstlast(cache::LowStorageRK2NCache, u) = (nothing, nothing) @@ -42,392 +18,115 @@ function initialize!(integrator, cache::LowStorageRK2NCache) integrator.kshortsize = 1 resize!(integrator.k, integrator.kshortsize) integrator.k[1] = k - integrator.f(k, integrator.uprev, integrator.p, integrator.t) # FSAL for interpolation + integrator.f(k, integrator.uprev, integrator.p, integrator.t) return OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) end -@muladd function perform_step!(integrator, cache::LowStorageRK2NCache, repeat_step = false) - (; t, dt, u, f, p) = integrator - (; k, tmp, williamson_condition, stage_limiter!, step_limiter!, thread) = cache - (; A2end, B1, B2end, c2end) = cache.tab - - # u1 - f(k, u, p, t) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - @.. broadcast = false thread = thread tmp = dt * k - @.. broadcast = false thread = thread u = u + B1 * tmp - # other stages - for i in eachindex(A2end) - if williamson_condition - f(ArrayFuse(tmp, u, (A2end[i], dt, B2end[i])), u, p, t + c2end[i] * dt) - else - @.. broadcast = false thread = thread tmp = A2end[i] * tmp - stage_limiter!(u, integrator, p, t + c2end[i] * dt) - f(k, u, p, t + c2end[i] * dt) - @.. broadcast = false thread = thread tmp = tmp + dt * k - @.. broadcast = false thread = thread u = u + B2end[i] * tmp - end - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - end - stage_limiter!(u, integrator, p, t + dt) - step_limiter!(u, integrator, p, t + dt) +function perform_step!(integrator, cache::LowStorageRK2NCache, repeat_step = false) + return _perform_step_iip!(integrator, cache, cache.tab) end -# 2C low storage methods -function initialize!(integrator, cache::LowStorageRK2CConstantCache) - integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal +function initialize!(integrator, cache::LowStorageRKTableau{TwoC}) + integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) integrator.kshortsize = 1 integrator.k = typeof(integrator.k)(undef, integrator.kshortsize) - - # Avoid undefined entries if k is an array of arrays integrator.fsallast = zero(integrator.fsalfirst) return integrator.k[1] = integrator.fsalfirst end -@muladd function perform_step!( - integrator, cache::LowStorageRK2CConstantCache, - repeat_step = false - ) - (; t, dt, u, f, p) = integrator - (; A2end, B1, B2end, c2end) = cache - - # u1 - k = integrator.fsalfirst = f(u, p, t) - integrator.k[1] = integrator.fsalfirst - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - u = u + B1 * dt * k - - # other stages - for i in eachindex(A2end) - tmp = u + A2end[i] * dt * k - k = f(tmp, p, t + c2end[i] * dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - u = u + B2end[i] * dt * k - end - - integrator.u = u +function perform_step!(integrator, cache::LowStorageRKTableau{TwoC}, repeat_step = false) + return _perform_step_oop!(integrator, cache) end function initialize!(integrator, cache::LowStorageRK2CCache) - (; k, fsalfirst) = cache - integrator.kshortsize = 1 resize!(integrator.k, integrator.kshortsize) integrator.k[1] = integrator.fsalfirst - integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) # FSAL for interpolation + integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) return OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) end -@muladd function perform_step!(integrator, cache::LowStorageRK2CCache, repeat_step = false) - (; t, dt, u, f, p) = integrator - (; k, fsalfirst, tmp, stage_limiter!, step_limiter!, thread) = cache - (; A2end, B1, B2end, c2end) = cache.tab - - # u1 - @.. broadcast = false thread = thread k = integrator.fsalfirst - @.. broadcast = false thread = thread u = u + B1 * dt * k - - # other stages - for i in eachindex(A2end) - @.. broadcast = false thread = thread tmp = u + A2end[i] * dt * k - f(k, tmp, p, t + c2end[i] * dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - @.. broadcast = false thread = thread u = u + B2end[i] * dt * k - end - step_limiter!(u, integrator, p, t + dt) - f(k, u, p, t + dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) +function perform_step!(integrator, cache::LowStorageRK2CCache, repeat_step = false) + return _perform_step_iip!(integrator, cache, cache.tab) end # 3S low storage methods function initialize!(integrator, cache::LowStorageRK3SConstantCache) - integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal + integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) integrator.kshortsize = 1 integrator.k = typeof(integrator.k)(undef, integrator.kshortsize) - - # Avoid undefined entries if k is an array of arrays integrator.fsallast = zero(integrator.fsalfirst) return integrator.k[1] = integrator.fsalfirst end -@muladd function perform_step!( - integrator, cache::LowStorageRK3SConstantCache, - repeat_step = false - ) - (; t, dt, uprev, u, f, p) = integrator - (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end) = cache - - # u1 - tmp = u - u = tmp + β1 * dt * integrator.fsalfirst - - # other stages - for i in eachindex(γ12end) - k = f(u, p, t + c2end[i] * dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - tmp = tmp + δ2end[i] * u - u = γ12end[i] * u + γ22end[i] * tmp + γ32end[i] * uprev + β2end[i] * dt * k - end - - integrator.fsallast = f(u, p, t + dt) # For interpolation, then FSAL'd - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - integrator.k[1] = integrator.fsalfirst - integrator.u = u +function perform_step!(integrator, cache::LowStorageRK3SConstantCache, repeat_step = false) + return _perform_step_oop!(integrator, cache) end function initialize!(integrator, cache::LowStorageRK3SCache) - (; k, fsalfirst) = cache - integrator.kshortsize = 1 resize!(integrator.k, integrator.kshortsize) integrator.k[1] = integrator.fsalfirst - integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) # FSAL for interpolation + integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) return OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) end -@muladd function perform_step!(integrator, cache::LowStorageRK3SCache, repeat_step = false) - (; t, dt, uprev, u, f, p) = integrator - (; k, fsalfirst, tmp, stage_limiter!, step_limiter!, thread) = cache - (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end) = cache.tab - - # u1 - @.. broadcast = false thread = thread tmp = u - @.. broadcast = false thread = thread u = tmp + β1 * dt * integrator.fsalfirst - - # other stages - for i in eachindex(γ12end) - f(k, u, p, t + c2end[i] * dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - @.. broadcast = false thread = thread tmp = tmp + δ2end[i] * u - @.. broadcast = false thread = thread u = γ12end[i] * u + γ22end[i] * tmp + - γ32end[i] * uprev + - β2end[i] * dt * k - end - - step_limiter!(u, integrator, p, t + dt) - f(k, u, p, t + dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) +function perform_step!(integrator, cache::LowStorageRK3SCache, repeat_step = false) + return _perform_step_iip!(integrator, cache, cache.tab) end -# 3S+ low storage methods: 3S methods adding another memory location for the embedded method (non-FSAL version) function initialize!(integrator, cache::LowStorageRK3SpConstantCache) - integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal + integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) integrator.kshortsize = 1 integrator.k = typeof(integrator.k)(undef, integrator.kshortsize) - - # Avoid undefined entries if k is an array of arrays integrator.fsallast = zero(integrator.fsalfirst) return integrator.k[1] = integrator.fsalfirst end -@muladd function perform_step!( - integrator, cache::LowStorageRK3SpConstantCache, - repeat_step = false - ) - (; t, dt, uprev, u, f, p) = integrator - (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end) = cache - - # u1 - integrator.fsalfirst = f(uprev, p, t) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - integrator.k[1] = integrator.fsalfirst - tmp = uprev - u = tmp + β1 * dt * integrator.fsalfirst - # Initialize utilde for JET - utilde = u - if integrator.opts.adaptive - utilde = bhat1 * dt * integrator.fsalfirst - end - - # other stages - for i in eachindex(γ12end) - k = f(u, p, t + c2end[i] * dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - tmp = tmp + δ2end[i] * u - u = γ12end[i] * u + γ22end[i] * tmp + γ32end[i] * uprev + β2end[i] * dt * k - if integrator.opts.adaptive - utilde = utilde + bhat2end[i] * dt * k - end - end - - if integrator.opts.adaptive - atmp = calculate_residuals( - utilde, uprev, u, integrator.opts.abstol, - integrator.opts.reltol, integrator.opts.internalnorm, t - ) - OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t)) - end - - integrator.u = u +function perform_step!(integrator, cache::LowStorageRK3SpConstantCache, repeat_step = false) + return _perform_step_oop!(integrator, cache) end function initialize!(integrator, cache::LowStorageRK3SpCache) - (; k, fsalfirst) = cache - integrator.kshortsize = 1 resize!(integrator.k, integrator.kshortsize) return integrator.k[1] = integrator.fsalfirst end -@muladd function perform_step!(integrator, cache::LowStorageRK3SpCache, repeat_step = false) - (; t, dt, uprev, u, f, p) = integrator - (; k, tmp, utilde, atmp, stage_limiter!, step_limiter!, thread) = cache - (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end) = cache.tab - - # u1 - f(integrator.fsalfirst, uprev, p, t) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - @.. broadcast = false thread = thread tmp = uprev - @.. broadcast = false thread = thread u = tmp + β1 * dt * integrator.fsalfirst - if integrator.opts.adaptive - @.. broadcast = false thread = thread utilde = bhat1 * dt * integrator.fsalfirst - end - - # other stages - for i in eachindex(γ12end) - stage_limiter!(u, integrator, p, t + c2end[i] * dt) - f(k, u, p, t + c2end[i] * dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - @.. broadcast = false thread = thread tmp = tmp + δ2end[i] * u - @.. broadcast = false thread = thread u = γ12end[i] * u + γ22end[i] * tmp + - γ32end[i] * uprev + β2end[i] * dt * k - if integrator.opts.adaptive - @.. broadcast = false thread = thread utilde = utilde + bhat2end[i] * dt * k - end - end - - stage_limiter!(u, integrator, p, t + dt) - step_limiter!(u, integrator, p, t + dt) - - if integrator.opts.adaptive - calculate_residuals!( - atmp, utilde, uprev, u, integrator.opts.abstol, - integrator.opts.reltol, integrator.opts.internalnorm, t, - thread - ) - OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t)) - end +function perform_step!(integrator, cache::LowStorageRK3SpCache, repeat_step = false) + return _perform_step_iip!(integrator, cache, cache.tab) end -# 3S+ FSAL low storage methods: 3S methods adding another memory location for the embedded method (FSAL version) function initialize!(integrator, cache::LowStorageRK3SpFSALConstantCache) - integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) # Pre-start fsal + integrator.fsalfirst = integrator.f(integrator.uprev, integrator.p, integrator.t) OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) integrator.kshortsize = 2 integrator.k = typeof(integrator.k)(undef, integrator.kshortsize) - - # Avoid undefined entries if k is an array of arrays integrator.fsallast = zero(integrator.fsalfirst) integrator.k[1] = integrator.fsalfirst return integrator.k[2] = integrator.fsallast end -@muladd function perform_step!( - integrator, cache::LowStorageRK3SpFSALConstantCache, - repeat_step = false +function perform_step!( + integrator, cache::LowStorageRK3SpFSALConstantCache, repeat_step = false ) - (; t, dt, uprev, u, f, p) = integrator - (; γ12end, γ22end, γ32end, δ2end, β1, β2end, c2end, bhat1, bhat2end, bhatfsal) = cache - - # u1 - tmp = uprev - u = tmp + β1 * dt * integrator.fsalfirst - # Initialize utilde for JET - utilde = u - if integrator.opts.adaptive - utilde = bhat1 * dt * integrator.fsalfirst - end - - # other stages - for i in eachindex(γ12end) - k = f(u, p, t + c2end[i] * dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - tmp = tmp + δ2end[i] * u - u = γ12end[i] * u + γ22end[i] * tmp + γ32end[i] * uprev + β2end[i] * dt * k - if integrator.opts.adaptive - utilde = utilde + bhat2end[i] * dt * k - end - end - - # FSAL - integrator.fsallast = f(u, p, t + dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - - if integrator.opts.adaptive - utilde = utilde + bhatfsal * dt * integrator.fsallast - atmp = calculate_residuals( - utilde, uprev, u, integrator.opts.abstol, - integrator.opts.reltol, integrator.opts.internalnorm, t - ) - OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t)) - end - - integrator.k[1] = integrator.fsalfirst - integrator.k[2] = integrator.fsallast - integrator.u = u + return _perform_step_oop!(integrator, cache) end function initialize!(integrator, cache::LowStorageRK3SpFSALCache) - (; k, fsalfirst) = cache - integrator.kshortsize = 2 resize!(integrator.k, integrator.kshortsize) integrator.k[1] = integrator.fsalfirst integrator.k[2] = integrator.fsallast - integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) # FSAL for interpolation + integrator.f(integrator.fsalfirst, integrator.uprev, integrator.p, integrator.t) return OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) end -@muladd function perform_step!( - integrator, cache::LowStorageRK3SpFSALCache, - repeat_step = false - ) - (; t, dt, uprev, u, f, p) = integrator - (; k, tmp, utilde, atmp, stage_limiter!, step_limiter!, thread) = cache - (; - γ12end, γ22end, γ32end, δ2end, β1, β2end, - c2end, bhat1, bhat2end, bhatfsal, - ) = cache.tab - - # u1 - @.. broadcast = false thread = thread tmp = uprev - @.. broadcast = false thread = thread u = tmp + β1 * dt * integrator.fsalfirst - if integrator.opts.adaptive - @.. broadcast = false thread = thread utilde = bhat1 * dt * integrator.fsalfirst - end - - # other stages - for i in eachindex(γ12end) - stage_limiter!(u, integrator, p, t + c2end[i] * dt) - f(k, u, p, t + c2end[i] * dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - @.. broadcast = false thread = thread tmp = tmp + δ2end[i] * u - @.. broadcast = false thread = thread u = γ12end[i] * u + γ22end[i] * tmp + - γ32end[i] * uprev + β2end[i] * dt * k - if integrator.opts.adaptive - @.. broadcast = false thread = thread utilde = utilde + bhat2end[i] * dt * k - end - end - - stage_limiter!(u, integrator, p, t + dt) - step_limiter!(u, integrator, p, t + dt) - - # FSAL - f(k, u, p, t + dt) - OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) - - if integrator.opts.adaptive - @.. broadcast = false thread = thread utilde = utilde + bhatfsal * dt * k - calculate_residuals!( - atmp, utilde, uprev, u, integrator.opts.abstol, - integrator.opts.reltol, integrator.opts.internalnorm, t, - thread - ) - OrdinaryDiffEqCore.set_EEst!(integrator, integrator.opts.internalnorm(atmp, t)) - end +function perform_step!(integrator, cache::LowStorageRK3SpFSALCache, repeat_step = false) + return _perform_step_iip!(integrator, cache, cache.tab) end # 2R+ low storage methods diff --git a/lib/OrdinaryDiffEqLowStorageRK/src/low_storage_tableaus.jl b/lib/OrdinaryDiffEqLowStorageRK/src/low_storage_tableaus.jl new file mode 100644 index 00000000000..b997135d3ed --- /dev/null +++ b/lib/OrdinaryDiffEqLowStorageRK/src/low_storage_tableaus.jl @@ -0,0 +1,14 @@ +@enum LowStorageRKForm TwoN TwoC + +struct LowStorageRKTableau{form, N, T, T2} <: OrdinaryDiffEqConstantCache + A2end::NTuple{N, T} + B1::T + B2end::NTuple{N, T} + c2end::NTuple{N, T2} +end + +function LowStorageRKTableau{form}( + A2end::NTuple{N, T}, B1::T, B2end::NTuple{N, T}, c2end::NTuple{N, T2} + ) where {form, N, T, T2} + return LowStorageRKTableau{form, N, T, T2}(A2end, B1, B2end, c2end) +end