diff --git a/src/integrators/integrator_utils.jl b/src/integrators/integrator_utils.jl index d62fa57a..3904fb92 100644 --- a/src/integrators/integrator_utils.jl +++ b/src/integrators/integrator_utils.jl @@ -10,18 +10,6 @@ function build_adaptive_controller_cache(alg::A, ::Type{T}) where {A, T} return beta1, beta2, qmax, qmin, gamma, qoldinit, qold end -function build_adaptive_tsit5_controller_cache(::Type{T}) where {T} - beta1 = T(7 / 50) - beta2 = T(2 / 25) - qmax = T(10.0) - qmin = T(1 / 5) - gamma = T(9 / 10) - qoldinit = T(1e-4) - qold = qoldinit - - return beta1, beta2, qmax, qmin, gamma, qoldinit, qold -end - @inline function savevalues!(integrator::DiffEqBase.AbstractODEIntegrator{AlgType, IIP, S, T }, ts, us, diff --git a/src/integrators/types.jl b/src/integrators/types.jl index aa007a1e..fe68ce1f 100644 --- a/src/integrators/types.jl +++ b/src/integrators/types.jl @@ -8,8 +8,9 @@ function Adapt.adapt_structure(to, prob::ODEProblem{<:Any, <:Any, iip}) where {i adapt(to, prob.kwargs)...) end -mutable struct GPUTsit5Integrator{IIP, S, T, ST, P, F, TS, CB} <: - DiffEqBase.AbstractODEIntegrator{GPUTsit5, IIP, S, T} +mutable struct GPUTsit5Integrator{IIP, S, T, ST, P, F, TS, CB, AlgType} <: + DiffEqBase.AbstractODEIntegrator{AlgType, IIP, S, T} + alg::AlgType f::F # eom uprev::S # previous state u::S # current state @@ -58,8 +59,9 @@ DiffEqBase.isinplace(::GPUT5I{IIP}) where {IIP} = IIP ## Adaptive TimeStep Integrator -mutable struct GPUATsit5Integrator{IIP, S, T, ST, P, F, N, TOL, Q, TS, CB} <: - DiffEqBase.AbstractODEIntegrator{GPUTsit5, IIP, S, T} +mutable struct GPUATsit5Integrator{IIP, S, T, ST, P, F, N, TOL, Q, TS, CB, AlgType} <: + DiffEqBase.AbstractODEIntegrator{AlgType, IIP, S, T} + alg::AlgType f::F # eom uprev::S # previous state u::S # current state @@ -113,8 +115,9 @@ end end ## Vern7 -mutable struct GPUV7Integrator{IIP, S, T, ST, P, F, TS, CB, TabType} <: - DiffEqBase.AbstractODEIntegrator{GPUVern7, IIP, S, T} +mutable struct GPUV7Integrator{IIP, S, T, ST, P, F, TS, CB, TabType, AlgType} <: + DiffEqBase.AbstractODEIntegrator{AlgType, IIP, S, T} + alg::AlgType f::F # eom uprev::S # previous state u::S # current state @@ -156,8 +159,9 @@ const GPUV7I = GPUV7Integrator _ode_interpolant(Θ, integrator.dt, integrator.uprev, integrator) end -mutable struct GPUAV7Integrator{IIP, S, T, ST, P, F, N, TOL, Q, TS, CB, TabType} <: - DiffEqBase.AbstractODEIntegrator{GPUVern7, IIP, S, T} +mutable struct GPUAV7Integrator{IIP, S, T, ST, P, F, N, TOL, Q, TS, CB, TabType, AlgType} <: + DiffEqBase.AbstractODEIntegrator{AlgType, IIP, S, T} + alg::AlgType f::F # eom uprev::S # previous state u::S # current state @@ -208,8 +212,9 @@ end ## Vern9 -mutable struct GPUV9Integrator{IIP, S, T, ST, P, F, TS, CB, TabType} <: - DiffEqBase.AbstractODEIntegrator{GPUVern9, IIP, S, T} +mutable struct GPUV9Integrator{IIP, S, T, ST, P, F, TS, CB, TabType, AlgType} <: + DiffEqBase.AbstractODEIntegrator{AlgType, IIP, S, T} + alg::AlgType f::F # eom uprev::S # previous state u::S # current state @@ -251,8 +256,9 @@ const GPUV9I = GPUV9Integrator _ode_interpolant(Θ, integrator.dt, integrator.uprev, integrator) end -mutable struct GPUAV9Integrator{IIP, S, T, ST, P, F, N, TOL, Q, TS, CB, TabType} <: - DiffEqBase.AbstractODEIntegrator{GPUVern9, IIP, S, T} +mutable struct GPUAV9Integrator{IIP, S, T, ST, P, F, N, TOL, Q, TS, CB, TabType, AlgType} <: + DiffEqBase.AbstractODEIntegrator{AlgType, IIP, S, T} + alg::AlgType f::F # eom uprev::S # previous state u::S # current state @@ -304,11 +310,11 @@ end ####################################################################################### # Initialization of Integrators ####################################################################################### -@inline function gputsit5_init(f::F, IIP::Bool, u0::S, t0::T, dt::T, +@inline function gputsit5_init(alg::AlgType, f::F, IIP::Bool, u0::S, t0::T, dt::T, p::P, tstops::TS, callback::CB, save_everystep::Bool, - saveat::ST) where {F, P, T, S <: AbstractArray{T}, + saveat::ST) where {AlgType, F, P, T, S <: AbstractArray{T}, TS, CB, ST} cs, as, rs = SimpleDiffEq._build_tsit5_caches(T) @@ -317,24 +323,28 @@ end vector_event_last_time = 0 last_event_error = zero(eltype(S)) - integ = GPUT5I{IIP, S, T, ST, P, F, TS, CB}(f, copy(u0), copy(u0), copy(u0), t0, t0, t0, - dt, - sign(dt), p, true, tstops, 1, callback, - save_everystep, saveat, 1, 1, - event_last_time, - vector_event_last_time, - last_event_error, - copy(u0), copy(u0), copy(u0), copy(u0), - copy(u0), - copy(u0), copy(u0), cs, as, rs, - DiffEqBase.ReturnCode.Default) + integ = GPUT5I{IIP, S, T, ST, P, F, TS, CB, AlgType}(alg, f, copy(u0), copy(u0), + copy(u0), t0, t0, t0, + dt, + sign(dt), p, true, tstops, 1, + callback, + save_everystep, saveat, 1, 1, + event_last_time, + vector_event_last_time, + last_event_error, + copy(u0), copy(u0), copy(u0), + copy(u0), + copy(u0), + copy(u0), copy(u0), cs, as, rs, + DiffEqBase.ReturnCode.Default) end -@inline function gpuatsit5_init(f::F, IIP::Bool, u0::S, t0::T, tf::T, dt::T, p::P, +@inline function gpuatsit5_init(alg::AlgType, f::F, IIP::Bool, u0::S, t0::T, tf::T, dt::T, + p::P, abstol::TOL, reltol::TOL, internalnorm::N, tstops::TS, callback::CB, - saveat::ST) where {F, P, S, T, N, TOL, TS, CB, ST} + saveat::ST) where {AlgType, F, P, S, T, N, TOL, TS, CB, ST} cs, as, btildes, rs = SimpleDiffEq._build_atsit5_caches(T) !IIP && @assert S <: SArray @@ -344,39 +354,49 @@ end vector_event_last_time = 0 last_event_error = zero(eltype(S)) - integ = GPUAT5I{IIP, S, T, ST, P, F, N, TOL, typeof(qoldinit), TS, CB}(f, copy(u0), - copy(u0), - copy(u0), t0, t0, - t0, - tf, dt, - dt, - sign(tf - t0), p, - true, tstops, 1, - callback, false, - saveat, 1, 1, - event_last_time, - vector_event_last_time, - last_event_error, - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), cs, as, - btildes, - rs, qoldinit, - abstol, - reltol, - internalnorm, - DiffEqBase.ReturnCode.Default) + integ = GPUAT5I{IIP, S, T, ST, P, F, N, TOL, typeof(qoldinit), TS, CB, AlgType}(alg, f, + copy(u0), + copy(u0), + copy(u0), + t0, t0, + t0, + tf, dt, + dt, + sign(tf - + t0), + p, + true, + tstops, + 1, + callback, + false, + saveat, + 1, 1, + event_last_time, + vector_event_last_time, + last_event_error, + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + cs, as, + btildes, + rs, + qoldinit, + abstol, + reltol, + internalnorm, + DiffEqBase.ReturnCode.Default) end -@inline function gpuvern7_init(f::F, IIP::Bool, u0::S, t0::T, dt::T, +@inline function gpuvern7_init(alg::AlgType, f::F, IIP::Bool, u0::S, t0::T, dt::T, p::P, tstops::TS, callback::CB, save_everystep::Bool, - saveat::ST) where {F, P, T, S <: AbstractArray{T}, + saveat::ST) where {AlgType, F, P, T, S <: AbstractArray{T}, TS, CB, ST} tab = Vern7Tableau(T, T) @@ -385,28 +405,35 @@ end vector_event_last_time = 0 last_event_error = zero(eltype(S)) - integ = GPUV7I{IIP, S, T, ST, P, F, TS, CB, typeof(tab)}(f, copy(u0), copy(u0), - copy(u0), - t0, t0, t0, dt, - sign(dt), p, true, tstops, 1, - callback, - save_everystep, saveat, 1, 1, - event_last_time, - vector_event_last_time, - last_event_error, - copy(u0), copy(u0), copy(u0), - copy(u0), - copy(u0), - copy(u0), copy(u0), copy(u0), - copy(u0), copy(u0), tab, - DiffEqBase.ReturnCode.Default) + integ = GPUV7I{IIP, S, T, ST, P, F, TS, CB, typeof(tab), AlgType}(alg, f, copy(u0), + copy(u0), + copy(u0), + t0, t0, t0, dt, + sign(dt), p, true, + tstops, 1, + callback, + save_everystep, + saveat, 1, 1, + event_last_time, + vector_event_last_time, + last_event_error, + copy(u0), copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), copy(u0), + copy(u0), + copy(u0), copy(u0), + tab, + DiffEqBase.ReturnCode.Default) end -@inline function gpuavern7_init(f::F, IIP::Bool, u0::S, t0::T, tf::T, dt::T, p::P, +@inline function gpuavern7_init(alg::AlgType, f::F, IIP::Bool, u0::S, t0::T, tf::T, dt::T, + p::P, abstol::TOL, reltol::TOL, internalnorm::N, tstops::TS, callback::CB, - saveat::ST) where {F, P, S, T, N, TOL, TS, CB, ST} + saveat::ST) where {AlgType, F, P, S, T, N, TOL, TS, CB, ST} !IIP && @assert S <: SArray tab = Vern7Tableau(T, T) @@ -416,53 +443,54 @@ end vector_event_last_time = 0 last_event_error = zero(eltype(S)) - integ = GPUAV7I{IIP, S, T, ST, P, F, N, TOL, typeof(qoldinit), TS, CB, typeof(tab)}(f, - copy(u0), - copy(u0), - copy(u0), - t0, - t0, - t0, - tf, - dt, - dt, - sign(tf - - t0), - p, - true, - tstops, - 1, - callback, - false, - saveat, - 1, - 1, - event_last_time, - vector_event_last_time, - last_event_error, - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - tab, - qoldinit, - abstol, - reltol, - internalnorm, - DiffEqBase.ReturnCode.Default) + integ = GPUAV7I{IIP, S, T, ST, P, F, N, TOL, typeof(qoldinit), TS, CB, typeof(tab), + AlgType}(alg, f, + copy(u0), + copy(u0), + copy(u0), + t0, + t0, + t0, + tf, + dt, + dt, + sign(tf - + t0), + p, + true, + tstops, + 1, + callback, + false, + saveat, + 1, + 1, + event_last_time, + vector_event_last_time, + last_event_error, + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + tab, + qoldinit, + abstol, + reltol, + internalnorm, + DiffEqBase.ReturnCode.Default) end -@inline function gpuvern9_init(f::F, IIP::Bool, u0::S, t0::T, dt::T, +@inline function gpuvern9_init(alg::AlgType, f::F, IIP::Bool, u0::S, t0::T, dt::T, p::P, tstops::TS, callback::CB, save_everystep::Bool, - saveat::ST) where {F, P, T, S <: AbstractArray{T}, + saveat::ST) where {AlgType, F, P, T, S <: AbstractArray{T}, TS, CB, ST} tab = Vern9Tableau(T, T) @@ -471,27 +499,34 @@ end vector_event_last_time = 0 last_event_error = zero(eltype(S)) - integ = GPUV9I{IIP, S, T, ST, P, F, TS, CB, typeof(tab)}(f, copy(u0), copy(u0), - copy(u0), - t0, t0, t0, dt, - sign(dt), p, true, tstops, 1, - callback, - save_everystep, saveat, 1, 1, - event_last_time, - vector_event_last_time, - last_event_error, - copy(u0), copy(u0), copy(u0), - copy(u0), copy(u0), copy(u0), - copy(u0), copy(u0), - copy(u0), copy(u0), tab, - DiffEqBase.ReturnCode.Default) + integ = GPUV9I{IIP, S, T, ST, P, F, TS, CB, typeof(tab), AlgType}(alg, f, copy(u0), + copy(u0), + copy(u0), + t0, t0, t0, dt, + sign(dt), p, true, + tstops, 1, + callback, + save_everystep, + saveat, 1, 1, + event_last_time, + vector_event_last_time, + last_event_error, + copy(u0), copy(u0), + copy(u0), + copy(u0), copy(u0), + copy(u0), + copy(u0), copy(u0), + copy(u0), copy(u0), + tab, + DiffEqBase.ReturnCode.Default) end -@inline function gpuavern9_init(f::F, IIP::Bool, u0::S, t0::T, tf::T, dt::T, p::P, +@inline function gpuavern9_init(alg::AlgType, f::F, IIP::Bool, u0::S, t0::T, tf::T, dt::T, + p::P, abstol::TOL, reltol::TOL, internalnorm::N, tstops::TS, callback::CB, - saveat::ST) where {F, P, S, T, N, TOL, TS, CB, ST} + saveat::ST) where {AlgType, F, P, S, T, N, TOL, TS, CB, ST} !IIP && @assert S <: SArray tab = Vern9Tableau(T, T) @@ -501,44 +536,45 @@ end vector_event_last_time = 0 last_event_error = zero(eltype(S)) - integ = GPUAV9I{IIP, S, T, ST, P, F, N, TOL, typeof(qoldinit), TS, CB, typeof(tab)}(f, - copy(u0), - copy(u0), - copy(u0), - t0, - t0, - t0, - tf, - dt, - dt, - sign(tf - - t0), - p, - true, - tstops, - 1, - callback, - false, - saveat, - 1, - 1, - event_last_time, - vector_event_last_time, - last_event_error, - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - copy(u0), - tab, - qoldinit, - abstol, - reltol, - internalnorm, - DiffEqBase.ReturnCode.Default) + integ = GPUAV9I{IIP, S, T, ST, P, F, N, TOL, typeof(qoldinit), TS, CB, typeof(tab), + AlgType}(alg, f, + copy(u0), + copy(u0), + copy(u0), + t0, + t0, + t0, + tf, + dt, + dt, + sign(tf - + t0), + p, + true, + tstops, + 1, + callback, + false, + saveat, + 1, + 1, + event_last_time, + vector_event_last_time, + last_event_error, + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + copy(u0), + tab, + qoldinit, + abstol, + reltol, + internalnorm, + DiffEqBase.ReturnCode.Default) end diff --git a/src/perform_step/gpu_rodas4_perform_step.jl b/src/perform_step/gpu_rodas4_perform_step.jl index f8003a95..5c5912e4 100644 --- a/src/perform_step/gpu_rodas4_perform_step.jl +++ b/src/perform_step/gpu_rodas4_perform_step.jl @@ -7,7 +7,7 @@ integ.uprev = integ.u uprev = integ.u @unpack a21, a31, a32, a41, a42, a43, a51, a52, a53, a54, C21, C31, C32, C41, C42, C43, - C51, C52, C53, C54, C61, C62, C63, C64, C65, gamma, c2, c3, c4, d1, d2, d3, + C51, C52, C53, C54, C61, C62, C63, C64, C65, γ, c2, c3, c4, d1, d2, d3, d4 = integ.tab integ.tprev = t @@ -57,7 +57,7 @@ dtd2 = dt * d2 dtd3 = dt * d3 dtd4 = dt * d4 - dtgamma = dt * gamma + dtgamma = dt * γ # Starting W = J - I * inv(dtgamma) @@ -188,9 +188,15 @@ end reltol = integ.reltol @unpack a21, a31, a32, a41, a42, a43, a51, a52, a53, a54, C21, C31, C32, C41, C42, C43, - C51, C52, C53, C54, C61, C62, C63, C64, C65, gamma, c2, c3, c4, d1, d2, d3, + C51, C52, C53, C54, C61, C62, C63, C64, C65, γ, c2, c3, c4, d1, d2, d3, d4 = integ.tab + @unpack h21, h22, h23, h24, h25, h31, h32, h33, h34, h35 = integ.tab + + # Jacobian + J = f.jac(uprev, p, t) + dT = f.tgrad(uprev, p, t) + if integ.u_modified k1 = f(uprev, p, t) integ.u_modified = false @@ -203,10 +209,6 @@ end while EEst > convert(T, 1.0) dt < convert(T, 1.0f-14) && error("dt