Skip to content

Commit

Permalink
Merge branch 'master' into upstream
Browse files Browse the repository at this point in the history
  • Loading branch information
Shreyas-Ekanathan authored Aug 21, 2024
2 parents 83c55c9 + 3d09fa2 commit de394f3
Show file tree
Hide file tree
Showing 9 changed files with 753 additions and 62 deletions.
3 changes: 3 additions & 0 deletions lib/OrdinaryDiffEqFIRK/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,15 @@ version = "1.1.0"
[deps]
DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
GenericLinearAlgebra = "14197337-ba66-59df-a3e3-ca00e7dcff7a"
GenericSchur = "c145ed77-6b09-5dd9-b285-bf645a82121e"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
OrdinaryDiffEqDifferentiation = "4302a76b-040a-498a-8c04-15b101fed76b"
OrdinaryDiffEqNonlinearSolve = "127b3ac7-2247-4354-8eb6-78cf4e7c58e8"
Polynomials = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961"
Expand Down
3 changes: 2 additions & 1 deletion lib/OrdinaryDiffEqFIRK/src/OrdinaryDiffEqFIRK.jl
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import OrdinaryDiffEqCore: alg_order, calculate_residuals!,
get_current_adaptive_order, get_fsalfirstlast,
isfirk
using MuladdMacro, DiffEqBase, RecursiveArrayTools
using Polynomials, GenericLinearAlgebra, GenericSchur
using SciMLOperators: AbstractSciMLOperator
using LinearAlgebra: I, UniformScaling, mul!, lu
import LinearSolve
Expand All @@ -42,6 +43,6 @@ include("firk_tableaus.jl")
include("firk_perform_step.jl")
include("integrator_interface.jl")

export RadauIIA3, RadauIIA5, RadauIIA9
export RadauIIA3, RadauIIA5, RadauIIA9, AdaptiveRadau

end
2 changes: 2 additions & 0 deletions lib/OrdinaryDiffEqFIRK/src/alg_utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,12 @@ qmax_default(alg::Union{RadauIIA3, RadauIIA5, RadauIIA9}) = 8
alg_order(alg::RadauIIA3) = 3
alg_order(alg::RadauIIA5) = 5
alg_order(alg::RadauIIA9) = 9
alg_order(alg::AdaptiveRadau) = 9

isfirk(alg::RadauIIA3) = true
isfirk(alg::RadauIIA5) = true
isfirk(alg::RadauIIA9) = true
isfirk(alg::AdaptiveRadau) = true

alg_adaptive_order(alg::RadauIIA3) = 1
alg_adaptive_order(alg::RadauIIA5) = 3
Expand Down
39 changes: 39 additions & 0 deletions lib/OrdinaryDiffEqFIRK/src/algorithms.jl
Original file line number Diff line number Diff line change
Expand Up @@ -153,3 +153,42 @@ function RadauIIA9(; chunk_size = Val{0}(), autodiff = Val{true}(),
controller,
step_limiter!)
end

struct AdaptiveRadau{CS, AD, F, P, FDT, ST, CJ, Tol, C1, C2, StepLimiter} <:
OrdinaryDiffEqNewtonAdaptiveAlgorithm{CS, AD, FDT, ST, CJ}
linsolve::F
precs::P
smooth_est::Bool
extrapolant::Symbol
κ::Tol
maxiters::Int
fast_convergence_cutoff::C1
new_W_γdt_cutoff::C2
controller::Symbol
step_limiter!::StepLimiter
num_stages::Int
end

function AdaptiveRadau(; chunk_size = Val{0}(), autodiff = Val{true}(),
standardtag = Val{true}(), concrete_jac = nothing,
diff_type = Val{:forward}, num_stages = 3,
linsolve = nothing, precs = DEFAULT_PRECS,
extrapolant = :dense, fast_convergence_cutoff = 1 // 5,
new_W_γdt_cutoff = 1 // 5,
controller = :Predictive, κ = nothing, maxiters = 10, smooth_est = true,
step_limiter! = trivial_limiter!)
AdaptiveRadau{_unwrap_val(chunk_size), _unwrap_val(autodiff), typeof(linsolve),
typeof(precs), diff_type, _unwrap_val(standardtag), _unwrap_val(concrete_jac),
typeof(κ), typeof(fast_convergence_cutoff),
typeof(new_W_γdt_cutoff), typeof(step_limiter!)}(linsolve,
precs,
smooth_est,
extrapolant,
κ,
maxiters,
fast_convergence_cutoff,
new_W_γdt_cutoff,
controller,
step_limiter!, num_stages)
end

161 changes: 161 additions & 0 deletions lib/OrdinaryDiffEqFIRK/src/firk_caches.jl
Original file line number Diff line number Diff line change
Expand Up @@ -470,3 +470,164 @@ function alg_cache(alg::RadauIIA9, u, rate_prototype, ::Type{uEltypeNoUnits},
linsolve1, linsolve2, linsolve3, rtol, atol, dt, dt,
Convergence, alg.step_limiter!)
end

mutable struct AdaptiveRadauConstantCache{F, Tab, Tol, Dt, U, JType} <:
OrdinaryDiffEqConstantCache
uf::F
tab::Tab
κ::Tol
ηold::Tol
iter::Int
cont::AbstractVector{U}
dtprev::Dt
W_γdt::Dt
status::NLStatus
J::JType
end

function alg_cache(alg::AdaptiveRadau, u, rate_prototype, ::Type{uEltypeNoUnits},
::Type{uBottomEltypeNoUnits},
::Type{tTypeNoUnits}, uprev, uprev2, f, t, dt, reltol, p, calck,
::Val{false}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
uf = UDerivativeWrapper(f, t, p)
uToltype = constvalue(uBottomEltypeNoUnits)
num_stages = alg.num_stages
tab = adaptiveRadauTableau(uToltype, constvalue(tTypeNoUnits), num_stages)

cont = Vector{typeof(u)}(undef, num_stages - 1)
for i in 1: (num_stages - 1)
cont[i] = zero(u)
end

κ = alg.κ !== nothing ? convert(uToltype, alg.κ) : convert(uToltype, 1 // 100)
J = false .* _vec(rate_prototype) .* _vec(rate_prototype)'

AdaptiveRadauConstantCache(uf, tab, κ, one(uToltype), 10000, cont, dt, dt,
Convergence, J)
end

mutable struct AdaptiveRadauCache{uType, cuType, uNoUnitsType, rateType, JType, W1Type, W2Type,
UF, JC, F1, F2, Tab, Tol, Dt, rTol, aTol, StepLimiter} <:
OrdinaryDiffEqMutableCache
u::uType
uprev::uType
z::AbstractVector{uType}
w::AbstractVector{uType}
dw1::uType
ubuff::uType
dw2::AbstractVector{cuType}
cubuff::AbstractVector{cuType}
cont::AbstractVector{uType}
du1::rateType
fsalfirst::rateType
k::AbstractVector{rateType}
fw::AbstractVector{rateType}
J::JType
W1::W1Type #real
W2::AbstractVector{W2Type} #complex
uf::UF
tab::Tab
κ::Tol
ηold::Tol
iter::Int
tmp::AbstractVector{uType}
atmp::uNoUnitsType
jac_config::JC
linsolve1::F1 #real
linsolve2::AbstractVector{F2} #complex
rtol::rTol
atol::aTol
dtprev::Dt
W_γdt::Dt
status::NLStatus
step_limiter!::StepLimiter
end

function alg_cache(alg::AdaptiveRadau, u, rate_prototype, ::Type{uEltypeNoUnits},
::Type{uBottomEltypeNoUnits},
::Type{tTypeNoUnits}, uprev, uprev2, f, t, dt, reltol, p, calck,
::Val{true}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits}
uf = UJacobianWrapper(f, t, p)
uToltype = constvalue(uBottomEltypeNoUnits)
alg.num_stages = num_stages
tab = adaptiveRadauTableau(uToltype, constvalue(tTypeNoUnits), num_stages)

κ = alg.κ !== nothing ? convert(uToltype, alg.κ) : convert(uToltype, 1 // 100)

z = Vector{typeof(u)}(undef, num_stages)
w = Vector{typeof(u)}(undef, num_stages)
for i in 1:s
z[i] = w[i] = zero(u)
end

dw1 = zero(u)
ubuff = zero(u)
dw2 = Vector{typeof(u)}(undef, floor(Int, num_stages/2))
for i in 1 : floor(Int, num_stages/2)
dw2[i] = similar(u, Complex{eltype(u)})
recursivefill!(dw[i], false)
end
cubuff = Vector{typeof(u)}(undef, floor(Int, num_stages/2))
for i in 1 :floor(Int, num_stages/2)
cubuff[i] = similar(u, Complex{eltype(u)})
recursivefill!(cubuff[i], false)
end

cont = Vector{typeof(u)}(undef, num_stages - 1)
for i in 1: (num_stages - 1)
cont[i] = zero(u)
end

fsalfirst = zero(rate_prototype)
fw = Vector{typeof(rate_prototype)}(undef, num_stages)
k = Vector{typeof(rate_prototype)}(undef, num_stages)
for i in 1: num_stages
k[i] = fw[i] = zero(rate_prototype)
end

J, W1 = build_J_W(alg, u, uprev, p, t, dt, f, uEltypeNoUnits, Val(true))
if J isa AbstractSciMLOperator
error("Non-concrete Jacobian not yet supported by RadauIIA5.")
end
W2 = vector{typeof(Complex{W1})}(undef, floor(Int, num_stages/2))
for i in 1 : floor(Int, num_stages/2)
W2[i] = similar(J, Complex{eltype(W1)})
recursivefill!(w2[i], false)
end

du1 = zero(rate_prototype)

tmp = Vector{typeof(u)}(undef, binomial(num_stages,2))
for i in 1 : binomial(num_stages,2)
tmp[i] = zero(u)
end

atmp = similar(u, uEltypeNoUnits)
recursivefill!(atmp, false)

jac_config = build_jac_config(alg, f, uf, du1, uprev, u, tmp, dw1)

linprob = LinearProblem(W1, _vec(ubuff); u0 = _vec(dw1))
linsolve1 = init(linprob, alg.linsolve, alias_A = true, alias_b = true,
assumptions = LinearSolve.OperatorAssumptions(true))

linsolve2 = Vector{typeof(linsolve1)}(undef, floor(Int, num_stages/2))
for i in 1 : floor(int, num_stages/2)
linprob = LinearProblem(W2[i], _vec(cubuff[i]); u0 = _vec(dw2[i]))
linsolve2 = init(linprob, alg.linsolve, alias_A = true, alias_b = true,
assumptions = LinearSolve.OperatorAssumptions(true))
end

rtol = reltol isa Number ? reltol : zero(reltol)
atol = reltol isa Number ? reltol : zero(reltol)

AdaptiveRadauCache(u, uprev,
z, w, dw1, ubuff, dw2, cubuff, cont,
du1, fsalfirst, k, fw,
J, W1, W2,
uf, tab, κ, one(uToltype), 10000,
tmp, atmp, jac_config,
linsolve1, linsolve2, rtol, atol, dt, dt,
Convergence, alg.step_limiter!)
end

Loading

0 comments on commit de394f3

Please sign in to comment.