diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 72d19e21..e25f4f89 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -15,7 +15,6 @@ jobs: - Core version: - '1' - - '1.6' steps: - uses: actions/checkout@v4 - uses: julia-actions/setup-julia@v1 diff --git a/Project.toml b/Project.toml index 239fa0e5..b17ec754 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "BoundaryValueDiffEq" uuid = "764a87c0-6b3e-53db-9096-fe964310641d" -version = "4.1.0" +version = "4.2.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -21,6 +21,12 @@ SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77" UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" +[weakdeps] +ODEInterface = "54ca160b-1b9f-5127-a996-1867f4bc2a2c" + +[extensions] +BoundaryValueDiffEqODEInterfaceExt = "ODEInterface" + [compat] ADTypes = "0.2" Adapt = "3" @@ -28,18 +34,22 @@ ArrayInterface = "7" ConcreteStructs = "0.2" DiffEqBase = "6.94.2" ForwardDiff = "0.10" -NonlinearSolve = "1" +NonlinearSolve = "2" +ODEInterface = "0.5" +PreallocationTools = "0.4" RecursiveArrayTools = "2.38.10" Reexport = "0.2, 1.0" -SciMLBase = "1.70" +SciMLBase = "2" Setfield = "1" +SparseDiffTools = "2.6" TruncatedStacktraces = "1" UnPack = "1" -julia = "1.6" +julia = "1.9" [extras] DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d" NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" +ODEInterface = "54ca160b-1b9f-5127-a996-1867f4bc2a2c" OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" @@ -47,4 +57,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["StaticArrays", "Random", "DiffEqDevTools", "OrdinaryDiffEq", "Test", "NonlinearSolve", "SafeTestsets"] +test = ["StaticArrays", "Random", "DiffEqDevTools", "OrdinaryDiffEq", "Test", "NonlinearSolve", "SafeTestsets", "ODEInterface"] diff --git a/ext/BoundaryValueDiffEqODEInterfaceExt.jl b/ext/BoundaryValueDiffEqODEInterfaceExt.jl new file mode 100644 index 00000000..b06648e5 --- /dev/null +++ b/ext/BoundaryValueDiffEqODEInterfaceExt.jl @@ -0,0 +1,113 @@ +module BoundaryValueDiffEqODEInterfaceExt + +using SciMLBase, BoundaryValueDiffEq, ODEInterface +import ODEInterface: OptionsODE, OPT_ATOL, OPT_RTOL, OPT_METHODCHOICE, OPT_DIAGNOSTICOUTPUT, + OPT_ERRORCONTROL, OPT_SINGULARTERM, OPT_MAXSTEPS, OPT_BVPCLASS, OPT_SOLMETHOD, + OPT_RHS_CALLMODE, RHS_CALL_INSITU, evalSolution +import ODEInterface: Bvpm2, bvpm2_init, bvpm2_solve, bvpm2_destroy, bvpm2_get_x +import ODEInterface: bvpsol + +function _test_bvpm2_bvpsol_problem_criteria(_, ::SciMLBase.StandardBVProblem, alg::Symbol) + throw(ArgumentError("$(alg) does not support standard BVProblem. Only TwoPointBVProblem is supported.")) +end +function _test_bvpm2_bvpsol_problem_criteria(prob, ::TwoPointBVProblem, alg::Symbol) + @assert isinplace(prob) "$(alg) only supports inplace TwoPointBVProblem!" +end + +#------ +# BVPM2 +#------ +## TODO: We can specify Drhs using forwarddiff if we want to +function SciMLBase.__solve(prob::BVProblem, alg::BVPM2; dt = 0.0, reltol = 1e-3, kwargs...) + _test_bvpm2_bvpsol_problem_criteria(prob, prob.problem_type, :BVPM2) + + has_initial_guess = prob.u0 isa AbstractVector{<:AbstractArray} + no_odes, n, u0 = if has_initial_guess + length(first(prob.u0)), (length(prob.u0) - 1), reduce(hcat, prob.u0) + else + dt ≤ 0 && throw(ArgumentError("dt must be positive")) + length(prob.u0), Int(cld((prob.tspan[2] - prob.tspan[1]), dt)), prob.u0 + end + + mesh = collect(range(prob.tspan[1], stop = prob.tspan[2], length = n + 1)) + + no_left_bc = length(first(prob.f.bcresid_prototype.x)) + + initial_guess = Bvpm2() + bvpm2_init(initial_guess, no_odes, no_left_bc, mesh, u0, eltype(u0)[], + alg.max_num_subintervals) + + bvp2m_f(t, u, du) = prob.f(du, u, prob.p, t) + bvp2m_bc(ya, yb, bca, bcb) = prob.bc((bca, bcb), (ya, yb), prob.p) + + opt = OptionsODE(OPT_RTOL => reltol, OPT_METHODCHOICE => alg.method_choice, + OPT_DIAGNOSTICOUTPUT => alg.diagnostic_output, + OPT_SINGULARTERM => alg.singular_term, OPT_ERRORCONTROL => alg.error_control) + + sol, retcode, stats = bvpm2_solve(initial_guess, bvp2m_f, bvp2m_bc, opt) + retcode = retcode ≥ 0 ? ReturnCode.Success : ReturnCode.Failure + + x_mesh = bvpm2_get_x(sol) + sol_final = DiffEqBase.build_solution(prob, alg, x_mesh, + eachcol(evalSolution(sol, x_mesh)); retcode, stats) + + bvpm2_destroy(initial_guess) + bvpm2_destroy(sol) + + return sol_final +end + +#------- +# BVPSOL +#------- +function SciMLBase.__solve(prob::BVProblem, alg::BVPSOL; maxiters = 1000, reltol = 1e-3, + dt = 0.0, verbose = true, kwargs...) + _test_bvpm2_bvpsol_problem_criteria(prob, prob.problem_type, :BVPSOL) + @assert isa(prob.p, SciMLBase.NullParameters) "BVPSOL only supports NullParameters!" + @assert isa(prob.u0, AbstractVector{<:AbstractArray}) "BVPSOL requires a vector of initial guesses!" + n, u0 = (length(prob.u0) - 1), reduce(hcat, prob.u0) + mesh = collect(range(prob.tspan[1], stop = prob.tspan[2], length = n + 1)) + + opt = OptionsODE(OPT_RTOL => reltol, OPT_MAXSTEPS => maxiters, + OPT_BVPCLASS => alg.bvpclass, OPT_SOLMETHOD => alg.sol_method, + OPT_RHS_CALLMODE => RHS_CALL_INSITU) + + f!(t, u, du) = prob.f(du, u, prob.p, t) + function bc!(ya, yb, r) + ra = first(prob.f.bcresid_prototype.x) + rb = last(prob.f.bcresid_prototype.x) + prob.bc((ra, rb), (ya, yb), prob.p) + r[1:length(ra)] .= ra + r[(length(ra) + 1):(length(ra) + length(rb))] .= rb + return r + end + + sol_t, sol_x, retcode, stats = bvpsol(f!, bc!, mesh, u0, alg.odesolver, opt) + + if verbose + if retcode == -3 + @warn "Integrator failed to complete the trajectory" + elseif retcode == -4 + @warn "Gauss Newton method failed to converge" + elseif retcode == -5 + @warn "Given initial values inconsistent with separable linear bc" + elseif retcode == -6 + @warn """Iterative refinement faild to converge for `sol_method=0` + Termination since multiple shooting condition or + condition of Jacobian is too bad for `sol_method=1`""" + elseif retcode == -8 + @warn "Condensing algorithm for linear block system fails, try `sol_method=1`" + elseif retcode == -9 + @warn "Sparse linear solver failed" + elseif retcode == -10 + @warn "Real or integer work-space exhausted" + elseif retcode == -11 + @warn "Rank reduction failed - resulting rank is zero" + end + end + + return DiffEqBase.build_solution(prob, alg, sol_t, eachcol(sol_x); + retcode = retcode ≥ 0 ? ReturnCode.Success : ReturnCode.Failure, stats) +end + +end diff --git a/src/BoundaryValueDiffEq.jl b/src/BoundaryValueDiffEq.jl index 619a0e3b..157e6f69 100644 --- a/src/BoundaryValueDiffEq.jl +++ b/src/BoundaryValueDiffEq.jl @@ -9,7 +9,7 @@ import ArrayInterface: matrix_colors, parameterless_type import ConcreteStructs: @concrete import DiffEqBase: solve import ForwardDiff: pickchunksize -import RecursiveArrayTools: DiffEqArray +import RecursiveArrayTools: ArrayPartition, DiffEqArray import SciMLBase: AbstractDiffEqInterpolation import SparseDiffTools: AbstractSparseADType import TruncatedStacktraces: @truncate_stacktrace @@ -23,12 +23,21 @@ include("mirk_tableaus.jl") include("cache.jl") include("collocation.jl") include("nlprob.jl") -include("solve.jl") +include("solve/single_shooting.jl") +include("solve/mirk.jl") include("adaptivity.jl") include("interpolation.jl") +function SciMLBase.__solve(prob::BVProblem, alg::BoundaryValueDiffEqAlgorithm, args...; + kwargs...) + cache = init(prob, alg, args...; kwargs...) + return solve!(cache) +end + export Shooting export MIRK2, MIRK3, MIRK4, MIRK5, MIRK6 export MIRKJacobianComputationAlgorithm +# From ODEInterface.jl +export BVPM2, BVPSOL end diff --git a/src/adaptivity.jl b/src/adaptivity.jl index cec9c7b8..ab6e3eb4 100644 --- a/src/adaptivity.jl +++ b/src/adaptivity.jl @@ -24,11 +24,11 @@ function interval(mesh, t) end """ - mesh_selector!(cache::MIRKCache{T}) + mesh_selector!(cache::MIRKCache) Generate new mesh based on the defect. """ -@views function mesh_selector!(cache::MIRKCache{T}) where {T} +@views function mesh_selector!(cache::MIRKCache{iip, T}) where {iip, T} @unpack M, order, defect, mesh, mesh_dt = cache (_, MxNsub, abstol, _, _), kwargs = __split_mirk_kwargs(; cache.kwargs...) N = length(cache.mesh) @@ -81,11 +81,12 @@ Generate new mesh based on the defect. end """ - redistribute!(cache::MIRKCache{T}, Nsub_star, ŝ, mesh, mesh_dt) where {T} + redistribute!(cache::MIRKCache, Nsub_star, ŝ, mesh, mesh_dt) Generate a new mesh based on the `ŝ`. """ -function redistribute!(cache::MIRKCache{T}, Nsub_star, ŝ, mesh, mesh_dt) where {T} +function redistribute!(cache::MIRKCache{iip, T}, Nsub_star, ŝ, mesh, + mesh_dt) where {iip, T} N = length(mesh) ζ = sum(ŝ .* mesh_dt) / Nsub_star k, i = 1, 0 @@ -138,14 +139,14 @@ end half_mesh!(cache::MIRKCache) = half_mesh!(cache.mesh, cache.mesh_dt) """ - defect_estimate!(cache::MIRKCache{T}) + defect_estimate!(cache::MIRKCache) defect_estimate use the discrete solution approximation Y, plus stages of the RK method in 'k_discrete', plus some new stages in 'k_interp' to construct an interpolant """ -@views function defect_estimate!(cache::MIRKCache{T}) where {T} - @unpack M, stage, f!, alg, mesh, mesh_dt, defect = cache +@views function defect_estimate!(cache::MIRKCache{iip, T}) where {iip, T} + @unpack M, stage, f, alg, mesh, mesh_dt, defect = cache @unpack s_star, τ_star = cache.ITU # Evaluate at the first sample point @@ -157,16 +158,24 @@ an interpolant for i in 1:(length(mesh) - 1) dt = mesh_dt[i] - yᵢ₁ = cache.y[i].du - yᵢ₂ = cache.y[i + 1].du z, z′ = sum_stages!(cache, w₁, w₁′, i) - f!(yᵢ₁, z, cache.p, mesh[i] + τ_star * dt) + if iip + yᵢ₁ = cache.y[i].du + f(yᵢ₁, z, cache.p, mesh[i] + τ_star * dt) + else + yᵢ₁ = f(z, cache.p, mesh[i] + τ_star * dt) + end yᵢ₁ .= (z′ .- yᵢ₁) ./ (abs.(yᵢ₁) .+ T(1)) est₁ = maximum(abs, yᵢ₁) z, z′ = sum_stages!(cache, w₂, w₂′, i) - f!(yᵢ₂, z, cache.p, mesh[i] + (T(1) - τ_star) * dt) + if iip + 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) + end yᵢ₂ .= (z′ .- yᵢ₂) ./ (abs.(yᵢ₂) .+ T(1)) est₂ = maximum(abs, yᵢ₂) @@ -182,11 +191,10 @@ end `interp_setup!` prepare the extra stages in ki_interp for interpolant construction. Here, the ki_interp is the stages in one subinterval. """ -@views function interp_setup!(cache::MIRKCache{T}) where {T} +@views function interp_setup!(cache::MIRKCache{iip, T}) where {iip, T} @unpack x_star, s_star, c_star, v_star = cache.ITU - @unpack k_interp, k_discrete, f!, stage, new_stages, y, p, mesh, mesh_dt = cache + @unpack k_interp, k_discrete, f, stage, new_stages, y, p, mesh, mesh_dt = cache - [fill!(s, zero(eltype(s))) for s in new_stages] for r in 1:(s_star - stage) idx₁ = ((1:stage) .- 1) .* (s_star - stage) .+ r idx₂ = ((1:(r - 1)) .+ stage .- 1) .* (s_star - stage) .+ r @@ -203,7 +211,11 @@ Here, the ki_interp is the stages in one subinterval. new_stages[i] .= new_stages[i] .* mesh_dt[i] .+ (1 - v_star[r]) .* vec(y[i].du) .+ v_star[r] .* vec(y[i + 1].du) - f!(k_interp[i][:, r], new_stages[i], p, mesh[i] + c_star[r] * mesh_dt[i]) + if iip + f(k_interp[i][:, r], new_stages[i], p, mesh[i] + c_star[r] * mesh_dt[i]) + else + k_interp[i][:, r] .= f(new_stages[i], p, mesh[i] + c_star[r] * mesh_dt[i]) + end end end diff --git a/src/algorithms.jl b/src/algorithms.jl index a8aa25bc..5aff9f08 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -1,5 +1,5 @@ -const DEFAULT_NLSOLVE_SHOOTING = TrustRegion(; autodiff = Val(true)) -const DEFAULT_NLSOLVE_MIRK = NewtonRaphson(; autodiff = Val(true)) +const DEFAULT_NLSOLVE_SHOOTING = NewtonRaphson() +const DEFAULT_NLSOLVE_MIRK = NewtonRaphson() const DEFAULT_JACOBIAN_ALGORITHM_MIRK = MIRKJacobianComputationAlgorithm() # Algorithms @@ -50,3 +50,50 @@ for order in (2, 3, 4, 5, 6) end end end + +""" + BVPM2(; max_num_subintervals = 3000, method_choice = 4, diagnostic_output = 1, + error_control = 1, singular_term = nothing) + BVPM2(max_num_subintervals::Int, method_choice::Int, diagnostic_output::Int, + error_control::Int, singular_term) + +Fortran code for solving two-point boundary value problems. For detailed documentation, see +[ODEInterface.jl](https://github.com/luchr/ODEInterface.jl/blob/master/doc/SolverOptions.md#bvpm2). + +!!! warning + Only supports inplace two-point boundary value problems, with very limited forms of + input structures! + +!!! note + Only available if the `ODEInterface` package is loaded. +""" +Base.@kwdef struct BVPM2{S} <: BoundaryValueDiffEqAlgorithm + max_num_subintervals::Int = 3000 + method_choice::Int = 4 + diagnostic_output::Int = -1 + error_control::Int = 1 + singular_term::S = nothing +end + +""" + BVPSOL(; bvpclass = 2, sol_method = 0, odesolver = nothing) + BVPSOL(bvpclass::Int, sol_methods::Int, odesolver) + +A FORTRAN77 code which solves highly nonlinear two point boundary value problems using a +local linear solver (condensing algorithm) or a global sparse linear solver for the solution +of the arising linear subproblems, by Peter Deuflhard, Georg Bader, Lutz Weimann. +For detailed documentation, see +[ODEInterface.jl](https://github.com/luchr/ODEInterface.jl/blob/master/doc/SolverOptions.md#bvpsol). + +!!! warning + Only supports inplace two-point boundary value problems, with very limited forms of + input structures! + +!!! note + Only available if the `ODEInterface` package is loaded. +""" +Base.@kwdef struct BVPSOL{O} <: BoundaryValueDiffEqAlgorithm + bvpclass::Int = 2 + sol_method::Int = 0 + odesolver::O = nothing +end diff --git a/src/cache.jl b/src/cache.jl index c4a77cef..6d70d9e9 100644 --- a/src/cache.jl +++ b/src/cache.jl @@ -1,10 +1,10 @@ -@concrete struct MIRKCache{T} +@concrete struct MIRKCache{iip, T} order::Int # The order of MIRK method stage::Int # The state of MIRK method M::Int # The number of equations in_size - f! # FIXME: After supporting OOP functions - bc! # FIXME: After supporting OOP functions + f + bc prob # BVProblem problem_type # StandardBVProblem p # Parameters @@ -27,7 +27,7 @@ kwargs end -Base.eltype(::MIRKCache{T}) where {T} = T +Base.eltype(::MIRKCache{iip, T}) where {iip, T} = T """ expand_cache!(cache::MIRKCache) @@ -47,6 +47,8 @@ function expand_cache!(cache::MIRKCache) return cache end +__append_similar!(::Nothing, n, _) = nothing + function __append_similar!(x::AbstractVector{<:AbstractArray}, n, _) N = n - length(x) N == 0 && return x diff --git a/src/collocation.jl b/src/collocation.jl index ae0ed056..8ab995f3 100644 --- a/src/collocation.jl +++ b/src/collocation.jl @@ -4,18 +4,8 @@ function __initial_state_from_prob(u0::AbstractVector{<:AbstractVector}, _) [copy(vec(u)) for u in u0] end -# Auxiliary functions for evaluation -function eval_bc_residual!(residual::AbstractArray, _, bc!, y, p, mesh, u) - return bc!(residual, y, p, mesh) -end -function eval_bc_residual!(residual::AbstractArray, ::TwoPointBVProblem, bc!, y, p, mesh, u) - y₁ = first(y) - y₂ = last(y) - return bc!(residual, (y₁, y₂), p, (first(mesh), last(mesh))) -end - function Φ!(residual, cache::MIRKCache, y, u, p = cache.p) - return Φ!(residual, cache.fᵢ_cache, cache.k_discrete, cache.f!, cache.TU, + return Φ!(residual, cache.fᵢ_cache, cache.k_discrete, cache.f, cache.TU, y, u, p, cache.mesh, cache.mesh_dt, cache.stage) end @@ -44,3 +34,36 @@ end __maybe_matmul!(residᵢ, K[:, 1:stage], b[1:stage], -h, T(1)) end end + +function Φ(cache::MIRKCache, y, u, p = cache.p) + return Φ(cache.fᵢ_cache, cache.k_discrete, cache.f, cache.TU, y, u, p, cache.mesh, + cache.mesh_dt, cache.stage) +end + +@views function Φ(fᵢ_cache, k_discrete, f, TU::MIRKTableau, y, u, p, mesh, mesh_dt, + stage::Int) + @unpack c, v, x, b = TU + residuals = [similar(yᵢ) for yᵢ in y[1:(end - 1)]] + tmp = get_tmp(fᵢ_cache, u) + T = eltype(u) + for i in eachindex(k_discrete) + K = get_tmp(k_discrete[i], u) + residᵢ = residuals[i] + h = mesh_dt[i] + + yᵢ = get_tmp(y[i], u) + yᵢ₊₁ = get_tmp(y[i + 1], u) + + for r in 1:stage + @. tmp = (1 - v[r]) * yᵢ + v[r] * yᵢ₊₁ + __maybe_matmul!(tmp, K[:, 1:(r - 1)], x[r, 1:(r - 1)], h, T(1)) + K[:, r] .= f(tmp, p, mesh[i] + c[r] * h) + end + + # Update residual + @. residᵢ = yᵢ₊₁ - yᵢ + __maybe_matmul!(residᵢ, K[:, 1:stage], b[1:stage], -h, T(1)) + end + + return residuals +end diff --git a/src/nlprob.jl b/src/nlprob.jl index 008ba163..0b6468d1 100644 --- a/src/nlprob.jl +++ b/src/nlprob.jl @@ -1,75 +1,254 @@ -import SparseDiffTools: __init_𝒥 +function construct_nlproblem(cache::MIRKCache{iip}, y::AbstractVector) where {iip} + loss_bc = if iip + function loss_bc_internal!(resid::AbstractVector, u::AbstractVector, p = cache.p) + y_ = recursive_unflatten!(cache.y, u) + eval_bc_residual!(resid, cache.problem_type, cache.bc, y_, p, cache.mesh) + return resid + end + else + function loss_bc_internal(u::AbstractVector, p = cache.p) + y_ = recursive_unflatten!(cache.y, u) + return eval_bc_residual(cache.problem_type, cache.bc, y_, p, cache.mesh) + end + end -function construct_nlproblem(cache::MIRKCache, y::AbstractVector) - function loss_bc!(resid::AbstractVector, u::AbstractVector, p = cache.p) - y_ = recursive_unflatten!(cache.y, u) - eval_bc_residual!(resid, cache.problem_type, cache.bc!, y_, p, cache.mesh, u) - return resid + loss_collocation = if iip + function loss_collocation_internal!(resid::AbstractVector, u::AbstractVector, + p = cache.p) + y_ = recursive_unflatten!(cache.y, u) + resids = [get_tmp(r, u) for r in cache.residual[2:end]] + Φ!(resids, cache, y_, u, p) + recursive_flatten!(resid, resids) + return resid + end + else + function loss_collocation_internal(u::AbstractVector, p = cache.p) + y_ = recursive_unflatten!(cache.y, u) + resids = Φ(cache, y_, u, p) + xxx = mapreduce(vec, vcat, resids) + return xxx + end end - function loss_collocation!(resid::AbstractVector, u::AbstractVector, p = cache.p) - y_ = recursive_unflatten!(cache.y, u) - resids = [get_tmp(r, u) for r in cache.residual[2:end]] - Φ!(resids, cache, y_, u, p) - recursive_flatten!(resid, resids) - return resid + loss = if !(cache.problem_type isa TwoPointBVProblem) + if iip + function loss_internal!(resid::AbstractVector, u::AbstractVector, p = cache.p) + y_ = recursive_unflatten!(cache.y, u) + resids = [get_tmp(r, u) for r in cache.residual] + eval_bc_residual!(resids[1], cache.problem_type, cache.bc, y_, p, + cache.mesh) + Φ!(resids[2:end], cache, y_, u, p) + recursive_flatten!(resid, resids) + return resid + end + else + function loss_internal(u::AbstractVector, p = cache.p) + y_ = recursive_unflatten!(cache.y, u) + resid_bc = eval_bc_residual(cache.problem_type, cache.bc, y_, p, cache.mesh) + resid_co = Φ(cache, y_, u, p) + return vcat(resid_bc, mapreduce(vec, vcat, resid_co)) + end + end + else + # Reordering for 2 point BVP + if iip + function loss_internal_2point!(resid::AbstractVector, u::AbstractVector, + p = cache.p) + y_ = recursive_unflatten!(cache.y, u) + resids = [get_tmp(r, u) for r in cache.residual] + eval_bc_residual!(resids[1], cache.problem_type, cache.bc, y_, p, + cache.mesh) + Φ!(resids[2:end], cache, y_, u, p) + recursive_flatten_twopoint!(resid, resids) + return resid + end + else + function loss_internal_2point(u::AbstractVector, p = cache.p) + y_ = recursive_unflatten!(cache.y, u) + resid_bc = eval_bc_residual(cache.problem_type, cache.bc, y_, p, cache.mesh) + resid_co = Φ(cache, y_, u, p) + return vcat(resid_bc.x[1], mapreduce(vec, vcat, resid_co), resid_bc.x[2]) + end + end end - function loss!(resid::AbstractVector, u::AbstractVector, p) - y_ = recursive_unflatten!(cache.y, u) - resids = [get_tmp(r, u) for r in cache.residual] - eval_bc_residual!(resids[1], cache.problem_type, cache.bc!, y_, p, cache.mesh, u) - Φ!(resids[2:end], cache, y_, u, p) - recursive_flatten!(resid, resids) - return resid + return generate_nlprob(cache, y, loss_bc, loss_collocation, loss, cache.problem_type) +end + +function construct_sparse_banded_jac_prototype(y, M, N) + l = sum(i -> min(2M + i, M * N) - max(1, i - 1) + 1, 1:(M * (N - 1))) + Is = Vector{Int}(undef, l) + Js = Vector{Int}(undef, l) + idx = 1 + for i in 1:(M * (N - 1)), j in max(1, i - 1):min(2M + i, M * N) + Is[idx] = i + Js[idx] = j + idx += 1 + end + col_colorvec = Vector{Int}(undef, M * N) + for i in eachindex(col_colorvec) + col_colorvec[i] = mod1(i, min(2M + 1, M * N) + 1) + end + row_colorvec = Vector{Int}(undef, M * (N - 1)) + for i in eachindex(row_colorvec) + row_colorvec[i] = mod1(i, min(2M + 1, M * N) + 1) end - @unpack nlsolve, jac_alg = cache.alg + y_ = similar(y, length(Is)) + return (sparse(adapt(parameterless_type(y), Is), adapt(parameterless_type(y), Js), + y_, M * (N - 1), M * N), col_colorvec, row_colorvec) +end - resid = similar(y) +# Two Point Specialization +function construct_sparse_banded_jac_prototype(y::ArrayPartition, M, N) + l = sum(i -> min(2M + i, M * N) - max(1, i - 1) + 1, 1:(M * (N - 1))) + l_top = M * length(y.x[1].x[1]) + l_bot = M * length(y.x[1].x[2]) + + Is = Vector{Int}(undef, l + l_top + l_bot) + Js = Vector{Int}(undef, l + l_top + l_bot) + idx = 1 + + for i in 1:length(y.x[1].x[1]), j in 1:M + Is[idx] = i + Js[idx] = j + idx += 1 + end + + for i in 1:(M * (N - 1)), j in max(1, i - 1):min(2M + i, M * N) + Is[idx] = i + length(y.x[1].x[1]) + Js[idx] = j + idx += 1 + end + + for i in 1:length(y.x[1].x[2]), j in 1:M + Is[idx] = i + length(y.x[1].x[1]) + M * (N - 1) + Js[idx] = j + M * (N - 1) + idx += 1 + end - resid_bc, resid_collocation = @view(resid[1:(cache.M)]), @view(resid[(cache.M + 1):end]) + col_colorvec = Vector{Int}(undef, M * N) + for i in eachindex(col_colorvec) + col_colorvec[i] = mod1(i, min(2M + 1, M * N) + 1) + end + row_colorvec = Vector{Int}(undef, M * N) + for i in eachindex(row_colorvec) + row_colorvec[i] = mod1(i, min(2M + 1, M * N) + 1) + end + + y_ = similar(y, length(Is)) + return (sparse(adapt(parameterless_type(y), Is), adapt(parameterless_type(y), Js), + y_, M * N, M * N), col_colorvec, row_colorvec) +end + +function generate_nlprob(cache::MIRKCache{iip}, y, loss_bc, loss_collocation, loss, + _) where {iip} + @unpack nlsolve, jac_alg = cache.alg + N = length(cache.mesh) + + resid_bc = cache.prob.f.bcresid_prototype === nothing ? similar(y, cache.M) : + cache.prob.f.bcresid_prototype + resid_collocation = similar(y, cache.M * (N - 1)) sd_bc = jac_alg.bc_diffmode isa AbstractSparseADType ? SymbolicsSparsityDetection() : NoSparsityDetection() - cache_bc = sparse_jacobian_cache(jac_alg.bc_diffmode, sd_bc, loss_bc!, resid_bc, y) - N = length(cache.mesh) + if iip + cache_bc = sparse_jacobian_cache(jac_alg.bc_diffmode, sd_bc, loss_bc, resid_bc, y) + else + cache_bc = sparse_jacobian_cache(jac_alg.bc_diffmode, sd_bc, loss_bc, y; + fx = resid_bc) + end + sd_collocation = if jac_alg.collocation_diffmode isa AbstractSparseADType - Jₛ = construct_sparse_banded_jac_prototype(y, cache.M, N) - JacPrototypeSparsityDetection(; jac_prototype = Jₛ) + Jₛ, cvec, rvec = construct_sparse_banded_jac_prototype(y, cache.M, N) + PrecomputedJacobianColorvec(; jac_prototype = Jₛ, row_colorvec = rvec, + col_colorvec = cvec) else NoSparsityDetection() end - cache_collocation = sparse_jacobian_cache(jac_alg.collocation_diffmode, sd_collocation, - loss_collocation!, resid_collocation, y) - jac_prototype = vcat(__init_𝒥(cache_bc), __init_𝒥(cache_collocation)) + if iip + cache_collocation = sparse_jacobian_cache(jac_alg.collocation_diffmode, + sd_collocation, loss_collocation, resid_collocation, y) + else + cache_collocation = sparse_jacobian_cache(jac_alg.collocation_diffmode, + sd_collocation, loss_collocation, y; fx = resid_collocation) + end + + jac_prototype = vcat(init_jacobian(cache_bc), + jac_alg.collocation_diffmode isa AbstractSparseADType ? Jₛ : + init_jacobian(cache_collocation)) - function jac!(J, x, p) - # TODO: Pass `p` into `loss_bc!` and `loss_collocation!` - sparse_jacobian!(@view(J[1:(cache.M), :]), jac_alg.bc_diffmode, cache_bc, loss_bc!, - resid_bc, x) - sparse_jacobian!(@view(J[(cache.M + 1):end, :]), jac_alg.collocation_diffmode, - cache_collocation, loss_collocation!, resid_collocation, x) - return J + # TODO: Pass `p` into `loss_bc` and `loss_collocation`. Currently leads to a Tag + # mismatch for ForwardDiff + jac = if iip + function jac_internal!(J, x, p) + sparse_jacobian!(@view(J[1:(cache.M), :]), jac_alg.bc_diffmode, cache_bc, + loss_bc, resid_bc, x) + sparse_jacobian!(@view(J[(cache.M + 1):end, :]), jac_alg.collocation_diffmode, + cache_collocation, loss_collocation, resid_collocation, x) + return J + end + else + J_ = jac_prototype + function jac_internal(x, p) + sparse_jacobian!(@view(J_[1:(cache.M), :]), jac_alg.bc_diffmode, cache_bc, + loss_bc, x) + sparse_jacobian!(@view(J_[(cache.M + 1):end, :]), jac_alg.collocation_diffmode, + cache_collocation, loss_collocation, x) + return J_ + end end - return NonlinearProblem(NonlinearFunction{true}(loss!; jac = jac!, jac_prototype), - y, cache.p) + return NonlinearProblem(NonlinearFunction{iip}(loss; jac, jac_prototype), y, cache.p) end -function construct_sparse_banded_jac_prototype(y, M, N) - l = sum(i -> min(2M + i, M * N) - max(1, i - 1) + 1, 1:(M * (N - 1))) - Is = Vector{Int}(undef, l) - Js = Vector{Int}(undef, l) - idx = 1 - for i in 1:(M * (N - 1)), j in max(1, i - 1):min(2M + i, M * N) - Is[idx] = i - Js[idx] = j - idx += 1 +function generate_nlprob(cache::MIRKCache{iip}, y, loss_bc, loss_collocation, loss, + ::TwoPointBVProblem) where {iip} + @unpack nlsolve, jac_alg = cache.alg + N = length(cache.mesh) + + if !iip && cache.prob.f.bcresid_prototype === nothing + y_ = recursive_unflatten!(cache.y, y) + resid_ = cache.bc((y_[1], y_[end]), cache.p) + resid = ArrayPartition(ArrayPartition(resid_), similar(y, cache.M * (N - 1))) + else + resid = ArrayPartition(cache.prob.f.bcresid_prototype, + similar(y, cache.M * (N - 1))) end - y_ = similar(y, length(Is)) - return sparse(adapt(parameterless_type(y), Is), adapt(parameterless_type(y), Js), - y_, M * (N - 1), M * N) + + sd = if jac_alg.diffmode isa AbstractSparseADType + Jₛ, cvec, rvec = construct_sparse_banded_jac_prototype(resid, cache.M, N) + PrecomputedJacobianColorvec(; jac_prototype = Jₛ, row_colorvec = rvec, + col_colorvec = cvec) + else + NoSparsityDetection() + end + + if iip + diffcache = sparse_jacobian_cache(jac_alg.diffmode, sd, loss, resid, y) + else + diffcache = sparse_jacobian_cache(jac_alg.diffmode, sd, loss, y; fx = resid) + end + + jac_prototype = jac_alg.diffmode isa AbstractSparseADType ? Jₛ : + init_jacobian(diffcache) + + # TODO: Pass `p` into `loss_bc` and `loss_collocation`. Currently leads to a Tag + # mismatch for ForwardDiff + jac = if iip + function jac_internal!(J, x, p) + sparse_jacobian!(J, jac_alg.diffmode, diffcache, loss, resid, x) + return J + end + else + J_ = jac_prototype + function jac_internal(x, p) + sparse_jacobian!(J_, jac_alg.diffmode, diffcache, loss, x) + return J_ + end + end + + return NonlinearProblem(NonlinearFunction{iip}(loss; jac, jac_prototype), y, cache.p) end diff --git a/src/solve.jl b/src/solve/mirk.jl similarity index 71% rename from src/solve.jl rename to src/solve/mirk.jl index ab827be0..65978c4e 100644 --- a/src/solve.jl +++ b/src/solve/mirk.jl @@ -1,34 +1,7 @@ -function SciMLBase.__solve(prob::BVProblem, alg; kwargs...) - # If dispatch not directly defined - cache = init(prob, alg; kwargs...) - return solve!(cache) -end - -# Shooting Methods - -function SciMLBase.__solve(prob::BVProblem, alg::Shooting; kwargs...) - iip = isinplace(prob) - bc = prob.bc - u0 = deepcopy(prob.u0) - function loss!(resid, u0, p) - tmp_prob = ODEProblem{iip}(prob.f, u0, prob.tspan, p) - internal_sol = solve(tmp_prob, alg.ode_alg; kwargs...) - bc(resid, internal_sol, prob.p, internal_sol.t) - return nothing - end - opt = solve(NonlinearProblem(NonlinearFunction{true}(loss!), u0, prob.p), alg.nlsolve; - kwargs...) - sol_prob = ODEProblem{iip}(prob.f, opt.u, prob.tspan, prob.p) - sol = solve(sol_prob, alg.ode_alg; kwargs...) - return DiffEqBase.solution_new_retcode(sol, - sol.retcode == opt.retcode ? ReturnCode.Success : ReturnCode.Failure) -end - -# MIRK Methods - -function SciMLBase.__init(prob::BVProblem, alg::AbstractMIRK; dt = 0.0, abstol = 1e-3, - adaptive = true, kwargs...) +function SciMLBase.__init(prob::BVProblem, alg::AbstractMIRK; dt = 0.0, + abstol = 1e-3, adaptive = true, kwargs...) has_initial_guess = prob.u0 isa AbstractVector{<:AbstractArray} + iip = isinplace(prob) (T, M, n) = if has_initial_guess # If user provided a vector of initial guesses _u0 = first(prob.u0) @@ -67,10 +40,28 @@ function SciMLBase.__init(prob::BVProblem, alg::AbstractMIRK; dt = 0.0, abstol = k_interp = adaptive ? [similar(X, M, ITU.s_star - stage) for _ in 1:n] : [similar(X, 0, 0) for _ in 1:n] - # FIXME: Here we are making the assumption that size(first(residual)) == size(first(y)) - # This won't hold true for underconstrained or overconstrained problems - resid₁_size = size(X) - residual = [maybe_allocate_diffcache(vec(copy(yᵢ)), chunksize, alg.jac_alg) for yᵢ in y₀] + resid₁_size = if prob.f.bcresid_prototype === nothing + size(X) + elseif prob.f.bcresid_prototype isa ArrayPartition + size.(prob.f.bcresid_prototype.x) + else + size(prob.f.bcresid_prototype) + end + + if iip + if prob.f.bcresid_prototype === nothing + residual = [maybe_allocate_diffcache(vec(copy(yᵢ)), chunksize, alg.jac_alg) + for yᵢ in y₀] + else + residual = vcat([ + maybe_allocate_diffcache(vec(copy(prob.f.bcresid_prototype)), + chunksize, alg.jac_alg)], + [maybe_allocate_diffcache(vec(copy(yᵢ)), chunksize, alg.jac_alg) + for yᵢ in y₀[2:end]]) + end + else + residual = nothing + end defect = adaptive ? [similar(X, M) for _ in 1:n] : [similar(X, 0) for _ in 1:n] @@ -79,23 +70,46 @@ function SciMLBase.__init(prob::BVProblem, alg::AbstractMIRK; dt = 0.0, abstol = # Transform the functions to handle non-vector inputs f, bc = if X isa AbstractVector prob.f, prob.bc - else - function vecf(du, u, p, t) + elseif iip + function vecf!(du, u, p, t) du_ = reshape(du, size(X)) x_ = reshape(u, size(X)) prob.f(du_, x_, p, t) return du end - function vecbc(resid, sol, p, t) + function vecbc!(resid, sol, p, t) resid_ = reshape(resid, resid₁_size) sol_ = map(s -> reshape(s, size(X)), sol) prob.bc(resid_, sol_, p, t) return resid end + function vecbc!((resida, residb), (ua, ub), p) + resida_ = reshape(resida, resid₁_size[1]) + residb_ = reshape(residb, resid₁_size[2]) + ua_ = reshape(ua, size(X)) + ub_ = reshape(ub, size(X)) + prob.bc((resida_, residb_), (ua_, ub_), p) + return (resida, residb) + end + vecf!, vecbc! + else + function vecf(u, p, t) + x_ = reshape(u, size(X)) + return vec(prob.f(x_, p, t)) + end + function vecbc(sol, p, t) + sol_ = map(s -> reshape(s, size(X)), sol) + return vec(prob.bc(sol_, p, t)) + end + function vecbc((ua, ub), p) + ua_ = reshape(ua, size(X)) + ub_ = reshape(ub, size(X)) + return vec.(prob.bc((ua_, ub_), p)) + end vecf, vecbc end - return MIRKCache{T}(alg_order(alg), stage, M, size(X), f, bc, prob, + return MIRKCache{iip, T}(alg_order(alg), stage, M, size(X), f, bc, prob, prob.problem_type, prob.p, alg, TU, ITU, mesh, mesh_dt, k_discrete, k_interp, y, y₀, residual, fᵢ_cache, fᵢ₂_cache, defect, new_stages, (; defect_threshold, MxNsub, abstol, dt, adaptive, kwargs...)) @@ -157,6 +171,6 @@ function SciMLBase.solve!(cache::MIRKCache) end u = [reshape(y, cache.in_size) for y in cache.y₀] - return DiffEqBase.build_solution(prob, alg, mesh, - u; interp = MIRKInterpolation(mesh, u, cache), retcode = info) + return DiffEqBase.build_solution(prob, alg, cache.mesh, + u; interp = MIRKInterpolation(cache.mesh, u, cache), retcode = info) end diff --git a/src/solve/single_shooting.jl b/src/solve/single_shooting.jl new file mode 100644 index 00000000..8b8cc9ab --- /dev/null +++ b/src/solve/single_shooting.jl @@ -0,0 +1,28 @@ +# TODO: Differentiate between nlsolve kwargs and odesolve kwargs +# TODO: Support Non-Vector Inputs +function SciMLBase.__solve(prob::BVProblem, alg::Shooting; kwargs...) + iip = isinplace(prob) + bc = prob.bc + u0 = deepcopy(prob.u0) + loss_fn = if iip + function loss!(resid, u0, p) + tmp_prob = ODEProblem{iip}(prob.f, u0, prob.tspan, p) + internal_sol = solve(tmp_prob, alg.ode_alg; kwargs...) + eval_bc_residual!(resid, prob.problem_type, bc, internal_sol, p) + return nothing + end + else + function loss(u0, p) + tmp_prob = ODEProblem(prob.f, u0, prob.tspan, p) + internal_sol = solve(tmp_prob, alg.ode_alg; kwargs...) + return eval_bc_residual(prob.problem_type, bc, internal_sol, p) + end + end + opt = solve(NonlinearProblem(NonlinearFunction{iip}(loss_fn; prob.f.jac_prototype, + resid_prototype = prob.f.bcresid_prototype), u0, prob.p), alg.nlsolve; + kwargs...) + sol_prob = ODEProblem{iip}(prob.f, opt.u, prob.tspan, prob.p) + sol = solve(sol_prob, alg.ode_alg; kwargs...) + return DiffEqBase.solution_new_retcode(sol, + sol.retcode == opt.retcode ? ReturnCode.Success : ReturnCode.Failure) +end diff --git a/src/types.jl b/src/types.jl index 4396cc93..ecbe1ae8 100644 --- a/src/types.jl +++ b/src/types.jl @@ -33,23 +33,39 @@ end @truncate_stacktrace MIRKInterpTableau 1 # Sparsity Detection -@static if VERSION < v"1.9" - # Sparse Linear Solvers in LinearSolve.jl are a bit flaky on older versions - Base.@kwdef struct MIRKJacobianComputationAlgorithm{BD, CD} - bc_diffmode::BD = AutoForwardDiff() - collocation_diffmode::CD = AutoForwardDiff() - end -else - Base.@kwdef struct MIRKJacobianComputationAlgorithm{BD, CD} - bc_diffmode::BD = AutoForwardDiff() - collocation_diffmode::CD = AutoSparseForwardDiff() +@concrete struct MIRKJacobianComputationAlgorithm + bc_diffmode + collocation_diffmode + diffmode +end + +function MIRKJacobianComputationAlgorithm(diffmode = missing; + collocation_diffmode = missing, bc_diffmode = missing) + if diffmode !== missing + @assert collocation_diffmode === missing && bc_diffmode === missing + return MIRKJacobianComputationAlgorithm(diffmode, diffmode, diffmode) + else + @static if VERSION < v"1.9" + diffmode = AutoForwardDiff() + bc_diffmode = bc_diffmode === missing ? AutoForwardDiff() : bc_diffmode + collocation_diffmode = collocation_diffmode === missing ? + AutoForwardDiff() : collocation_diffmode + else + diffmode = AutoSparseForwardDiff() + bc_diffmode = bc_diffmode === missing ? AutoForwardDiff() : bc_diffmode + collocation_diffmode = collocation_diffmode === missing ? + AutoSparseForwardDiff() : collocation_diffmode + end + return MIRKJacobianComputationAlgorithm(bc_diffmode, collocation_diffmode, + collocation_diffmode) end end __needs_diffcache(::Union{AutoForwardDiff, AutoSparseForwardDiff}) = true __needs_diffcache(_) = false function __needs_diffcache(jac_alg::MIRKJacobianComputationAlgorithm) - return __needs_diffcache(jac_alg.bc_diffmode) || + return __needs_diffcache(jac_alg.diffmode) || + __needs_diffcache(jac_alg.bc_diffmode) || __needs_diffcache(jac_alg.collocation_diffmode) end diff --git a/src/utils.jl b/src/utils.jl index a5d569d7..c61d18f6 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -15,6 +15,18 @@ end end return y end +@views function recursive_flatten_twopoint!(y::AbstractVector, x::Vector{<:AbstractArray}) + x_, xiter = Iterators.peel(x) + # x_ will be an ArrayPartition + copyto!(y[1:length(x_.x[1])], x_.x[1]) + i = length(x_.x[1]) + for xᵢ in xiter + copyto!(y[(i + 1):(i + length(xᵢ))], xᵢ) + i += length(xᵢ) + end + copyto!(y[(i + 1):(i + length(x_.x[2]))], x_.x[2]) + return y +end @views function recursive_unflatten!(y::Vector{<:AbstractArray}, x::AbstractVector) i = 0 @@ -57,3 +69,22 @@ function __maybe_matmul!(z, A, b, α = eltype(z)(1), β = eltype(z)(0)) end return z end + +## Easier to dispatch +eval_bc_residual(pt, bc, sol, p) = eval_bc_residual(pt, bc, sol, p, sol.t) +eval_bc_residual(_, bc, sol, p, t) = bc(sol, p, t) +function eval_bc_residual(::TwoPointBVProblem, bc, sol, p, t) + ua = sol isa AbstractVector ? sol[1] : sol(first(t)) + ub = sol isa AbstractVector ? sol[end] : sol(last(t)) + resid₀, resid₁ = bc((ua, ub), p) + return ArrayPartition(resid₀, resid₁) +end + +eval_bc_residual!(resid, pt, bc!, sol, p) = eval_bc_residual!(resid, pt, bc!, sol, p, sol.t) +eval_bc_residual!(resid, _, bc!, sol, p, t) = bc!(resid, sol, p, t) +@views function eval_bc_residual!(resid, ::TwoPointBVProblem, bc!, sol, p, t) + ua = sol isa AbstractVector ? sol[1] : sol(first(t)) + ub = sol isa AbstractVector ? sol[end] : sol(last(t)) + bc!((resid.x[1], resid.x[2]), (ua, ub), p) + return resid +end diff --git a/test/ensemble.jl b/test/ensemble.jl index 28919a94..e0771993 100644 --- a/test/ensemble.jl +++ b/test/ensemble.jl @@ -1,4 +1,4 @@ -using BoundaryValueDiffEq, Random +using BoundaryValueDiffEq, Random, Test function ode!(du, u, p, t) du[1] = u[2] @@ -10,21 +10,21 @@ function bc!(residual, u, p, t) residual[2] = u[end][1] end -function prob_func(prob, i, repeat) - remake(prob, p = [rand()]) -end +prob_func(prob, i, repeat) = remake(prob, p = [rand()]) initial_guess = [0.0, 1.0] tspan = (0, pi / 2) p = [rand()] bvp = BVProblem(ode!, bc!, initial_guess, tspan, p) -ensemble_prob = EnsembleProblem(bvp, prob_func = prob_func) +ensemble_prob = EnsembleProblem(bvp; prob_func) @testset "$(solver)" for solver in (MIRK2, MIRK3, MIRK4, MIRK5, MIRK6) jac_algs = [MIRKJacobianComputationAlgorithm(), MIRKJacobianComputationAlgorithm(; bc_diffmode = AutoFiniteDiff(), collocation_diffmode = AutoSparseFiniteDiff())] for jac_alg in jac_algs - @test_nowarn solve(ensemble_prob, solver(; jac_alg), trajectories = 10, dt = 0.1) + # Not sure why it is throwing so many warnings + sol = solve(ensemble_prob, solver(; jac_alg); trajectories = 10, dt = 0.1) + @test sol.converged end end diff --git a/test/mirk_convergence_tests.jl b/test/mirk_convergence_tests.jl index 355280be..b14164d6 100644 --- a/test/mirk_convergence_tests.jl +++ b/test/mirk_convergence_tests.jl @@ -6,44 +6,62 @@ for order in (2, 3, 4, 5, 6) end # First order test -function func_1!(du, u, p, t) +function f1!(du, u, p, t) du[1] = u[2] du[2] = 0 end +f1(u, p, t) = [u[2], 0] -# Not able to change the initial condition. -# Hard coded solution. -func_1 = ODEFunction(func_1!, analytic = (u0, p, t) -> [5 - t, -1]) +# Second order linear test +function f2!(du, u, p, t) + du[1] = u[2] + du[2] = -u[1] +end +f2(u, p, t) = [u[2], -u[1]] function boundary!(residual, u, p, t) residual[1] = u[1][1] - 5 residual[2] = u[end][1] end +boundary(u, p, t) = [u[1][1] - 5, u[end][1]] -function boundary_two_point!(residual, u, p, t) - ua = u[1] - ub = u[end] - residual[1] = ua[1] - 5 - residual[2] = ub[1] -end - -# Second order linear test -function func_2!(du, u, p, t) - du[1] = u[2] - du[2] = -u[1] +function boundary_two_point!((resida, residb), (ua, ub), p) + resida[1] = ua[1] - 5 + residb[1] = ub[1] end +boundary_two_point((ua, ub), p) = [ua[1] - 5, ub[1]] # Not able to change the initial condition. # Hard coded solution. -func_2 = ODEFunction(func_2!, - analytic = (u0, p, t) -> [5 * (cos(t) - cot(5) * sin(t)), - 5 * (-cos(t) * cot(5) - sin(t))]) +odef1! = ODEFunction(f1!, analytic = (u0, p, t) -> [5 - t, -1]) +odef1 = ODEFunction(f1, analytic = (u0, p, t) -> [5 - t, -1]) + +odef2! = ODEFunction(f2!, + analytic = (u0, p, t) -> [ + 5 * (cos(t) - cot(5) * sin(t)), + 5 * (-cos(t) * cot(5) - sin(t)), + ]) +odef2 = ODEFunction(f2, + analytic = (u0, p, t) -> [ + 5 * (cos(t) - cot(5) * sin(t)), + 5 * (-cos(t) * cot(5) - sin(t)), + ]) + +bcresid_prototype = (Array{Float64}(undef, 1), Array{Float64}(undef, 1)) + tspan = (0.0, 5.0) u0 = [5.0, -3.5] -probArr = [BVProblem(func_1, boundary!, u0, tspan), - BVProblem(func_2, boundary!, u0, tspan), - TwoPointBVProblem(func_1, boundary_two_point!, u0, tspan), - TwoPointBVProblem(func_2, boundary_two_point!, u0, tspan)] + +probArr = [ + BVProblem(odef1!, boundary!, u0, tspan), + BVProblem(odef1, boundary, u0, tspan), + BVProblem(odef2!, boundary!, u0, tspan), + BVProblem(odef2, boundary, u0, tspan), + TwoPointBVProblem(odef1!, boundary_two_point!, u0, tspan; bcresid_prototype), + TwoPointBVProblem(odef1, boundary_two_point, u0, tspan; bcresid_prototype), + TwoPointBVProblem(odef2!, boundary_two_point!, u0, tspan; bcresid_prototype), + TwoPointBVProblem(odef2, boundary_two_point, u0, tspan; bcresid_prototype), +]; testTol = 0.2 affineTol = 1e-2 @@ -52,7 +70,7 @@ dts = 1 .// 2 .^ (3:-1:1) @info "Collocation method (MIRK)" @testset "Affineness" begin - @testset "Problem: $i" for i in (1, 3) + @testset "Problem: $i" for i in (1, 2, 5, 6) prob = probArr[i] @testset "MIRK$order" for order in (2, 3, 4, 5, 6) @time sol = solve(prob, mirk_solver(Val(order)), dt = 0.2) @@ -62,7 +80,7 @@ dts = 1 .// 2 .^ (3:-1:1) end @testset "Convergence on Linear" begin - @testset "Problem: $i" for i in (2, 4) + @testset "Problem: $i" for i in (3, 4, 7, 8) prob = probArr[i] @testset "MIRK$order" for (i, order) in enumerate((2, 3, 4, 5, 6)) @time sim = test_convergence(dts, prob, mirk_solver(Val(order)); @@ -82,13 +100,14 @@ function simplependulum!(du, u, p, t) du[2] = -(g / L) * sin(θ) end -function bc1!(residual, u, p, t) +# FIXME: This is a really bad test. Needs interpolation +function bc_pendulum!(residual, u, p, t) residual[1] = u[end ÷ 2][1] + π / 2 # the solution at the middle of the time span should be -pi/2 residual[2] = u[end][1] - π / 2 # the solution at the end of the time span should be pi/2 end u0 = MVector{2}([pi / 2, pi / 2]) -bvp1 = BVProblem(simplependulum!, bc1!, u0, tspan) +bvp1 = BVProblem(simplependulum!, bc_pendulum!, u0, tspan) jac_alg = MIRKJacobianComputationAlgorithm(; bc_diffmode = AutoFiniteDiff(), collocation_diffmode = AutoSparseFiniteDiff()) diff --git a/test/non_vector_inputs.jl b/test/non_vector_inputs.jl index eb81d5bc..3a4ddddf 100644 --- a/test/non_vector_inputs.jl +++ b/test/non_vector_inputs.jl @@ -1,27 +1,52 @@ -using BoundaryValueDiffEq, DiffEqBase, DiffEqDevTools, LinearAlgebra, Test +using BoundaryValueDiffEq, LinearAlgebra, OrdinaryDiffEq, Test for order in (2, 3, 4, 5, 6) s = Symbol("MIRK$(order)") @eval mirk_solver(::Val{$order}) = $(s)() end -function func_1!(du, u, p, t) +function f1!(du, u, p, t) du[1, 1] = u[1, 2] du[1, 2] = 0 end +function f1(u, p, t) + return [u[1, 2] 0] +end + function boundary!(residual, u, p, t) residual[1, 1] = u[1][1, 1] - 5 residual[1, 2] = u[end][1, 1] end +function boundary!((resida, residb), (ua, ub), p) + resida[1, 1] = ua[1, 1] - 5 + residb[1, 1] = ub[1, 1] +end + +function boundary(u, p, t) + return [u[1][1, 1] - 5 u[end][1, 1]] +end + +function boundary((ua, ub), p) + return (reshape([ua[1, 1] - 5], (1, 1)), reshape([ub[1, 1]], (1, 1))) +end + tspan = (0.0, 5.0) u0 = [5.0 -3.5] -prob = BVProblem(func_1!, boundary!, u0, tspan) +probs = [ + BVProblem(f1!, boundary!, u0, tspan), + TwoPointBVProblem(f1!, boundary!, u0, tspan; + bcresid_prototype = (Array{Float64}(undef, 1, 1), Array{Float64}(undef, 1, 1))), + BVProblem(f1, boundary, u0, tspan), + TwoPointBVProblem(f1, boundary, u0, tspan), +]; @testset "Affineness" begin @testset "MIRK$order" for order in (2, 3, 4, 5, 6) - @time sol = solve(prob, mirk_solver(Val(order)); dt = 0.2) - @test norm(diff(first.(sol.u)) .+ 0.2, Inf) + abs(sol[1][1] - 5) < 0.01 + for prob in probs + @time sol = solve(prob, mirk_solver(Val(order)); dt = 0.2) + @test norm(diff(first.(sol.u)) .+ 0.2, Inf) + abs(sol[1][1] - 5) < 0.01 + end end end diff --git a/test/odeinterface_ex7.jl b/test/odeinterface_ex7.jl new file mode 100644 index 00000000..fc900b21 --- /dev/null +++ b/test/odeinterface_ex7.jl @@ -0,0 +1,48 @@ +using Test, BoundaryValueDiffEq, LinearAlgebra, ODEInterface, Random + +# Adaptation of https://github.com/luchr/ODEInterface.jl/blob/958b6023d1dabf775033d0b89c5401b33100bca3/examples/BasicExamples/ex7.jl +function ex7_f!(du, u, p, t) + ϵ = p[1] + u₁, λ = u + du[1] = (sin(t)^2 + λ * sin(t)^4 / u₁) / ϵ + du[2] = 0 + return nothing +end + +function ex7_2pbc!((resa, resb), (ua, ub), p) + resa[1] = ua[1] - 1 + resb[1] = ub[1] - 1 + return nothing +end + +u0 = [0.5, 1.0] +p = [0.1] +tspan = (-π / 2, π / 2) + +tpprob = TwoPointBVProblem(ex7_f!, ex7_2pbc!, u0, tspan, p; + bcresid_prototype = (zeros(1), zeros(1))) + +@info "BVPM2" + +sol_bvpm2 = solve(tpprob, BVPM2(); dt = π / 20) +@test SciMLBase.successful_retcode(sol_bvpm2) +resid_f = (Array{Float64, 1}(undef, 1), Array{Float64, 1}(undef, 1)) +ex7_2pbc!(resid_f, (sol_bvpm2(tspan[1]), sol_bvpm2(tspan[2])), nothing) +@test norm(resid_f) < 1e-6 + +function ex7_f2!(du, u, p, t) + u₁, λ = u + du[1] = (sin(t)^2 + λ * sin(t)^4 / u₁) / 0.1 + du[2] = 0 + return nothing +end + +@info "BVPSOL" + +initial_u0 = [sol_bvpm2(t) .+ rand() for t in tspan[1]:(π / 20):tspan[2]] +tpprob = TwoPointBVProblem(ex7_f2!, ex7_2pbc!, initial_u0, tspan; + bcresid_prototype = (zeros(1), zeros(1))) + +# Just test that it runs. BVPSOL only works with linearly separable BCs. +# TODO: Implement appolo reentry example from ODEInterface.jl +sol_bvpsol = solve(tpprob, BVPSOL(); dt = π / 20) diff --git a/test/orbital.jl b/test/orbital.jl index 5a2a1294..4bb387fe 100644 --- a/test/orbital.jl +++ b/test/orbital.jl @@ -1,8 +1,7 @@ # Lambert's Problem +using BoundaryValueDiffEq, OrdinaryDiffEq, LinearAlgebra, Test -using BoundaryValueDiffEq -using DiffEqBase, OrdinaryDiffEq, LinearAlgebra, NonlinearSolve -using Test +@info "Testing Lambert's Problem" y0 = [ -4.7763169762853989E+06, @@ -28,7 +27,7 @@ t1 = 86400 * 2.3577522023524125E+04 tspan = (t0, t1) # ODE solver -function orbital(dy, y, p, t) +function orbital!(dy, y, p, t) r2 = (y[1]^2 + y[2]^2 + y[3]^2) r3 = r2^(3 / 2) w = 1 + 1.5J2 * (req * req / r2) * (1 - 5y[3] * y[3] / r2) @@ -42,45 +41,51 @@ function orbital(dy, y, p, t) end function bc!_generator(resid, sol, init_val) - resid[1] = sol[1][1] - init_val[1] - resid[2] = sol[1][2] - init_val[2] - resid[3] = sol[1][3] - init_val[3] - resid[4] = sol[end][1] - init_val[4] - resid[5] = sol[end][2] - init_val[5] - resid[6] = sol[end][3] - init_val[6] + resid[1] = sol(t0)[1] - init_val[1] + resid[2] = sol(t0)[2] - init_val[2] + resid[3] = sol(t0)[3] - init_val[3] + resid[4] = sol(t1)[1] - init_val[4] + resid[5] = sol(t1)[2] - init_val[5] + resid[6] = sol(t1)[3] - init_val[6] end + +function bc!_generator_2p((resid0, resid1), (ua, ub), init_val) + resid0[1] = ua[1] - init_val[1] + resid0[2] = ua[2] - init_val[2] + resid0[3] = ua[3] - init_val[3] + resid1[1] = ub[1] - init_val[4] + resid1[2] = ub[2] - init_val[5] + resid1[3] = ub[3] - init_val[6] +end + cur_bc! = (resid, sol, p, t) -> bc!_generator(resid, sol, init_val) +cur_bc_2point! = (resid, sol, p) -> bc!_generator_2p(resid, sol, init_val) resid_f = Array{Float64}(undef, 6) - -### Test the IVP Near the true solution +resid_f_2p = (Array{Float64, 1}(undef, 3), Array{Float64, 1}(undef, 3)) TestTol = 0.05 ### Now use the BVP solver to get closer -bvp = BVProblem(orbital, cur_bc!, y0, tspan) -nlsolve = NewtonRaphson(; autodiff = Val(false), diff_type = Val(:central)) -@time sol = solve(bvp, - Shooting(DP5(); nlsolve), - force_dtmin = true, - abstol = 1e-13, - reltol = 1e-13) -cur_bc!(resid_f, sol, nothing, sol.t) -@test norm(resid_f, Inf) < TestTol - -nlsolve = NewtonRaphson(; autodiff = Val(false), diff_type = Val(:forward)) -@time sol = solve(bvp, - Shooting(DP5(); nlsolve), - force_dtmin = true, - abstol = 1e-13, - reltol = 1e-13) -cur_bc!(resid_f, sol, nothing, sol.t) -@test norm(resid_f, Inf) < TestTol +bvp = BVProblem(orbital!, cur_bc!, y0, tspan) +for autodiff in (AutoForwardDiff(), AutoFiniteDiff(; fdtype = Val(:central)), + AutoSparseForwardDiff(), AutoFiniteDiff(; fdtype = Val(:forward)), + AutoSparseFiniteDiff()) + nlsolve = NewtonRaphson(; autodiff) + @time sol = solve(bvp, Shooting(DP5(); nlsolve); force_dtmin = true, abstol = 1e-13, + reltol = 1e-13) + cur_bc!(resid_f, sol, nothing, sol.t) + @test norm(resid_f, Inf) < TestTol +end -nlsolve = NewtonRaphson(; autodiff = Val(true)) -@time sol = solve(bvp, - Shooting(DP5(); nlsolve), - force_dtmin = true, - abstol = 1e-13, - reltol = 1e-13) -cur_bc!(resid_f, sol, nothing, sol.t) -@test norm(resid_f, Inf) < TestTol +### Using the TwoPoint BVP Structure +bvp = TwoPointBVProblem(orbital!, cur_bc_2point!, y0, tspan; + bcresid_prototype = (Array{Float64}(undef, 3), Array{Float64}(undef, 3))) +for autodiff in (AutoForwardDiff(), AutoFiniteDiff(; fdtype = Val(:central)), + AutoSparseForwardDiff(), AutoFiniteDiff(; fdtype = Val(:forward)), + AutoSparseFiniteDiff()) + nlsolve = NewtonRaphson(; autodiff) + @time sol = solve(bvp, Shooting(DP5(); nlsolve); force_dtmin = true, abstol = 1e-13, + reltol = 1e-13) + cur_bc_2point!(resid_f_2p, (sol(t0), sol(t1)), nothing) + @test norm(vcat(resid_f_2p...), Inf) < TestTol +end diff --git a/test/runtests.jl b/test/runtests.jl index 9d7a3797..af8b3bbd 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,5 +1,3 @@ -using BoundaryValueDiffEq -using DiffEqBase, OrdinaryDiffEq, DiffEqDevTools using Test, SafeTestsets @testset "Boundary Value Problem Tests" begin @@ -24,6 +22,12 @@ using Test, SafeTestsets end end + @time @testset "ODE Interface Solvers" begin + @time @safetestset "ODE Interface Tests" begin + include("odeinterface_ex7.jl") + end + end + @time @testset "Non Vector Inputs Tests" begin @time @safetestset "Non Vector Inputs" begin include("non_vector_inputs.jl") diff --git a/test/shooting_tests.jl b/test/shooting_tests.jl index 23f218c8..131af201 100644 --- a/test/shooting_tests.jl +++ b/test/shooting_tests.jl @@ -1,32 +1,89 @@ -using BoundaryValueDiffEq -using DiffEqBase, OrdinaryDiffEq, DiffEqDevTools -using Test, LinearAlgebra +using BoundaryValueDiffEq, LinearAlgebra, OrdinaryDiffEq, Test @info "Shooting method" -function f(du, u, p, t) - (x, v) = u - du[1] = v - du[2] = -x +@info "Multi-Point BVProblem" # Not really but using that API + +tspan = (0.0, 100.0) +u0 = [0.0, 1.0] + +# Inplace +function f1!(du, u, p, t) + du[1] = u[2] + du[2] = -u[1] + return nothing end -function bc!(resid, sol, p, t) - resid[1] = sol[1][1] - resid[2] = sol[end][1] - 1 +function bc1!(resid, sol, p, t) + t₀, t₁ = first(t), last(t) + resid[1] = sol(t₀)[1] + resid[2] = sol(t₁)[1] - 1 + return nothing end -tspan = (0.0, 100.0) -u0 = [0.0, 1.0] -bvp = BVProblem(f, bc!, u0, tspan) +bvp1 = BVProblem(f1!, bc1!, u0, tspan) +@test SciMLBase.isinplace(bvp1) resid_f = Array{Float64}(undef, 2) -sol = solve(bvp, Shooting(Tsit5())) -bc!(resid_f, sol, nothing, sol.t) +sol = solve(bvp1, Shooting(Tsit5()); abstol = 1e-6, reltol = 1e-6) +@test SciMLBase.successful_retcode(sol) +bc1!(resid_f, sol, nothing, sol.t) +@test norm(resid_f) < 1e-6 + +# Out of Place +f1(u, p, t) = [u[2], -u[1]] + +function bc1(sol, p, t) + t₀, t₁ = first(t), last(t) + return [sol(t₀)[1], sol(t₁)[1] - 1] +end + +@test_throws SciMLBase.NonconformingFunctionsError BVProblem(f1!, bc1, u0, tspan) +@test_throws SciMLBase.NonconformingFunctionsError BVProblem(f1, bc1!, u0, tspan) + +bvp2 = BVProblem(f1, bc1, u0, tspan) +@test !SciMLBase.isinplace(bvp2) +sol = solve(bvp2, Shooting(Tsit5()); abstol = 1e-6, reltol = 1e-6) +@test SciMLBase.successful_retcode(sol) +resid_f = bc1(sol, nothing, sol.t) +@test norm(resid_f) < 1e-6 + +@info "Two Point BVProblem" # Not really but using that API + +# Inplace +function bc2!((resida, residb), (ua, ub), p) + resida[1] = ua[1] + residb[1] = ub[1] - 1 + return nothing +end + +bvp3 = TwoPointBVProblem(f1!, bc2!, u0, tspan; + bcresid_prototype = (Array{Float64}(undef, 1), Array{Float64}(undef, 1))) +@test SciMLBase.isinplace(bvp3) +sol = solve(bvp3, Shooting(Tsit5()); abstol = 1e-6, reltol = 1e-6) +@test SciMLBase.successful_retcode(sol) +resid_f = (Array{Float64, 1}(undef, 1), Array{Float64, 1}(undef, 1)) +bc2!(resid_f, (sol(tspan[1]), sol(tspan[2])), nothing) +@test_broken norm(resid_f) < 1e-6 +@test norm(resid_f) < 1e-4 + +# Out of Place +function bc2((ua, ub), p) + return ([ua[1]], [ub[1] - 1]) +end + +bvp4 = TwoPointBVProblem(f1, bc2, u0, tspan) +@test !SciMLBase.isinplace(bvp4) +sol = solve(bvp4, Shooting(Tsit5()); abstol = 1e-6, reltol = 1e-6) +@test SciMLBase.successful_retcode(sol) +resid_f = reduce(vcat, bc2((sol(tspan[1]), sol(tspan[2])), nothing)) @test norm(resid_f) < 1e-6 #Test for complex values u0 = [0.0, 1.0] .+ 1im -bvp = BVProblem(f, bc!, u0, tspan) +bvp = BVProblem(f1!, bc1!, u0, tspan) +resid_f = Array{ComplexF64}(undef, 2) +sol = solve(bvp, Shooting(Tsit5(); nlsolve = NewtonRaphson(; autodiff = AutoFiniteDiff())); + abstol = 1e-6, reltol = 1e-6) resid_f = Array{ComplexF64}(undef, 2) -sol = solve(bvp, Shooting(Tsit5(); nlsolve = NewtonRaphson(; autodiff = false))) -bc!(resid_f, sol, nothing, sol.t) +bc1!(resid_f, sol, nothing, sol.t) @test norm(resid_f) < 1e-6 diff --git a/test/vectorofvector_initials.jl b/test/vectorofvector_initials.jl index 247aac2c..afadc897 100644 --- a/test/vectorofvector_initials.jl +++ b/test/vectorofvector_initials.jl @@ -52,7 +52,8 @@ tspan = (0.0, T) prob = ODEProblem(TC!, u0, tspan, dt = 0.01) sol = solve(prob, Rodas4P(), reltol = 1e-12, abstol = 1e-12, saveat = 0.5) -#The BVP set up +# The BVP set up +# This is not really kind of Two-Point BVP we support. function bc_po!(residual, u, p, t) residual[1] = u[1][1] - u[end][1] residual[2] = u[1][2] - u[end][2] @@ -60,14 +61,10 @@ function bc_po!(residual, u, p, t) end #This is the part of the code that has problems -bvp1 = TwoPointBVProblem(TC!, bc_po!, sol.u, tspan) +bvp1 = BVProblem(TC!, bc_po!, sol.u, tspan) sol6 = solve(bvp1, MIRK6(); dt = 0.5) @test SciMLBase.successful_retcode(sol6.retcode) -@static if VERSION ≥ v"1.9" - # 1.6 runs without sparsity support. This takes over 2 hrs to run there :( - # Setup to test the mesh_selector! code - bvp1 = TwoPointBVProblem(TC!, bc_po!, zero(first(sol.u)), tspan) - sol6 = solve(bvp1, MIRK6(); dt = 0.1, abstol = 1e-16) - @test SciMLBase.successful_retcode(sol6.retcode) -end +bvp1 = BVProblem(TC!, bc_po!, zero(first(sol.u)), tspan) +sol6 = solve(bvp1, MIRK6(); dt = 0.1, abstol = 1e-16) +@test SciMLBase.successful_retcode(sol6.retcode)