From be2298f281a3f11fa4d06adaa2f28da7929f80ab Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Wed, 11 Oct 2023 19:12:11 -0400 Subject: [PATCH] Fix tests and make Multiple Shooting Type Stable --- src/algorithms.jl | 11 ++ src/solve/multiple_shooting.jl | 256 +++++++--------------------- src/sparse_jacobians.jl | 95 ++++++++++- src/types.jl | 11 +- test/mirk/ensemble.jl | 6 +- test/mirk/mirk_convergence_tests.jl | 4 +- test/misc/non_vector_inputs.jl | 2 +- test/shooting/orbital.jl | 36 ++-- 8 files changed, 203 insertions(+), 218 deletions(-) diff --git a/src/algorithms.jl b/src/algorithms.jl index faafae4e..cb825cd2 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -29,6 +29,17 @@ Significantly more stable than Single Shooting. grid_coarsening end +function concretize_jacobian_algorithm(alg::MultipleShooting, prob) + jac_alg = concrete_jacobian_algorithm(alg.jac_alg, prob, alg) + return MultipleShooting(alg.ode_alg, alg.nlsolve, jac_alg, alg.nshoots, + alg.grid_coarsening) +end + +function update_nshoots(alg::MultipleShooting, nshoots::Int) + return MultipleShooting(alg.ode_alg, alg.nlsolve, alg.jac_alg, nshoots, + alg.grid_coarsening) +end + function MultipleShooting(nshoots::Int, ode_alg; nlsolve = NewtonRaphson(), grid_coarsening = true, jac_alg = BVPJacobianAlgorithm()) @assert grid_coarsening isa Bool || grid_coarsening isa Function || diff --git a/src/solve/multiple_shooting.jl b/src/solve/multiple_shooting.jl index 4c9403c2..e9bb791f 100644 --- a/src/solve/multiple_shooting.jl +++ b/src/solve/multiple_shooting.jl @@ -1,33 +1,28 @@ -function __solve(prob::BVProblem, alg::MultipleShooting; odesolve_kwargs = (;), +function __solve(prob::BVProblem, _alg::MultipleShooting; odesolve_kwargs = (;), nlsolve_kwargs = (;), ensemblealg = EnsembleThreads(), verbose = true, kwargs...) @unpack f, tspan = prob - bc = prob.f.bc - has_initial_guess = prob.u0 isa AbstractVector{<:AbstractArray} - @set! alg.jac_alg = concrete_jacobian_algorithm(alg.jac_alg, prob, alg) - _u0 = has_initial_guess ? first(prob.u0) : prob.u0 - N, u0_size, nshoots, iip = length(_u0), size(_u0), alg.nshoots, isinplace(prob) - if prob.f.bcresid_prototype === nothing - if prob.problem_type isa TwoPointBVProblem - # This can only happen if the problem is !iip - bcresid_prototype = ArrayPartition(bc[1](_u0, prob.p), bc[2](_u0, prob.p)) - else - bcresid_prototype = similar(_u0) - end + + ig, T, N, Nig, u0 = __extract_problem_details(prob; dt = 0.1) + has_initial_guess = known(ig) + + bcresid_prototype, resid_size = __get_bcresid_prototype(prob, u0) + iip, bc, u0, u0_size = isinplace(prob), prob.f.bc, deepcopy(u0), size(u0) + + __alg = concretize_jacobian_algorithm(_alg, prob) + alg = if has_initial_guess && Nig != __alg.nshoots + 1 + verbose && + @warn "Initial guess length != `nshoots + 1`! Adapting to `nshoots = $(Nig - 1)`" + update_nshoots(__alg, Nig - 1) else - bcresid_prototype = prob.f.bcresid_prototype + __alg end + nshoots = alg.nshoots if prob.problem_type isa TwoPointBVProblem resida_len = length(bcresid_prototype.x[1]) residb_len = length(bcresid_prototype.x[2]) end - if has_initial_guess && length(prob.u0) != nshoots + 1 - nshoots = length(prob.u0) - 1 - verbose && - @warn "Initial guess length != `nshoots + 1`! Adapting to `nshoots = $(nshoots)`" - end - # We will use colored AD for this part! @views function solve_internal_odes!(resid_nodes, us, p, cur_nshoots, nodes) ts_ = Vector{Vector{typeof(first(tspan))}}(undef, cur_nshoots) @@ -52,7 +47,7 @@ function __solve(prob::BVProblem, alg::MultipleShooting; odesolve_kwargs = (;), ensemble_prob = EnsembleProblem(odeprob; prob_func, reduction, safetycopy = false, u_init = (; us = us_, ts = ts_, resid = resid_nodes)) - ensemble_sol = solve(ensemble_prob, alg.ode_alg, ensemblealg; odesolve_kwargs..., + ensemble_sol = __solve(ensemble_prob, alg.ode_alg, ensemblealg; odesolve_kwargs..., verbose, kwargs..., save_end = true, save_everystep = false, trajectories = cur_nshoots) @@ -66,7 +61,7 @@ function __solve(prob::BVProblem, alg::MultipleShooting; odesolve_kwargs = (;), # Just Recompute the last ODE Solution lastodeprob = ODEProblem{iip}(f, reshape(ub0, u0_size), (nodes[end - 1], nodes[end]), p) - sol_ode_last = solve(lastodeprob, alg.ode_alg; odesolve_kwargs..., verbose, + sol_ode_last = __solve(lastodeprob, alg.ode_alg; odesolve_kwargs..., verbose, kwargs..., save_everystep = false, saveat = (), save_end = true) ub = vec(sol_ode_last.u[end]) @@ -136,7 +131,9 @@ function __solve(prob::BVProblem, alg::MultipleShooting; odesolve_kwargs = (;), @views function jac_tp!(J::AbstractMatrix, us, p, resid_bc, resid_nodes::MaybeDiffCache, ode_jac_cache, bc_jac_cache::Tuple, ode_fn, bc_fn, cur_nshoot, nodes) - J isa SparseArrays.SparseMatrixCSC || fill!(J, 0) + # This is mostly a safety measure + fill!(J, 0) + J_bc = J[1:N, :] J_c = J[(N + 1):end, :] @@ -150,7 +147,7 @@ function __solve(prob::BVProblem, alg::MultipleShooting; odesolve_kwargs = (;), resida, residb = resid_bc.x J_bc[1:length(resida), 1:N] .= J_bc′[1:length(resida), 1:N] idxᵢ = (length(resida) + 1):(length(resida) + length(residb)) - J_bc[idxᵢ, (end - N + 1):end] .= J_bc′[idxᵢ, (end - N + 1):end] + J_bc[idxᵢ, (end - 2N + 1):(end - N)] .= J_bc′[idxᵢ, (end - N + 1):end] return nothing end @@ -158,6 +155,9 @@ function __solve(prob::BVProblem, alg::MultipleShooting; odesolve_kwargs = (;), @views function jac_mp!(J::AbstractMatrix, us, p, resid_bc, resid_nodes::MaybeDiffCache, ode_jac_cache, bc_jac_cache, ode_fn, bc_fn, cur_nshoot, nodes) + # This is mostly a safety measure + fill!(J, 0) + J_bc = J[1:N, :] J_c = J[(N + 1):end, :] @@ -174,15 +174,15 @@ function __solve(prob::BVProblem, alg::MultipleShooting; odesolve_kwargs = (;), # This gets all the nshoots except the final SingleShooting case all_nshoots = get_all_nshoots(alg.grid_coarsening, nshoots) - u_at_nodes, nodes = nothing, nothing + u_at_nodes, nodes = similar(u0, 0), typeof(first(tspan))[] for (i, cur_nshoot) in enumerate(all_nshoots) if i == 1 - nodes, u_at_nodes = multiple_shooting_initialize(prob, alg, has_initial_guess, - nshoots; odesolve_kwargs, verbose, kwargs...) + nodes, u_at_nodes = multiple_shooting_initialize(prob, alg, ig, nshoots; + odesolve_kwargs, verbose, kwargs...) else nodes, u_at_nodes = multiple_shooting_initialize(u_at_nodes, prob, alg, nodes, - cur_nshoot, all_nshoots[i - 1], has_initial_guess; odesolve_kwargs, verbose, + cur_nshoot, all_nshoots[i - 1]::Int, ig; odesolve_kwargs, verbose, kwargs...) end @@ -191,35 +191,23 @@ function __solve(prob::BVProblem, alg::MultipleShooting; odesolve_kwargs = (;), resid_nodes = __maybe_allocate_diffcache(resid_prototype.x[2], pickchunksize((cur_nshoot + 1) * N), alg.jac_alg.bc_diffmode) - if prob.problem_type isa TwoPointBVProblem - if alg.jac_alg.nonbc_diffmode isa AbstractSparseADType || - alg.jac_alg.bc_diffmode isa AbstractSparseADType - J_full, col_colorvec, row_colorvec, (J_c, J_bc_partial), col_colorvec_bc, row_colorvec_bc, = __generate_sparse_jacobian_prototype(alg, - prob.problem_type, bcresid_prototype, _u0, N, cur_nshoot) - end - elseif alg.jac_alg.nonbc_diffmode isa AbstractSparseADType - J_c, col_colorvec, row_colorvec, = __generate_sparse_jacobian_prototype(alg, - prob.problem_type, bcresid_prototype, _u0, N, cur_nshoot) + if alg.jac_alg.nonbc_diffmode isa AbstractSparseADType || + alg.jac_alg.bc_diffmode isa AbstractSparseADType + J_full, J_c, J_bc = __generate_sparse_jacobian_prototype(alg, prob.problem_type, + bcresid_prototype, u0, N, cur_nshoot) end ode_fn = (du, u) -> solve_internal_odes!(du, u, prob.p, cur_nshoot, nodes) - sd_ode = if alg.jac_alg.nonbc_diffmode isa AbstractSparseADType - PrecomputedJacobianColorvec(; jac_prototype = J_c, row_colorvec, col_colorvec) - else - NoSparsityDetection() - end + sd_ode = alg.jac_alg.nonbc_diffmode isa AbstractSparseADType ? + PrecomputedJacobianColorvec(J_c) : NoSparsityDetection() ode_jac_cache = sparse_jacobian_cache(alg.jac_alg.nonbc_diffmode, sd_ode, ode_fn, similar(u_at_nodes, cur_nshoot * N), u_at_nodes) bc_fn = (du, u) -> compute_bc_residual!(du, u, prob.p, cur_nshoot, nodes, resid_nodes) if prob.problem_type isa TwoPointBVProblem - sd_bc = if alg.jac_alg.bc_diffmode isa AbstractSparseADType - PrecomputedJacobianColorvec(; jac_prototype = J_bc_partial, - row_colorvec = row_colorvec_bc, col_colorvec = col_colorvec_bc) - else - NoSparsityDetection() - end + sd_bc = alg.jac_alg.bc_diffmode isa AbstractSparseADType ? + PrecomputedJacobianColorvec(J_bc) : NoSparsityDetection() bc_jac_cache_partial = sparse_jacobian_cache(alg.jac_alg.bc_diffmode, sd_bc, bc_fn, similar(bcresid_prototype), ArrayPartition(@view(u_at_nodes[1:N]), @@ -227,12 +215,10 @@ function __solve(prob::BVProblem, alg::MultipleShooting; odesolve_kwargs = (;), bc_jac_cache = (bc_jac_cache_partial, init_jacobian(bc_jac_cache_partial)) - jac_prototype = if alg.jac_alg.bc_diffmode isa AbstractSparseADType || - alg.jac_alg.nonbc_diffmode isa AbstractSparseADType + jac_prototype = if @isdefined(J_full) J_full else - # Dense AD being used! - fill!(similar(u_at_nodes, length(resid_prototype), length(u_at_nodes)), 0) + __zeros_like(u_at_nodes, length(resid_prototype), length(u_at_nodes)) end else sd_bc = alg.jac_alg.bc_diffmode isa AbstractSparseADType ? @@ -242,34 +228,42 @@ function __solve(prob::BVProblem, alg::MultipleShooting; odesolve_kwargs = (;), jac_prototype = vcat(init_jacobian(bc_jac_cache), init_jacobian(ode_jac_cache)) end + jac_fn = (J, us, p) -> jac!(J, us, p, similar(bcresid_prototype), resid_nodes, ode_jac_cache, bc_jac_cache, ode_fn, bc_fn, cur_nshoot, nodes) loss_function! = NonlinearFunction{true}((args...) -> loss!(args..., cur_nshoot, nodes); resid_prototype, jac = jac_fn, jac_prototype) nlprob = NonlinearProblem(loss_function!, u_at_nodes, prob.p) - sol_nlsolve = solve(nlprob, alg.nlsolve; nlsolve_kwargs..., verbose, kwargs...) - u_at_nodes = sol_nlsolve.u + sol_nlsolve = __solve(nlprob, alg.nlsolve; nlsolve_kwargs..., verbose, kwargs...) + # u_at_nodes = sol_nlsolve.u end single_shooting_prob = remake(prob; u0 = reshape(u_at_nodes[1:N], u0_size)) - return solve(single_shooting_prob, Shooting(alg.ode_alg; alg.nlsolve); odesolve_kwargs, - nlsolve_kwargs, verbose, kwargs...) + return __solve(single_shooting_prob, Shooting(alg.ode_alg; alg.nlsolve); + odesolve_kwargs, nlsolve_kwargs, verbose, kwargs...) end -@views function multiple_shooting_initialize(prob, alg::MultipleShooting, has_initial_guess, +@views function multiple_shooting_initialize(prob, alg::MultipleShooting, ::True, nshoots; odesolve_kwargs = (;), verbose = true, kwargs...) @unpack f, u0, tspan, p = prob @unpack ode_alg = alg nodes = range(tspan[1], tspan[2]; length = nshoots + 1) - N = has_initial_guess ? length(first(u0)) : length(u0) + N = length(first(u0)) - if has_initial_guess - u_at_nodes = similar(first(u0), (nshoots + 1) * N) - recursive_flatten!(u_at_nodes, u0) - return nodes, u_at_nodes - end + u_at_nodes = similar(first(u0), (nshoots + 1) * N) + recursive_flatten!(u_at_nodes, u0) + return nodes, u_at_nodes +end + +@views function multiple_shooting_initialize(prob, alg::MultipleShooting, ::False, + nshoots; odesolve_kwargs = (;), verbose = true, kwargs...) + @unpack f, u0, tspan, p = prob + @unpack ode_alg = alg + + nodes = range(tspan[1], tspan[2]; length = nshoots + 1) + N = length(u0) # Ensures type stability in case the parameters are dual numbers if !(typeof(p) <: SciMLBase.NullParameters) @@ -283,12 +277,13 @@ end # Assumes no initial guess for now start_prob = ODEProblem{isinplace(prob)}(f, u0, tspan, p) - sol = solve(start_prob, ode_alg; odesolve_kwargs..., verbose, kwargs..., saveat = nodes) + sol = __solve(start_prob, ode_alg; odesolve_kwargs..., verbose, kwargs..., + saveat = nodes) if SciMLBase.successful_retcode(sol) - u_at_nodes[1:N] .= sol.u[1] + u_at_nodes[1:N] .= vec(sol.u[1]) for i in 2:(nshoots + 1) - u_at_nodes[(N + (i - 2) * N) .+ (1:N)] .= sol.u[i] + u_at_nodes[(N + (i - 2) * N) .+ (1:N)] .= vec(sol.u[i]) end else @warn "Initialization using odesolve failed. Initializing using 0s. It is \ @@ -300,10 +295,10 @@ end end @views function multiple_shooting_initialize(u_at_nodes_prev, prob, alg, prev_nodes, - nshoots, old_nshoots, has_initial_guess; odesolve_kwargs = (;), kwargs...) + nshoots, old_nshoots, ig; odesolve_kwargs = (;), kwargs...) @unpack f, u0, tspan, p = prob nodes = range(tspan[1], tspan[2]; length = nshoots + 1) - N = has_initial_guess ? length(first(u0)) : length(u0) + N = known(ig) ? length(first(u0)) : length(u0) u_at_nodes = similar(u_at_nodes_prev, N + nshoots * N) u_at_nodes[1:N] .= u_at_nodes_prev[1:N] @@ -332,8 +327,8 @@ end ustart = u_at_nodes_prev[idxs_prev] odeprob = ODEProblem(f, ustart, (t0, tstop), p) - odesol = solve(odeprob, alg.ode_alg; odesolve_kwargs..., kwargs..., saveat = (), - save_end = true) + odesol = __solve(odeprob, alg.ode_alg; odesolve_kwargs..., kwargs..., + saveat = (), save_end = true) u_at_nodes[idxs] .= odesol.u[end] end @@ -361,124 +356,3 @@ end @assert !(1 in nshoots_vec) return nshoots_vec end - -""" - __generate_sparse_jacobian_prototype(::MultipleShooting, _, _, u0, N::Int, - nshoots::Int) - __generate_sparse_jacobian_prototype(::MultipleShooting, ::TwoPointBVProblem, - bcresid_prototype, u0, N::Int, nshoots::Int) - -For a Multi-Point Problem, returns the Jacobian Prototype for the Sparse Part. For a Two- -Point Problem, returns the Jacobian Prototype for the Entire Jacobian. - -Also returns the column and row color vectors for the Sparse Non-BC Part Jacobian. - -Returns the column and row color vectors for the Sparse BC Part Jacobian (if computed). - -Also returns the indices `Is` and `Js` used to construct the Sparse Jacobian. -""" -function __generate_sparse_jacobian_prototype(::MultipleShooting, _, _, u0, N::Int, - nshoots::Int) - # Sparse for Stitching solution together - Is = Vector{Int64}(undef, (N^2 + N) * nshoots) - Js = Vector{Int64}(undef, (N^2 + N) * nshoots) - - idx = 1 - for i in 1:nshoots - for (i₁, i₂) in Iterators.product(1:N, 1:N) - Is[idx] = i₁ + ((i - 1) * N) - Js[idx] = i₂ + ((i - 1) * N) - idx += 1 - end - Is[idx:(idx + N - 1)] .= (1:N) .+ ((i - 1) * N) - Js[idx:(idx + N - 1)] .= (1:N) .+ (i * N) - idx += N - end - - J_c = sparse(adapt(parameterless_type(u0), Is), adapt(parameterless_type(u0), Js), - similar(u0, length(Is))) - - col_colorvec = Vector{Int}(undef, N * (nshoots + 1)) - for i in eachindex(col_colorvec) - col_colorvec[i] = mod1(i, 2 * N) - end - row_colorvec = Vector{Int}(undef, N * nshoots) - for i in eachindex(row_colorvec) - row_colorvec[i] = mod1(i, 2 * N) - end - - return J_c, col_colorvec, row_colorvec, (J_c, nothing), nothing, nothing, Is, Js -end - -function __generate_sparse_jacobian_prototype(alg::MultipleShooting, ::TwoPointBVProblem, - bcresid_prototype, u0, N::Int, nshoots::Int) - resida, residb = bcresid_prototype.x - # Sparse for Stitching solution together - L = N * length(resida) + (N^2 + N) * nshoots + N * length(residb) - Is = Vector{Int64}(undef, L) - Js = Vector{Int64}(undef, L) - - idx = 1 - for row in 1:length(resida) - for j in 1:N - Is[idx] = row - Js[idx] = j - idx += 1 - end - end - for row in 1:length(residb) - for j in 1:N - Is[idx] = length(resida) + row - Js[idx] = j + (nshoots * N) - idx += 1 - end - end - J_c, col_colorvec, row_colorvec, _, _, _, Is′, Js′ = __generate_sparse_jacobian_prototype(alg, - nothing, nothing, u0, N, nshoots) - for (i, j) in zip(Is′, Js′) - Is[idx] = length(resida) + length(residb) + i - Js[idx] = j - idx += 1 - end - - col_colorvec_bc = Vector{Int}(undef, 2N) - row_colorvec_bc = Vector{Int}(undef, length(resida) + length(residb)) - col_colorvec_bc[1:N] .= 1:N - col_colorvec_bc[(N + 1):end] .= 1:N - for i in 1:max(length(resida), length(residb)) - if i ≤ length(resida) - row_colorvec_bc[i] = i - end - if i ≤ length(residb) - row_colorvec_bc[i + length(resida)] = i - end - end - - J = sparse(adapt(parameterless_type(u0), Is), adapt(parameterless_type(u0), Js), - similar(u0, length(Is))) - - Is_bc = Vector{Int64}(undef, N^2) - Js_bc = Vector{Int64}(undef, N^2) - idx = 1 - for i in 1:length(resida) - for j in 1:N - Is_bc[idx] = i - Js_bc[idx] = j - idx += 1 - end - end - for i in 1:length(residb) - for j in 1:N - Is_bc[idx] = i + length(resida) - Js_bc[idx] = j + N - idx += 1 - end - end - - J_bc = sparse(adapt(parameterless_type(u0), Is_bc), - adapt(parameterless_type(u0), Js_bc), - similar(u0, length(Is_bc))) - - return (J, col_colorvec, row_colorvec, (J_c, J_bc), col_colorvec_bc, row_colorvec_bc, - Is, Js) -end diff --git a/src/sparse_jacobians.jl b/src/sparse_jacobians.jl index 3be9d4ce..e9c88abb 100644 --- a/src/sparse_jacobians.jl +++ b/src/sparse_jacobians.jl @@ -2,7 +2,7 @@ function _sparse_like(I, J, x::AbstractArray, m = maximum(I), n = maximum(J)) I′ = adapt(parameterless_type(x), I) J′ = adapt(parameterless_type(x), J) - V = similar(x, length(I)) + V = __ones_like(x, length(I)) return sparse(I′, J′, V, m, n) end @@ -21,6 +21,11 @@ end col_colorvec end +Base.size(M::ColoredMatrix, args...) = size(M.M, args...) +Base.eltype(M::ColoredMatrix) = eltype(M.M) + +ColoredMatrix() = ColoredMatrix(nothing, nothing, nothing) + function SparseDiffTools.PrecomputedJacobianColorvec(M::ColoredMatrix) return PrecomputedJacobianColorvec(; jac_prototype = M.M, M.row_colorvec, M.col_colorvec) @@ -110,3 +115,91 @@ function __generate_sparse_jacobian_prototype(::MIRKCache, ::TwoPointBVProblem, end # For Multiple Shooting +""" + __generate_sparse_jacobian_prototype(::MultipleShooting, ::StandardBVProblem, + bcresid_prototype, u0, N::Int, nshoots::Int) + __generate_sparse_jacobian_prototype(::MultipleShooting, ::TwoPointBVProblem, + bcresid_prototype, u0, N::Int, nshoots::Int) + +Returns a 3-Tuple: + +* Entire Jacobian Prototype (if Two-Point Problem) else `nothing`. +* Sparse Non-BC Part Jacobian Prototype along with the column and row color vectors. +* Sparse BC Part Jacobian Prototype along with the column and row color vectors (if + Two-Point Problem) else `nothing`. +""" +function __generate_sparse_jacobian_prototype(::MultipleShooting, ::StandardBVProblem, + bcresid_prototype, u0, N::Int, nshoots::Int) + Is = Vector{Int}(undef, (N^2 + N) * nshoots) + Js = Vector{Int}(undef, (N^2 + N) * nshoots) + + idx = 1 + for i in 1:nshoots + for (i₁, i₂) in Iterators.product(1:N, 1:N) + Is[idx] = i₁ + ((i - 1) * N) + Js[idx] = i₂ + ((i - 1) * N) + idx += 1 + end + Is[idx:(idx + N - 1)] .= (1:N) .+ ((i - 1) * N) + Js[idx:(idx + N - 1)] .= (1:N) .+ (i * N) + idx += N + end + + J_c = _sparse_like(Is, Js, u0) + + col_colorvec = Vector{Int}(undef, size(J_c, 2)) + for i in eachindex(col_colorvec) + col_colorvec[i] = mod1(i, 2N) + end + row_colorvec = Vector{Int}(undef, size(J_c, 1)) + for i in eachindex(row_colorvec) + row_colorvec[i] = mod1(i, 2N) + end + + return nothing, ColoredMatrix(J_c, row_colorvec, col_colorvec), nothing +end + +function __generate_sparse_jacobian_prototype(alg::MultipleShooting, ::TwoPointBVProblem, + bcresid_prototype::ArrayPartition, u0, N::Int, nshoots::Int) + resida, residb = bcresid_prototype.x + L₁, L₂ = length(resida), length(residb) + + _, J_c, _ = __generate_sparse_jacobian_prototype(alg, StandardBVProblem(), + bcresid_prototype, u0, N, nshoots) + + Is_bc = Vector{Int}(undef, (L₁ + L₂) * N) + Js_bc = Vector{Int}(undef, (L₁ + L₂) * N) + idx = 1 + for i in 1:L₁, j in 1:N + Is_bc[idx] = i + Js_bc[idx] = j + idx += 1 + end + for i in 1:L₂, j in 1:N + Is_bc[idx] = i + L₁ + Js_bc[idx] = j + N + idx += 1 + end + + col_colorvec_bc = Vector{Int}(undef, 2N) + row_colorvec_bc = Vector{Int}(undef, L₁ + L₂) + col_colorvec_bc[1:N] .= 1:N + col_colorvec_bc[(N + 1):end] .= 1:N + for i in 1:max(L₁, L₂) + i ≤ L₁ && (row_colorvec_bc[i] = i) + i ≤ L₂ && (row_colorvec_bc[i + L₁] = i) + end + + J_bc = ColoredMatrix(_sparse_like(Is_bc, Js_bc, bcresid_prototype), row_colorvec_bc, + col_colorvec_bc) + + J_full = _sparse_like(Int[], Int[], u0, size(J_bc, 1) + size(J_c, 1), + size(J_c, 2)) + + J_full[(L₁ + L₂ + 1):end, :] .= J_c.M + J_full[1:L₁, 1:N] .= J_bc.M[1:L₁, 1:N] + J_full[(L₁ + 1):(L₁ + L₂), (end - 2N + 1):(end - N)] .= J_bc.M[(L₁ + 1):(L₁ + L₂), + (N + 1):(2N)] + + return J_full, J_c, J_bc +end diff --git a/src/types.jl b/src/types.jl index 02c2bf7a..f38898af 100644 --- a/src/types.jl +++ b/src/types.jl @@ -42,13 +42,14 @@ end function BVPJacobianAlgorithm(diffmode = missing; nonbc_diffmode = missing, bc_diffmode = missing) if diffmode !== missing - @assert nonbc_diffmode === missing && bc_diffmode === missing + bc_diffmode = bc_diffmode === missing ? diffmode : bc_diffmode + nonbc_diffmode = nonbc_diffmode === missing ? diffmode : nonbc_diffmode return BVPJacobianAlgorithm(diffmode, diffmode, diffmode) else diffmode = nothing bc_diffmode = bc_diffmode === missing ? nothing : bc_diffmode nonbc_diffmode = nonbc_diffmode === missing ? nothing : nonbc_diffmode - return BVPJacobianAlgorithm(bc_diffmode, nonbc_diffmode, nonbc_diffmode) + return BVPJacobianAlgorithm(bc_diffmode, nonbc_diffmode, diffmode) end end @@ -88,6 +89,12 @@ function concrete_jacobian_algorithm(jac_alg::BVPJacobianAlgorithm, ::TwoPointBV return BVPJacobianAlgorithm(bc_diffmode, nonbc_diffmode, diffmode) end +# This can cause Type Instability +function concretize_jacobian_algorithm(alg, prob) + @set! alg.jac_alg = concrete_jacobian_algorithm(alg.jac_alg, prob, alg) + return alg +end + function MIRKJacobianComputationAlgorithm(diffmode = missing; collocation_diffmode = missing, bc_diffmode = missing) Base.depwarn("`MIRKJacobianComputationAlgorithm` has been deprecated in favor of \ diff --git a/test/mirk/ensemble.jl b/test/mirk/ensemble.jl index e0771993..d03c2d94 100644 --- a/test/mirk/ensemble.jl +++ b/test/mirk/ensemble.jl @@ -19,9 +19,9 @@ bvp = BVProblem(ode!, bc!, initial_guess, tspan, p) 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())] + jac_algs = [BVPJacobianAlgorithm(), + BVPJacobianAlgorithm(; bc_diffmode = AutoFiniteDiff(), + nonbc_diffmode = AutoSparseFiniteDiff())] for jac_alg in jac_algs # Not sure why it is throwing so many warnings sol = solve(ensemble_prob, solver(; jac_alg); trajectories = 10, dt = 0.1) diff --git a/test/mirk/mirk_convergence_tests.jl b/test/mirk/mirk_convergence_tests.jl index 54707c75..cdfd879f 100644 --- a/test/mirk/mirk_convergence_tests.jl +++ b/test/mirk/mirk_convergence_tests.jl @@ -115,8 +115,8 @@ end u0 = MVector{2}([pi / 2, pi / 2]) bvp1 = BVProblem(simplependulum!, bc_pendulum!, u0, tspan) -jac_alg = MIRKJacobianComputationAlgorithm(; bc_diffmode = AutoFiniteDiff(), - collocation_diffmode = AutoSparseFiniteDiff()) +jac_alg = BVPJacobianAlgorithm(; bc_diffmode = AutoFiniteDiff(), + nonbc_diffmode = AutoSparseFiniteDiff()) # Using ForwardDiff might lead to Cache expansion warnings @test_nowarn solve(bvp1, MIRK2(; jac_alg); dt = 0.005) diff --git a/test/misc/non_vector_inputs.jl b/test/misc/non_vector_inputs.jl index b7332505..170d48f9 100644 --- a/test/misc/non_vector_inputs.jl +++ b/test/misc/non_vector_inputs.jl @@ -58,5 +58,5 @@ probs = [ end end - # TODO: Multiple Shooting + # FIXME: Add Multiple Shooting here once it supports non-vector inputs end diff --git a/test/shooting/orbital.jl b/test/shooting/orbital.jl index 65f644fc..d28c47d4 100644 --- a/test/shooting/orbital.jl +++ b/test/shooting/orbital.jl @@ -64,25 +64,24 @@ cur_bc_2point_b! = (resid, sol, p) -> bc!_generator_2p_b(resid, sol, init_val) resid_f = Array{Float64}(undef, 6) 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) 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) + nlsolve = TrustRegion(; autodiff) + @time sol = solve(bvp, Shooting(DP5(); nlsolve); force_dtmin = true, abstol = 1e-13, + reltol = 1e-13, verbose = false) cur_bc!(resid_f, sol, nothing, sol.t) - @test norm(resid_f, Inf) < TestTol + @info "Single Shooting Lambert's Problem: $(norm(resid_f, Inf))" + @test norm(resid_f, Inf) < 0.005 jac_alg = BVPJacobianAlgorithm(; nonbc_diffmode = autodiff) - @time sol = solve(bvp, MultipleShooting(10, DP5(); nlsolve, jac_alg); abstol = 1e-6, - reltol = 1e-6) - @test SciMLBase.successful_retcode(sol) + @time sol = solve(bvp, MultipleShooting(10, AutoVern7(Rodas5P()); nlsolve, jac_alg); + abstol = 1e-6, reltol = 1e-6, verbose = false) cur_bc!(resid_f, sol, nothing, sol.t) - @test norm(resid_f, Inf) < 1e-6 + @info "Multiple Shooting Lambert's Problem: $(norm(resid_f, Inf))" + @test norm(resid_f, Inf) < 0.005 end ### Using the TwoPoint BVP Structure @@ -91,18 +90,19 @@ bvp = TwoPointBVProblem(orbital!, (cur_bc_2point_a!, cur_bc_2point_b!), y0, tspa for autodiff in (AutoForwardDiff(), AutoFiniteDiff(; fdtype = Val(:central)), AutoSparseForwardDiff(), AutoFiniteDiff(; fdtype = Val(:forward)), AutoSparseFiniteDiff()) - nlsolve = NewtonRaphson(; autodiff) + nlsolve = TrustRegion(; autodiff) @time sol = solve(bvp, Shooting(DP5(); nlsolve); force_dtmin = true, abstol = 1e-13, - reltol = 1e-13) + reltol = 1e-13, verbose = false) cur_bc_2point_a!(resid_f_2p[1], sol(t0), nothing) cur_bc_2point_b!(resid_f_2p[2], sol(t1), nothing) - @test norm(reduce(vcat, resid_f_2p), Inf) < TestTol + @info "Single Shooting Lambert's Problem: $(norm(reduce(vcat, resid_f_2p), Inf))" + @test norm(reduce(vcat, resid_f_2p), Inf) < 0.005 - jac_alg = BVPJacobianAlgorithm(; nonbc_diffmode = autodiff) - @time sol = solve(bvp, MultipleShooting(10, DP5(); nlsolve, jac_alg); abstol = 1e-6, - reltol = 1e-6) - @test SciMLBase.successful_retcode(sol) + jac_alg = BVPJacobianAlgorithm(; nonbc_diffmode = autodiff, bc_diffmode = autodiff) + @time sol = solve(bvp, MultipleShooting(10, AutoVern7(Rodas5P()); nlsolve, jac_alg); + abstol = 1e-6, reltol = 1e-6, verbose = false) cur_bc_2point_a!(resid_f_2p[1], sol(t0), nothing) cur_bc_2point_b!(resid_f_2p[2], sol(t1), nothing) - @test norm(reduce(vcat, resid_f_2p), Inf) < TestTol + @info "Multiple Shooting Lambert's Problem: $(norm(reduce(vcat, resid_f_2p), Inf))" + @test norm(reduce(vcat, resid_f_2p), Inf) < 0.005 end