Add Lobatto IIIa-c and Radau IIa solvers

src/BoundaryValueDiffEq.jl

@@ -30,9 +30,12 @@ include("algorithms.jl")
 include("alg_utils.jl")
 
 include("mirk_tableaus.jl")
+include("lobatto_tableaus.jl")
+include("radau_tableaus.jl")
 
 include("solve/single_shooting.jl")
 include("solve/multiple_shooting.jl")
+include("solve/firk.jl")
 include("solve/mirk.jl")
 
 include("collocation.jl")
@@ -273,6 +276,10 @@ export Shooting, MultipleShooting
 export MIRK2, MIRK3, MIRK4, MIRK5, MIRK6
 export BVPM2, BVPSOL, COLNEW # From ODEInterface.jl
 
+export RadauIIa1, RadauIIa2, RadauIIa3, RadauIIa5, RadauIIa7
+export LobattoIIIa2, LobattoIIIa3, LobattoIIIa4, LobattoIIIa5
+export LobattoIIIb2, LobattoIIIb3, LobattoIIIb4, LobattoIIIb5
+export LobattoIIIc2, LobattoIIIc3, LobattoIIIc4, LobattoIIIc5
 export MIRKJacobianComputationAlgorithm, BVPJacobianAlgorithm
 
 end

src/adaptivity.jl

@@ -12,6 +12,129 @@ After we construct an interpolant, we use interp_eval to evaluate it.
     return y
 end
 
+@views function interp_eval!(y::AbstractArray, cache::FIRKCacheExpand{iip}, t, mesh, mesh_dt) where {iip}
+    i = findfirst(x -> x == y, cache.y₀.u)
+    interp_eval!(cache.y₀.u, i, cache::FIRKCacheExpand{iip}, t, mesh, mesh_dt)
+    return y
+end
+
+@views function interp_eval!(y::AbstractArray, i::Int, cache::FIRKCacheExpand{iip}, t, mesh, mesh_dt) where {iip}
+    j = interval(mesh, t)
+    h = mesh_dt[j]
+    lf = (length(cache.y₀) - 1) / (length(cache.y) - 1) # Cache length factor. We use a h corresponding to cache.y. Note that this assumes equidistributed mesh
+    if lf > 1
+        h *= lf
+    end
+    τ = (t - mesh[j]) / h
+
+    (; f, M, p, ITU, TU) = cache
+    (; c, a, b) = TU
+    (; q_coeff, stage) = ITU
+
+    K = zeros(eltype(cache.y[1].du), M, stage)
+
+    ctr_y0 = (i - 1) * (ITU.stage + 1) + 1
+    ctr_y = (j - 1) * (ITU.stage + 1) + 1
+
+    yᵢ = cache.y[ctr_y].du
+    yᵢ₊₁ = cache.y[ctr_y + ITU.stage + 1].du
+
+    if iip
+        dyᵢ = copy(yᵢ)
+        dyᵢ₊₁ = copy(yᵢ₊₁)
+
+        f(dyᵢ, yᵢ, cache.p, mesh[j])
+        f(dyᵢ₊₁, yᵢ₊₁, cache.p, mesh[j + 1])
+    else
+        dyᵢ = f(yᵢ, cache.p, mesh[j])
+        dyᵢ₊₁ = f(yᵢ₊₁, cache.p, mesh[j + 1])
+    end
+
+    # Load interpolation residual
+    for jj in 1:stage
+        K[:, jj] = cache.y[ctr_y + jj].du
+    end
+
+    z₁, z₁′ = eval_q(yᵢ, 0.5, h, q_coeff, K) # Evaluate q(x) at midpoints
+    S_coeffs = get_S_coeffs(h, yᵢ, yᵢ₊₁, z₁, dyᵢ, dyᵢ₊₁, z₁′)
+
+    #FIXME: need a better way to fix this
+    eltype(y) <: AbstractArray ? (y[ctr_y0] = S_interpolate(τ * h, S_coeffs)) : (y .= S_interpolate(τ * h, S_coeffs))
+    if ctr_y0 > length(y)
+        for (k, ci) in enumerate(c)
+            eltype(y) <: AbstractArray ? (y[ctr_y0 + k] = dS_interpolate(τ * h + (1 - τ * h) * ci, S_coeffs)) : (y = dS_interpolate(τ * h + (1 - τ * h) * ci, S_coeffs))
+        end
+    end
+
+    return eltype(y) <: AbstractArray ? y[ctr_y0] : y
+end
+
+@views function interp_eval!(y::AbstractArray, cache::FIRKCacheNested{iip}, t, mesh, mesh_dt) where {iip}
+    j = interval(mesh, t)
+    h = mesh_dt[j]
+    lf = (length(cache.y₀) - 1) / (length(cache.y) - 1) # Cache length factor. We use a h corresponding to cache.y. Note that this assumes equidistributed mesh
+    if lf > 1
+        h *= lf
+    end
+    τ = (t - mesh[j]) / h
+
+    (; f, M, p, k_discrete, ITU, TU, nest_cache, p_nestprob, prob) = cache
+    (; c, a, b) = TU
+    (; q_coeff, stage) = ITU
+
+    yᵢ = copy(cache.y[j].du)
+    yᵢ₊₁ = copy(cache.y[j + 1].du)
+
+    if iip
+        dyᵢ = copy(yᵢ)
+        dyᵢ₊₁ = copy(yᵢ₊₁)
+
+        f(dyᵢ, yᵢ, cache.p, mesh[j])
+        f(dyᵢ₊₁, yᵢ₊₁, cache.p, mesh[j + 1])
+    else
+        dyᵢ = f(yᵢ, cache.p, mesh[j])
+        dyᵢ₊₁ = f(yᵢ₊₁, cache.p, mesh[j + 1])
+    end
+
+    # Load interpolation residual
+    y_i = eltype(yᵢ) == Float64 ? yᵢ : [y.value for y in yᵢ]
+
+    p_nestprob[1:2] .= promote(mesh[j], mesh_dt[j], one(eltype(y_i)))[1:2]
+    p_nestprob[3:end] .= y_i
+
+    solve_cache!(nest_cache, k_discrete[j].du, p_nestprob)
+    K = nest_cache.u
+
+    z₁, z₁′ = eval_q(yᵢ, 0.5, h, q_coeff, K) # Evaluate q(x) at midpoints
+    S_coeffs = get_S_coeffs(h, yᵢ, yᵢ₊₁, z₁, dyᵢ, dyᵢ₊₁, z₁′)
+
+    y .= S_interpolate(τ * h, S_coeffs)
+
+    return y
+end
+
+function get_S_coeffs(h, yᵢ, yᵢ₊₁, dyᵢ, dyᵢ₊₁, ymid, dymid)
+    vals = vcat(yᵢ, yᵢ₊₁, dyᵢ, dyᵢ₊₁, ymid, dymid)
+    M = length(yᵢ)
+    A = s_constraints(M, h)
+    coeffs = reshape(A \ vals, 6, M)'
+    return coeffs
+end
+
+# S forward Interpolation
+function S_interpolate(t, coeffs)
+    ts = [t^(i - 1) for i in axes(coeffs, 2)]
+    return coeffs * ts
+end
+
+function dS_interpolate(t, S_coeffs)
+    ts = zeros(size(S_coeffs, 2))
+    for i in 2:size(S_coeffs, 2)
+        ts[i] = (i - 1) * t^(i - 2)
+    end
+    return S_coeffs * ts
+end
+
 """
     interval(mesh, t)
 
@@ -26,7 +149,7 @@ end
 
 Generate new mesh based on the defect.
 """
-@views function mesh_selector!(cache::MIRKCache{iip, T}) where {iip, T}
+@views function mesh_selector!(cache::Union{MIRKCache{iip, T}, FIRKCacheExpand{iip, T}, FIRKCacheNested{iip, T}}) where {iip, T}
     (; M, order, defect, mesh, mesh_dt) = cache
     (abstol, _, _), kwargs = __split_mirk_kwargs(; cache.kwargs...)
     N = length(cache.mesh)
@@ -134,7 +257,7 @@ function half_mesh!(mesh::Vector{T}, mesh_dt::Vector{T}) where {T}
     end
     return mesh, mesh_dt
 end
-half_mesh!(cache::MIRKCache) = half_mesh!(cache.mesh, cache.mesh_dt)
+half_mesh!(cache::Union{MIRKCache, FIRKCacheNested, FIRKCacheExpand}) = half_mesh!(cache.mesh, cache.mesh_dt)
 
 """
     defect_estimate!(cache::MIRKCache)
@@ -169,7 +292,7 @@ an interpolant
 
         z, z′ = sum_stages!(cache, w₂, w₂′, i)
         if iip
-            yᵢ₂ = cache.y[i + 1].du
+            yᵢ₂ = cache.y[i+1].du
             f(yᵢ₂, z, cache.p, mesh[i] + (T(1) - τ_star) * dt)
         else
             yᵢ₂ = f(z, cache.p, mesh[i] + (T(1) - τ_star) * dt)
@@ -183,6 +306,126 @@ an interpolant
     return maximum(Base.Fix1(maximum, abs), defect.u)
 end
 
+@views function defect_estimate!(cache::FIRKCacheExpand{iip, T}) where {iip, T}
+    (; f, M, stage, mesh, mesh_dt, defect, ITU) = cache
+    (; q_coeff, τ_star) = ITU
+
+    ctr = 1
+    K = zeros(eltype(cache.y[1].du), M, stage)
+    for i in 1:(length(mesh) - 1)
+        h = mesh_dt[i]
+
+        # Load interpolation residual
+        for j in 1:stage
+            K[:, j] = cache.y[ctr + j].du
+        end
+
+        # Defect estimate from q(x) at y_i + τ* * h
+        yᵢ₁ = copy(cache.y[ctr].du)
+        yᵢ₂ = copy(yᵢ₁)
+        z₁, z₁′ = eval_q(yᵢ₁, τ_star, h, q_coeff, K)
+        if iip
+            f(yᵢ₁, z₁, cache.p, mesh[i] + τ_star * h)
+        else
+            yᵢ₁ = f(z₁, cache.p, mesh[i] + τ_star * h)
+        end
+        yᵢ₁ .= (z₁′ .- yᵢ₁) ./ (abs.(yᵢ₁) .+ T(1))
+        est₁ = maximum(abs, yᵢ₁)
+
+        z₂, z₂′ = eval_q(yᵢ₂, (T(1) - τ_star), h, q_coeff, K)
+        # Defect estimate from q(x) at y_i + (1-τ*) * h
+        if iip
+            f(yᵢ₂, z₂, cache.p, mesh[i] + (T(1) - τ_star) * h)
+        else
+            yᵢ₂ = f(z₂, cache.p, mesh[i] + (T(1) - τ_star) * h)
+        end
+        yᵢ₂ .= (z₂′ .- yᵢ₂) ./ (abs.(yᵢ₂) .+ T(1))
+        est₂ = maximum(abs, yᵢ₂)
+
+        defect.u[i] .= est₁ > est₂ ? yᵢ₁ : yᵢ₂
+        ctr += stage + 1 # Advance one step
+    end
+
+    return maximum(Base.Fix1(maximum, abs), defect)
+end
+
+@views function defect_estimate!(cache::FIRKCacheNested{iip, T}) where {iip, T}
+    (; f, M, stage, mesh, mesh_dt, defect, TU, ITU, nest_cache, p_nestprob, prob) = cache
+    (; a, c) = TU
+    (; q_coeff, τ_star) = ITU
+
+    for i in 1:(length(mesh) - 1)
+        h = mesh_dt[i]
+        yᵢ₁ = copy(cache.y[i].du)
+        yᵢ₂ = copy(yᵢ₁)
+
+        K = copy(cache.k_discrete[i].du)
+
+        if minimum(abs.(K)) < 1e-2
+            K = fill(one(eltype(K)), size(K))
+        end
+
+        y_i = eltype(yᵢ₁) == Float64 ? yᵢ₁ : [y.value for y in yᵢ₁]
+
+        p_nestprob[1:2] .= promote(mesh[i], mesh_dt[i], one(eltype(y_i)))[1:2]
+        p_nestprob[3:end] = y_i
+        solve_cache!(nest_cache, K, p_nestprob)
+
+        # Defect estimate from q(x) at y_i + τ* * h
+        z₁, z₁′ = eval_q(yᵢ₁, τ_star, h, q_coeff, nest_cache.u)
+        if iip
+            f(yᵢ₁, z₁, cache.p, mesh[i] + τ_star * h)
+        else
+            yᵢ₁ = f(z₁, cache.p, mesh[i] + τ_star * h)
+        end
+        yᵢ₁ .= (z₁′ .- yᵢ₁) ./ (abs.(yᵢ₁) .+ T(1))
+        est₁ = maximum(abs, yᵢ₁)
+
+        # Defect estimate from q(x) at y_i + (1-τ*) * h
+        z₂, z₂′ = eval_q(yᵢ₂, (T(1) - τ_star), h, q_coeff, nest_cache.u)
+        if iip
+            f(yᵢ₂, z₂, cache.p, mesh[i] + (T(1) - τ_star) * h)
+        else
+            yᵢ₂ = f(z₂, cache.p, mesh[i] + (T(1) - τ_star) * h)
+        end
+        yᵢ₂ .= (z₂′ .- yᵢ₂) ./ (abs.(yᵢ₂) .+ T(1))
+        est₂ = maximum(abs, yᵢ₂)
+
+        defect.u[i] .= est₁ > est₂ ? yᵢ₁ : yᵢ₂
+    end
+
+    return maximum(Base.Fix1(maximum, abs), defect)
+end
+
+function get_q_coeffs(A, ki, h)
+    coeffs = A * ki
+    for i in axes(coeffs, 1)
+        coeffs[i] = coeffs[i] / (h^(i - 1))
+    end
+    return coeffs
+end
+
+function apply_q(y_i, τ, h, coeffs)
+    return y_i + sum(coeffs[i] * (τ * h)^(i) for i in axes(coeffs, 1))
+end
+
+function apply_q_prime(τ, h, coeffs)
+    return sum(i * coeffs[i] * (τ * h)^(i - 1) for i in axes(coeffs, 1))
+end
+
+function eval_q(y_i, τ, h, A, K)
+    M = size(K, 1)
+    q = zeros(M)
+    q′ = zeros(M)
+    for i in 1:M
+        ki = @view K[i, :]
+        coeffs = get_q_coeffs(A, ki, h)
+        q[i] = apply_q(y_i[i], τ, h, coeffs)
+        q′[i] = apply_q_prime(τ, h, coeffs)
+    end
+    return q, q′
+end
+
 """
     interp_setup!(cache::MIRKCache)
 

src/alg_utils.jl

@@ -4,8 +4,33 @@ for order in (2, 3, 4, 5, 6)
     @eval alg_stage(::$(alg)) = $(order - 1)
 end
 
+for stage in (1, 2, 3, 5, 7)
+    alg = Symbol("RadauIIa$(stage)")
+    @eval alg_order(::$(alg)) = $(2 * stage - 1)
+    @eval alg_stage(::$(alg)) = $stage
+end
+
+for stage in (2, 3, 4, 5)
+    alg = Symbol("LobattoIIIa$(stage)")
+    @eval alg_order(::$(alg)) = $(2 * stage - 2)
+    @eval alg_stage(::$(alg)) = $stage
+end
+
+for stage in (2, 3, 4, 5)
+    alg = Symbol("LobattoIIIb$(stage)")
+    @eval alg_order(::$(alg)) = $(2 * stage - 2)
+    @eval alg_stage(::$(alg)) = $stage
+end
+
+for stage in (2, 3, 4, 5)
+    alg = Symbol("LobattoIIIc$(stage)")
+    @eval alg_order(::$(alg)) = $(2 * stage - 2)
+    @eval alg_stage(::$(alg)) = $stage
+end
+
 SciMLBase.isautodifferentiable(::BoundaryValueDiffEqAlgorithm) = true
 SciMLBase.allows_arbitrary_number_types(::BoundaryValueDiffEqAlgorithm) = true
 SciMLBase.allowscomplex(alg::BoundaryValueDiffEqAlgorithm) = true
 
 SciMLBase.isadaptive(alg::AbstractMIRK) = true
+SciMLBase.isadaptive(alg::AbstractFIRK) = true