diff --git a/lib/OrdinaryDiffEqExponentialRK/Project.toml b/lib/OrdinaryDiffEqExponentialRK/Project.toml index 1ef7f583862..3b6a13397aa 100644 --- a/lib/OrdinaryDiffEqExponentialRK/Project.toml +++ b/lib/OrdinaryDiffEqExponentialRK/Project.toml @@ -1,7 +1,7 @@ name = "OrdinaryDiffEqExponentialRK" uuid = "e0540318-69ee-4070-8777-9e2de6de23de" authors = ["ParamThakkar123 "] -version = "2.1.0" +version = "2.1.1" [deps] FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" diff --git a/lib/OrdinaryDiffEqExponentialRK/src/exponential_rk_caches.jl b/lib/OrdinaryDiffEqExponentialRK/src/exponential_rk_caches.jl index 359819a1664..e7a68e6c284 100644 --- a/lib/OrdinaryDiffEqExponentialRK/src/exponential_rk_caches.jl +++ b/lib/OrdinaryDiffEqExponentialRK/src/exponential_rk_caches.jl @@ -443,7 +443,7 @@ for (Alg, Cache) in [ end end -@cache struct Exp4Cache{uType, rateType, JCType, FType, matType, JType, KsType} <: +@cache struct Exp4Cache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <: ExpRKCache u::uType uprev::uType @@ -458,6 +458,7 @@ end J::JType B::matType KsCache::KsType + ts::tsType end function alg_cache( alg::Exp4, u, rate_prototype, ::Type{uEltypeNoUnits}, @@ -489,10 +490,11 @@ function alg_cache( # Allocate caches for phiv_timestep maxiter = min(alg.m, n) KsCache = _phiv_timestep_caches(u, maxiter, 1) - return Exp4Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache) + ts = Vector{typeof(dt)}(undef, 3) + return Exp4Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache, ts) end -@cache struct EPIRK4s3ACache{uType, rateType, JCType, FType, matType, JType, KsType} <: +@cache struct EPIRK4s3ACache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <: ExpRKCache u::uType uprev::uType @@ -507,6 +509,7 @@ end J::JType B::matType KsCache::KsType + ts::tsType end function alg_cache( alg::EPIRK4s3A, u, rate_prototype, ::Type{uEltypeNoUnits}, @@ -537,10 +540,11 @@ function alg_cache( # Allocate caches for phiv_timestep maxiter = min(alg.m, n) KsCache = _phiv_timestep_caches(u, maxiter, 4) - return EPIRK4s3ACache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache) + ts = Vector{typeof(dt)}(undef, 2) + return EPIRK4s3ACache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache, ts) end -@cache struct EPIRK4s3BCache{uType, rateType, JCType, FType, matType, JType, KsType} <: +@cache struct EPIRK4s3BCache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <: ExpRKCache u::uType uprev::uType @@ -555,6 +559,7 @@ end J::JType B::matType KsCache::KsType + ts::tsType end function alg_cache( alg::EPIRK4s3B, u, rate_prototype, ::Type{uEltypeNoUnits}, @@ -585,7 +590,8 @@ function alg_cache( # Allocate caches for phiv_timestep maxiter = min(alg.m, n) KsCache = _phiv_timestep_caches(u, maxiter, 4) - return EPIRK4s3BCache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache) + ts = Vector{typeof(dt)}(undef, 2) + return EPIRK4s3BCache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache, ts) end @cache struct EPIRK5s3Cache{uType, rateType, JCType, FType, matType, JType, KsType} <: @@ -635,7 +641,7 @@ function alg_cache( return EPIRK5s3Cache(u, uprev, tmp, dz, k, rtmp, rtmp2, du1, jac_config, uf, J, B, KsCache) end -@cache struct EXPRB53s3Cache{uType, rateType, JCType, FType, matType, JType, KsType} <: +@cache struct EXPRB53s3Cache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <: ExpRKCache u::uType uprev::uType @@ -650,6 +656,7 @@ end J::JType B::matType KsCache::KsType + ts::tsType end function alg_cache( alg::EXPRB53s3, u, rate_prototype, ::Type{uEltypeNoUnits}, @@ -680,10 +687,11 @@ function alg_cache( # Allocate caches for phiv_timestep maxiter = min(alg.m, n) KsCache = _phiv_timestep_caches(u, maxiter, 4) - return EXPRB53s3Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache) + ts = Vector{typeof(dt)}(undef, 2) + return EXPRB53s3Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache, ts) end -@cache struct EPIRK5P1Cache{uType, rateType, JCType, FType, matType, JType, KsType} <: +@cache struct EPIRK5P1Cache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <: ExpRKCache u::uType uprev::uType @@ -698,6 +706,7 @@ end J::JType B::matType KsCache::KsType + ts::tsType end function alg_cache( alg::EPIRK5P1, u, rate_prototype, ::Type{uEltypeNoUnits}, @@ -728,10 +737,11 @@ function alg_cache( # Allocate caches for phiv_timestep maxiter = min(alg.m, n) KsCache = _phiv_timestep_caches(u, maxiter, 3) - return EPIRK5P1Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache) + ts = Vector{typeof(dt)}(undef, 3) + return EPIRK5P1Cache(u, uprev, tmp, dz, rtmp, rtmp2, du1, jac_config, uf, K, J, B, KsCache, ts) end -@cache struct EPIRK5P2Cache{uType, rateType, JCType, FType, matType, JType, KsType} <: +@cache struct EPIRK5P2Cache{uType, rateType, JCType, FType, matType, JType, KsType, tsType} <: ExpRKCache u::uType uprev::uType @@ -747,6 +757,7 @@ end J::JType B::matType KsCache::KsType + ts::tsType end function alg_cache( alg::EPIRK5P2, u, rate_prototype, ::Type{uEltypeNoUnits}, @@ -777,7 +788,8 @@ function alg_cache( # Allocate caches for phiv_timestep maxiter = min(alg.m, n) KsCache = _phiv_timestep_caches(u, maxiter, 3) - return EPIRK5P2Cache(u, uprev, tmp, dz, rtmp, rtmp2, dR, du1, jac_config, uf, K, J, B, KsCache) + ts = Vector{typeof(dt)}(undef, 3) + return EPIRK5P2Cache(u, uprev, tmp, dz, rtmp, rtmp2, dR, du1, jac_config, uf, K, J, B, KsCache, ts) end #################################### diff --git a/lib/OrdinaryDiffEqExponentialRK/src/exponential_rk_perform_step.jl b/lib/OrdinaryDiffEqExponentialRK/src/exponential_rk_perform_step.jl index d6ce5f3418c..3f3b0d01969 100644 --- a/lib/OrdinaryDiffEqExponentialRK/src/exponential_rk_perform_step.jl +++ b/lib/OrdinaryDiffEqExponentialRK/src/exponential_rk_perform_step.jl @@ -801,11 +801,13 @@ end function perform_step!(integrator, cache::Exp4Cache, repeat_step = false) (; t, dt, uprev, u, f, p) = integrator - (; tmp, rtmp, rtmp2, K, J, B, KsCache) = cache + (; tmp, rtmp, rtmp2, K, J, B, KsCache, ts) = cache calc_J!(J, integrator, cache) alg = unwrap_alg(integrator, true) f0 = integrator.fsalfirst # f(u0) is fsaled - ts = [dt / 3, 2dt / 3, dt] + ts[1] = dt / 3 + ts[2] = 2dt / 3 + ts[3] = dt kwargs = ( tol = integrator.opts.reltol, iop = alg.iop, opnorm = integrator.opts.internalopnorm, @@ -819,21 +821,23 @@ function perform_step!(integrator, cache::Exp4Cache, repeat_step = false) @inbounds for i in 1:3 K[:, i] ./= ts[i] end - mul!(rtmp, K, [-7 / 300, 97 / 150, -37 / 300]) # rtmp is now w4 + @views @.. broadcast = false rtmp = (-7 / 300) * K[:, 1] + (97 / 150) * K[:, 2] + + (-37 / 300) * K[:, 3] # rtmp is now w4 @muladd @.. broadcast = false tmp = uprev + dt * rtmp # tmp is now u4 mul!(rtmp2, J, rtmp) f(rtmp, tmp, p, t + dt) # TODO: what should be the time? OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) @muladd @.. broadcast = false @view(B[:, 2]) = rtmp - f0 - dt * rtmp2 # B[:,2] is now d4 # Partially update entities that use k1, k2, k3 - mul!(rtmp, K, [59 / 300, -7 / 75, 269 / 300]) # rtmp is now w7 + @views @.. broadcast = false rtmp = (59 / 300) * K[:, 1] + (-7 / 75) * K[:, 2] + + (269 / 300) * K[:, 3] # rtmp is now w7 @muladd @.. broadcast = false u = uprev + dt * @view(K[:, 3]) # Krylov for the first remainder d4 phiv_timestep!(K, ts, J, B; kwargs...) @inbounds for i in 1:3 K[:, i] ./= ts[i] end - mul!(rtmp2, K, [2 / 3, 2 / 3, 2 / 3]) + @views @.. broadcast = false rtmp2 = (2 / 3) * (K[:, 1] + K[:, 2] + K[:, 3]) rtmp .+= rtmp2 # w7 fully updated @muladd @.. broadcast = false tmp = uprev + dt * rtmp # tmp is now u7 mul!(rtmp2, J, rtmp) @@ -841,7 +845,7 @@ function perform_step!(integrator, cache::Exp4Cache, repeat_step = false) OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) @muladd @.. broadcast = false @view(B[:, 2]) = rtmp - f0 - dt * rtmp2 # B[:,2] is now d7 # Partially update entities that use k4, k5, k6 - mul!(rtmp, K, [1.0, -4 / 3, 1.0]) + @views @.. broadcast = false rtmp = K[:, 1] + (-4 / 3) * K[:, 2] + K[:, 3] axpy!(dt, rtmp, u) # Krylov for the second remainder d7 k7 = @view(K[:, 1]) @@ -892,7 +896,7 @@ end function perform_step!(integrator, cache::EPIRK4s3ACache, repeat_step = false) (; t, dt, uprev, u, f, p) = integrator - (; tmp, rtmp, rtmp2, K, J, B, KsCache) = cache + (; tmp, rtmp, rtmp2, K, J, B, KsCache, ts) = cache calc_J!(J, integrator, cache) alg = unwrap_alg(integrator, true) f0 = integrator.fsalfirst # f(u0) is fsaled @@ -905,21 +909,23 @@ function perform_step!(integrator, cache::EPIRK4s3ACache, repeat_step = false) # Compute U2 and U3 vertically B[:, 2] .= f0 - phiv_timestep!(K, [dt / 2, 2dt / 3], J, @view(B[:, 1:2]); kwargs...) + ts[1] = dt / 2 + ts[2] = 2dt / 3 + phiv_timestep!(K, ts, J, @view(B[:, 1:2]); kwargs...) ## U2 and R2 @.. broadcast = false tmp = uprev + @view(K[:, 1]) # tmp is now U2 f(rtmp, tmp, p, t + dt / 2) mul!(rtmp2, J, @view(K[:, 1])) @.. broadcast = false rtmp = rtmp - f0 - rtmp2 # rtmp is now R2 - B[:, 4] .= (32 / dt^2) * rtmp - B[:, 5] .= (-144 / dt^3) * rtmp + B[:, 4] .= (32 / dt^2) .* rtmp + B[:, 5] .= (-144 / dt^3) .* rtmp ## U3 and R3 @.. broadcast = false tmp = uprev + @view(K[:, 2]) # tmp is now U3 f(rtmp, tmp, p, t + 2dt / 3) mul!(rtmp2, J, @view(K[:, 2])) @.. broadcast = false rtmp = rtmp - f0 - rtmp2 # rtmp is now R3 - B[:, 4] .-= (13.5 / dt^2) * rtmp - B[:, 5] .+= (81 / dt^3) * rtmp + B[:, 4] .-= (13.5 / dt^2) .* rtmp + B[:, 5] .+= (81 / dt^3) .* rtmp OrdinaryDiffEqCore.increment_nf!(integrator.stats, 2) # Update u @@ -972,7 +978,7 @@ end function perform_step!(integrator, cache::EPIRK4s3BCache, repeat_step = false) (; t, dt, uprev, u, f, p) = integrator - (; tmp, rtmp, rtmp2, K, J, B, KsCache) = cache + (; tmp, rtmp, rtmp2, K, J, B, KsCache, ts) = cache calc_J!(J, integrator, cache) alg = unwrap_alg(integrator, true) f0 = integrator.fsalfirst # f(u0) is fsaled @@ -986,7 +992,9 @@ function perform_step!(integrator, cache::EPIRK4s3BCache, repeat_step = false) # Compute U2 and U3 vertically fill!(@view(B[:, 2]), zero(eltype(B))) B[:, 3] .= f0 - phiv_timestep!(K, [dt / 2, 3dt / 4], J, @view(B[:, 1:3]); kwargs...) + ts[1] = dt / 2 + ts[2] = 3dt / 4 + phiv_timestep!(K, ts, J, @view(B[:, 1:3]); kwargs...) K[:, 1] .*= 8 / (3 * dt) K[:, 2] .*= 16 / (9 * dt) ## U2 and R2 @@ -994,15 +1002,15 @@ function perform_step!(integrator, cache::EPIRK4s3BCache, repeat_step = false) f(rtmp, tmp, p, t + dt / 2) mul!(rtmp2, J, @view(K[:, 1])) @.. broadcast = false rtmp = rtmp - f0 - rtmp2 # rtmp is now R2 - B[:, 4] .= (54 / dt^2) * rtmp - B[:, 5] .= (-324 / dt^3) * rtmp + B[:, 4] .= (54 / dt^2) .* rtmp + B[:, 5] .= (-324 / dt^3) .* rtmp ## U3 and R3 @.. broadcast = false tmp = uprev + @view(K[:, 2]) # tmp is now U3 f(rtmp, tmp, p, t + 3dt / 4) mul!(rtmp2, J, @view(K[:, 2])) @.. broadcast = false rtmp = rtmp - f0 - rtmp2 # rtmp is now R3 - B[:, 4] .-= (16 / dt^2) * rtmp - B[:, 5] .+= (144 / dt^3) * rtmp + B[:, 4] .-= (16 / dt^2) .* rtmp + B[:, 5] .+= (144 / dt^3) .* rtmp OrdinaryDiffEqCore.increment_nf!(integrator.stats, 2) # Update u @@ -1091,7 +1099,7 @@ function perform_step!(integrator, cache::EPIRK5s3Cache, repeat_step = false) # Compute U3 horizontally B[:, 2] .= (53 / 5) .* f0 B[:, 3] .= (-648 / (5 * dt)) .* f0 - B[:, 4] .= (2916 / (5 * dt^2)) .* f0 + (32065 / (1152 * dt^2)) .* rtmp + B[:, 4] .= (2916 / (5 * dt^2)) .* f0 .+ (32065 / (1152 * dt^2)) .* rtmp phiv_timestep!(k, 4dt / 9, J, @view(B[:, 1:4]); kwargs...) ## Update B matrix using R2 B[:, 2] .= f0 @@ -1161,7 +1169,7 @@ end function perform_step!(integrator, cache::EXPRB53s3Cache, repeat_step = false) (; t, dt, uprev, u, f, p) = integrator - (; tmp, rtmp, rtmp2, K, J, B, KsCache) = cache + (; tmp, rtmp, rtmp2, K, J, B, KsCache, ts) = cache calc_J!(J, integrator, cache) alg = unwrap_alg(integrator, true) f0 = integrator.fsalfirst # f(u0) is fsaled @@ -1174,7 +1182,9 @@ function perform_step!(integrator, cache::EXPRB53s3Cache, repeat_step = false) # Compute the first group for U2 and U3 B[:, 2] .= f0 - phiv_timestep!(K, [dt / 2, 9dt / 10], J, @view(B[:, 1:2]); kwargs...) + ts[1] = dt / 2 + ts[2] = 9dt / 10 + phiv_timestep!(K, ts, J, @view(B[:, 1:2]); kwargs...) ## U2 and R2 @.. broadcast = false tmp = uprev + @view(K[:, 1]) # tmp is now U2 f(rtmp, tmp, p, t + dt / 2) @@ -1186,21 +1196,21 @@ function perform_step!(integrator, cache::EXPRB53s3Cache, repeat_step = false) # Compute the second group for U3 fill!(@view(B[:, 2]), zero(eltype(B))) B[:, 4] .= rtmp - phiv_timestep!(K, [dt / 2, 9dt / 10], J, @view(B[:, 1:4]); kwargs...) + phiv_timestep!(K, ts, J, @view(B[:, 1:4]); kwargs...) ## Update B using R2 B[:, 2] .= f0 B[:, 4] .= (18 / dt^2) .* rtmp B[:, 5] .= (-60 / dt^3) .* rtmp ## U3 and R3 - @views tmp .+= 216 / (25 * dt^2) .* K[:, 1] + 8 / dt^2 .* K[:, 2] # tmp is now U3 + @views tmp .+= 216 / (25 * dt^2) .* K[:, 1] .+ 8 / dt^2 .* K[:, 2] # tmp is now U3 f(rtmp, tmp, p, t + 9dt / 10) OrdinaryDiffEqCore.increment_nf!(integrator.stats, 1) tmp .-= uprev mul!(rtmp2, J, tmp) @.. broadcast = false rtmp = rtmp - f0 - rtmp2 # rtmp is now R3 ## Update B using R3 - B[:, 4] .-= (250 / (81 * dt^2)) * rtmp - B[:, 5] .+= (500 / (27 * dt^3)) * rtmp + B[:, 4] .-= (250 / (81 * dt^2)) .* rtmp + B[:, 5] .+= (500 / (27 * dt^3)) .* rtmp # Update u du = @view(K[:, 1]) @@ -1268,7 +1278,7 @@ end function perform_step!(integrator, cache::EPIRK5P1Cache, repeat_step = false) (; t, dt, uprev, u, f, p) = integrator - (; tmp, rtmp, rtmp2, K, J, B, KsCache) = cache + (; tmp, rtmp, rtmp2, K, J, B, KsCache, ts) = cache calc_J!(J, integrator, cache) alg = unwrap_alg(integrator, true) f0 = integrator.fsalfirst # f(u0) is fsaled @@ -1291,7 +1301,10 @@ function perform_step!(integrator, cache::EPIRK5P1Cache, repeat_step = false) # Compute the first column (f0) B[:, 2] .= f0 - phiv_timestep!(K, [g11, g21, g31], J, @view(B[:, 1:2]); kwargs...) + ts[1] = g11 + ts[2] = g21 + ts[3] = g31 + phiv_timestep!(K, ts, J, @view(B[:, 1:2]); kwargs...) ## U1 and R1 @.. broadcast = false tmp = uprev + @view(K[:, 1]) # tmp is now U1 f(rtmp, tmp, p, t + g11) @@ -1384,7 +1397,7 @@ end function perform_step!(integrator, cache::EPIRK5P2Cache, repeat_step = false) (; t, dt, uprev, u, f, p) = integrator - (; tmp, rtmp, rtmp2, dR, K, J, B, KsCache) = cache + (; tmp, rtmp, rtmp2, dR, K, J, B, KsCache, ts) = cache calc_J!(J, integrator, cache) alg = unwrap_alg(integrator, true) f0 = integrator.fsalfirst # f(u0) is fsaled @@ -1409,7 +1422,10 @@ function perform_step!(integrator, cache::EPIRK5P2Cache, repeat_step = false) # Compute the first column (f0) B[:, 2] .= f0 - phiv_timestep!(K, [g11, g21, g31], J, @view(B[:, 1:2]); kwargs...) + ts[1] = g11 + ts[2] = g21 + ts[3] = g31 + phiv_timestep!(K, ts, J, @view(B[:, 1:2]); kwargs...) ## U1 and R1 @.. broadcast = false tmp = uprev + @view(K[:, 1]) # tmp is now U1 f(rtmp, tmp, p, t + g11)