From eb3a028da4936da0ec314d90ab7707191bcadcf6 Mon Sep 17 00:00:00 2001 From: Avik Pal Date: Mon, 25 Dec 2023 18:45:55 -0500 Subject: [PATCH] Handle Dynamic Dispatches in Multiple Shooting --- ext/BoundaryValueDiffEqODEInterfaceExt.jl | 4 +- src/algorithms.jl | 10 - src/solve/multiple_shooting.jl | 62 +++-- test/shooting/basic_problems.jl | 305 +++++++++++----------- test/shooting/nonlinear_least_squares.jl | 69 ++++- test/shooting/orbital.jl | 31 +-- 6 files changed, 248 insertions(+), 233 deletions(-) diff --git a/ext/BoundaryValueDiffEqODEInterfaceExt.jl b/ext/BoundaryValueDiffEqODEInterfaceExt.jl index d15fd3ff..e478f298 100644 --- a/ext/BoundaryValueDiffEqODEInterfaceExt.jl +++ b/ext/BoundaryValueDiffEqODEInterfaceExt.jl @@ -89,7 +89,7 @@ function __solve(prob::BVProblem, alg::BVPM2; dt = 0.0, reltol = 1e-3, kwargs... bvpm2_destroy(obj) bvpm2_destroy(sol) - return SciMLBase.build_solution(prob, ivpsol) + return SciMLBase.build_solution(prob, ivpsol, nothing) end #------- @@ -180,7 +180,7 @@ function __solve(prob::BVProblem, alg::BVPSOL; maxiters = 1000, reltol = 1e-3, d map(x -> reshape(convert(Vector{eltype(u0_)}, x), u0_size), eachcol(sol_x)); retcode = retcode ≥ 0 ? ReturnCode.Success : ReturnCode.Failure, stats) - return SciMLBase.build_solution(prob, ivpsol) + return SciMLBase.build_solution(prob, ivpsol, nothing) end #------- diff --git a/src/algorithms.jl b/src/algorithms.jl index 12aaa6fc..738ce066 100644 --- a/src/algorithms.jl +++ b/src/algorithms.jl @@ -108,16 +108,6 @@ Significantly more stable than Single Shooting. grid_coarsening end -# function Base.show(io::IO, alg::MultipleShooting) -# print(io, "MultipleShooting(") -# modifiers = String[] -# alg.nlsolve !== nothing && push!(modifiers, "nlsolve = $(alg.nlsolve)") -# alg.jac_alg !== nothing && push!(modifiers, "jac_alg = $(alg.jac_alg)") -# alg.ode_alg !== nothing && push!(modifiers, "ode_alg = $(__nameof(alg.ode_alg))()") -# print(io, join(modifiers, ", ")) -# print(io, ")") -# 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, diff --git a/src/solve/multiple_shooting.jl b/src/solve/multiple_shooting.jl index f26f9924..5a841cbe 100644 --- a/src/solve/multiple_shooting.jl +++ b/src/solve/multiple_shooting.jl @@ -1,8 +1,11 @@ function __solve(prob::BVProblem, _alg::MultipleShooting; odesolve_kwargs = (;), nlsolve_kwargs = (;), ensemblealg = EnsembleThreads(), verbose = true, kwargs...) - @unpack f, tspan = prob + (; f, tspan) = prob - @assert (ensemblealg isa EnsembleSerial)||(ensemblealg isa EnsembleThreads) "Currently MultipleShooting only supports `EnsembleSerial` and `EnsembleThreads`!" + if !(ensemblealg isa EnsembleSerial) && !(ensemblealg isa EnsembleThreads) + throw(ArgumentError("Currently MultipleShooting only supports `EnsembleSerial` and \ + `EnsembleThreads`!")) + end ig, T, N, Nig, u0 = __extract_problem_details(prob; dt = 0.1) has_initial_guess = _unwrap_val(ig) @@ -30,11 +33,9 @@ function __solve(prob::BVProblem, _alg::MultipleShooting; odesolve_kwargs = (;), internal_ode_kwargs = (; verbose, kwargs..., odesolve_kwargs..., save_end = true) - function solve_internal_odes!(resid_nodes::T1, us::T2, p::T3, cur_nshoot::Int, - nodes::T4, odecache::C) where {T1, T2, T3, T4, C} - return __multiple_shooting_solve_internal_odes!(resid_nodes, us, cur_nshoot, - odecache, nodes, u0_size, N, ensemblealg) - end + solve_internal_odes! = @closure (resid_nodes, us, p, cur_nshoot, nodes, + odecache) -> __multiple_shooting_solve_internal_odes!(resid_nodes, us, cur_nshoot, + odecache, nodes, u0_size, N, ensemblealg, tspan) # This gets all the nshoots except the final SingleShooting case all_nshoots = __get_all_nshoots(alg.grid_coarsening, nshoots) @@ -96,7 +97,7 @@ function __solve_nlproblem!(::TwoPointBVProblem, alg::MultipleShooting, bcresid_ resid_prototype = vcat(bcresid_prototype[1], similar(u_at_nodes, cur_nshoot * N), bcresid_prototype[2]) - loss_fn = (du, u, p) -> __multiple_shooting_2point_loss!(du, u, p, cur_nshoot, + loss_fn = @closure (du, u, p) -> __multiple_shooting_2point_loss!(du, u, p, cur_nshoot, nodes, prob, solve_internal_odes!, resida_len, residb_len, N, bca, bcb, ode_cache_loss_fn) @@ -112,19 +113,18 @@ function __solve_nlproblem!(::TwoPointBVProblem, alg::MultipleShooting, bcresid_ jac_cache, alg.jac_alg.diffmode, alg.ode_alg, cur_nshoot, u0; internal_ode_kwargs...) - loss_fnₚ = (du, u) -> __multiple_shooting_2point_loss!(du, u, prob.p, cur_nshoot, - nodes, prob, solve_internal_odes!, resida_len, residb_len, N, bca, bcb, + loss_fnₚ = @closure (du, u) -> __multiple_shooting_2point_loss!(du, u, prob.p, + cur_nshoot, nodes, prob, solve_internal_odes!, resida_len, residb_len, N, bca, bcb, ode_cache_jac_fn) - jac_fn = (J, u, p) -> __multiple_shooting_2point_jacobian!(J, u, p, jac_cache, + jac_fn = @closure (J, u, p) -> __multiple_shooting_2point_jacobian!(J, u, p, jac_cache, loss_fnₚ, resid_prototype_cached, alg) - loss_function! = NonlinearFunction{true}(loss_fn; resid_prototype, jac = jac_fn, - jac_prototype) + loss_function! = __unsafe_nonlinearfunction{true}(loss_fn; resid_prototype, + jac = jac_fn, jac_prototype) # NOTE: u_at_nodes is updated inplace - nlprob = (M != N ? NonlinearLeastSquaresProblem : NonlinearProblem)(loss_function!, - u_at_nodes, prob.p) + nlprob = __internal_nlsolve_problem(prob, M, N, loss_function!, u_at_nodes, prob.p) __solve(nlprob, alg.nlsolve; kwargs..., alias_u0 = true) return nothing @@ -144,7 +144,7 @@ function __solve_nlproblem!(::StandardBVProblem, alg::MultipleShooting, bcresid_ resid_nodes = __maybe_allocate_diffcache(__resid_nodes, pickchunksize((cur_nshoot + 1) * N), alg.jac_alg.bc_diffmode) - loss_fn = (du, u, p) -> __multiple_shooting_mpoint_loss!(du, u, p, cur_nshoot, + loss_fn = @closure (du, u, p) -> __multiple_shooting_mpoint_loss!(du, u, p, cur_nshoot, nodes, prob, solve_internal_odes!, resid_len, N, f, bc, u0_size, prob.tspan, alg.ode_alg, u0, ode_cache_loss_fn) @@ -169,22 +169,21 @@ function __solve_nlproblem!(::StandardBVProblem, alg::MultipleShooting, bcresid_ jac_prototype = vcat(init_jacobian(bc_jac_cache), init_jacobian(ode_jac_cache)) # Define the functions now - ode_fn = (du, u) -> solve_internal_odes!(du, u, prob.p, cur_nshoot, nodes, + ode_fn = @closure (du, u) -> solve_internal_odes!(du, u, prob.p, cur_nshoot, nodes, ode_cache_ode_jac_fn) - bc_fn = (du, u) -> __multiple_shooting_mpoint_loss_bc!(du, u, prob.p, cur_nshoot, nodes, - prob, solve_internal_odes!, N, f, bc, u0_size, prob.tspan, alg.ode_alg, u0, - ode_cache_bc_jac_fn) + bc_fn = @closure (du, u) -> __multiple_shooting_mpoint_loss_bc!(du, u, prob.p, + cur_nshoot, nodes, prob, solve_internal_odes!, N, f, bc, u0_size, prob.tspan, + alg.ode_alg, u0, ode_cache_bc_jac_fn) - jac_fn = (J, u, p) -> __multiple_shooting_mpoint_jacobian!(J, u, p, + jac_fn = @closure (J, u, p) -> __multiple_shooting_mpoint_jacobian!(J, u, p, similar(bcresid_prototype), resid_nodes, ode_jac_cache, bc_jac_cache, ode_fn, bc_fn, alg, N, M) - loss_function! = NonlinearFunction{true}(loss_fn; resid_prototype, jac = jac_fn, - jac_prototype) + loss_function! = __unsafe_nonlinearfunction{true}(loss_fn; resid_prototype, + jac_prototype, jac = jac_fn) # NOTE: u_at_nodes is updated inplace - nlprob = (M != N ? NonlinearLeastSquaresProblem : NonlinearProblem)(loss_function!, - u_at_nodes, prob.p) + nlprob = __internal_nlsolve_problem(prob, M, N, loss_function!, u_at_nodes, prob.p) __solve(nlprob, alg.nlsolve; kwargs..., alias_u0 = true) return nothing @@ -224,7 +223,7 @@ end # Not using `EnsembleProblem` since it is hard to initialize the cache and stuff function __multiple_shooting_solve_internal_odes!(resid_nodes, us, cur_nshoots::Int, - odecache, nodes, u0_size, N::Int, ::EnsembleSerial) + odecache, nodes, u0_size, N::Int, ::EnsembleSerial, tspan) ts_ = Vector{Vector{typeof(first(tspan))}}(undef, cur_nshoots) us_ = Vector{Vector{typeof(us)}}(undef, cur_nshoots) @@ -242,7 +241,7 @@ function __multiple_shooting_solve_internal_odes!(resid_nodes, us, cur_nshoots:: end function __multiple_shooting_solve_internal_odes!(resid_nodes, us, cur_nshoots::Int, - odecache::Vector, nodes, u0_size, N::Int, ::EnsembleThreads) + odecache::Vector, nodes, u0_size, N::Int, ::EnsembleThreads, tspan) ts_ = Vector{Vector{typeof(first(tspan))}}(undef, cur_nshoots) us_ = Vector{Vector{typeof(us)}}(undef, cur_nshoots) @@ -364,7 +363,7 @@ end # NOTE: We don't check `u0 isa Function` since `u0` in-principle can be a callable # struct - u0_ = u0 isa AbstractArray ? u0 : [__initial_guess(u0, prob.p, t) for t in nodes] + u0_ = u0 isa VectorOfArray ? u0 : [__initial_guess(u0, prob.p, t) for t in nodes] N = length(first(u0_)) u_at_nodes = similar(first(u0_), (nshoots + 1) * N) @@ -399,14 +398,13 @@ end sol = solve!(odecache) if SciMLBase.successful_retcode(sol) - res = sol(nodes).u for i in 1:length(nodes) - u_at_nodes[(i - 1) * N .+ (1:N)] .= vec(res[i]) + u_at_nodes[(i - 1) * N .+ (1:N)] .= vec(sol(nodes[i])) end else @warn "Initialization using odesolve failed. Initializing using 0s. It is \ recommended to provide an initial guess function via \ - `u0 = (p, t)` or `u0 = (t)` in this case." + `u0 = (p, t)` in this case." fill!(u_at_nodes, 0) end @@ -478,4 +476,4 @@ end end @assert !(1 in nshoots_vec) return nshoots_vec -end +end \ No newline at end of file diff --git a/test/shooting/basic_problems.jl b/test/shooting/basic_problems.jl index fb4950ca..5fd51adc 100644 --- a/test/shooting/basic_problems.jl +++ b/test/shooting/basic_problems.jl @@ -1,17 +1,14 @@ using BoundaryValueDiffEq, LinearAlgebra, LinearSolve, OrdinaryDiffEq, Test, JET -# FIXME: Reenable the Multiple Shooting Tests - @testset "Basic Shooting Tests" begin SOLVERS = [ Shooting(Tsit5(), NewtonRaphson(; autodiff = AutoFiniteDiff())), Shooting(Tsit5(), NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 2))), Shooting(Tsit5()), - # MultipleShooting(10, Tsit5(), - # NewtonRaphson(; autodiff = AutoFiniteDiff())), - # MultipleShooting(10, Tsit5(), - # NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 2))), - # MultipleShooting(10, Tsit5()), + MultipleShooting(10, Tsit5(), NewtonRaphson(; autodiff = AutoFiniteDiff())), + MultipleShooting(10, Tsit5(), + NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 2))), + MultipleShooting(10, Tsit5()), ] JET_SKIP = [false, false, true, false, false, true] JET_BROKEN = [false, false, false, false, false, false] @@ -146,12 +143,9 @@ end SOLVERS = [ Shooting(Vern7(), NewtonRaphson(; autodiff = AutoFiniteDiff())), Shooting(Vern7()), - # MultipleShooting(10, Vern7(), - # NewtonRaphson(; autodiff = AutoFiniteDiff())), - # MultipleShooting(10, Vern7()), + MultipleShooting(10, Vern7(), NewtonRaphson(; autodiff = AutoFiniteDiff())), + MultipleShooting(10, Vern7()), ] - JET_SKIP = [false, true, false, true] - JET_BROKEN = [false, false, false, false] function f1!(du, u, p, t) du[1] = u[2] @@ -219,8 +213,10 @@ end @info "Flow in a Channel" for solver in [ - Shooting(AutoTsit5(Rosenbrock23())), - # MultipleShooting(10, AutoTsit5(Rosenbrock23())), + Shooting(AutoTsit5(Rosenbrock23()), + NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 8))), + MultipleShooting(10, AutoTsit5(Rosenbrock23()), + NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 8))), ] @info "Solver: $solver" sol = @time solve(flow_bvp, solver; abstol = 1e-8, reltol = 1e-8, @@ -231,148 +227,139 @@ end end end -# @testset "Ray Tracing" begin -# @inline v(x, y, z, p) = 1 / (4 + cos(p[1] * x) + sin(p[2] * y) - cos(p[3] * z)) -# @inline ux(x, y, z, p) = -p[1] * sin(p[1] * x) -# @inline uy(x, y, z, p) = p[2] * cos(p[2] * y) -# @inline uz(x, y, z, p) = p[3] * sin(p[3] * z) - -# function ray_tracing(u, p, t) -# du = similar(u) -# ray_tracing!(du, u, p, t) -# return du -# end - -# function ray_tracing!(du, u, p, t) -# x, y, z, ξ, η, ζ, T, S = u - -# nu = v(x, y, z, p) # Velocity of a sound wave, function of space; -# μx = ux(x, y, z, p) # ∂(slowness)/∂x, function of space -# μy = uy(x, y, z, p) # ∂(slowness)/∂y, function of space -# μz = uz(x, y, z, p) # ∂(slowness)/∂z, function of space - -# du[1] = S * nu * ξ -# du[2] = S * nu * η -# du[3] = S * nu * ζ - -# du[4] = S * μx -# du[5] = S * μy -# du[6] = S * μz - -# du[7] = S / nu -# du[8] = 0 - -# return nothing -# end - -# function ray_tracing_bc(sol, p, t) -# res = similar(first(sol)) -# ray_tracing_bc!(res, sol, p, t) -# return res -# end - -# function ray_tracing_bc!(res, sol, p, t) -# ua = sol(0.0) -# ub = sol(1.0) -# nu = v(ua[1], ua[2], ua[3], p) # Velocity of a sound wave, function of space; - -# res[1] = ua[1] - p[4] -# res[2] = ua[2] - p[5] -# res[3] = ua[3] - p[6] -# res[4] = ua[7] # T(0) = 0 -# res[5] = ua[4]^2 + ua[5]^2 + ua[6]^2 - 1 / nu^2 -# res[6] = ub[1] - p[7] -# res[7] = ub[2] - p[8] -# res[8] = ub[3] - p[9] -# return nothing -# end - -# function ray_tracing_bc_a(ua, p) -# resa = similar(ua, 5) -# ray_tracing_bc_a!(resa, ua, p) -# return resa -# end - -# function ray_tracing_bc_a!(resa, ua, p) -# nu = v(ua[1], ua[2], ua[3], p) # Velocity of a sound wave, function of space; - -# resa[1] = ua[1] - p[4] -# resa[2] = ua[2] - p[5] -# resa[3] = ua[3] - p[5] -# resa[4] = ua[7] -# resa[5] = ua[4]^2 + ua[5]^2 + ua[6]^2 - 1 / nu^2 - -# return nothing -# end - -# function ray_tracing_bc_b(ub, p) -# resb = similar(ub, 3) -# ray_tracing_bc_b!(resb, ub, p) -# return resb -# end - -# function ray_tracing_bc_b!(resb, ub, p) -# resb[1] = ub[1] - p[7] -# resb[2] = ub[2] - p[8] -# resb[3] = ub[3] - p[9] -# return nothing -# end - -# p = [0, 1, 2, 0, 0, 0, 4, 3, 2.0] - -# dx = p[7] - p[4] -# dy = p[8] - p[5] -# dz = p[9] - p[6] - -# u0 = zeros(8) -# u0[1:3] .= 0 # position -# u0[4] = dx / v(p[4], p[5], p[6], p) -# u0[5] = dy / v(p[4], p[5], p[6], p) -# u0[6] = dz / v(p[4], p[5], p[6], p) -# u0[8] = 1 - -# tspan = (0.0, 1.0) - -# prob_oop = BVProblem(BVPFunction{false}(ray_tracing, ray_tracing_bc), u0, tspan, p; -# nlls = Val(false)) -# prob_iip = BVProblem(BVPFunction{true}(ray_tracing!, ray_tracing_bc!), u0, tspan, p; -# nlls = Val(true)) -# prob_tp_oop = BVProblem(BVPFunction{false}(ray_tracing, -# (ray_tracing_bc_a, ray_tracing_bc_b); twopoint = Val(true)), u0, tspan, p; -# nlls = Val(true)) -# prob_tp_iip = BVProblem(BVPFunction{true}(ray_tracing!, -# (ray_tracing_bc_a!, ray_tracing_bc_b!); -# bcresid_prototype = (zeros(5), zeros(3)), twopoint = Val(true)), u0, tspan, p; -# nlls = Val(true)) - -# @info "Ray Tracing: Multiple Shooting" - -# alg_sp = MultipleShooting(10, AutoVern7(Rodas4P()); grid_coarsening = true, -# nlsolve = TrustRegion(), -# jac_alg = BVPJacobianAlgorithm(; bc_diffmode = AutoForwardDiff(; chunksize = 8), -# nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 8))) -# alg_dense = MultipleShooting(10, AutoVern7(Rodas4P()); grid_coarsening = true, -# nlsolve = TrustRegion(), -# jac_alg = BVPJacobianAlgorithm(; bc_diffmode = AutoForwardDiff(; chunksize = 8), -# nonbc_diffmode = AutoForwardDiff(; chunksize = 8))) -# alg_default = MultipleShooting(10, AutoVern7(Rodas4P()); grid_coarsening = true) - -# for (prob, alg) in Iterators.product((prob_oop, prob_iip, prob_tp_oop, prob_tp_iip), -# (alg_sp, alg_dense, alg_default)) -# @info "Solver: $alg" -# @time sol = solve(prob, alg; abstol = 1e-6, reltol = 1e-6, maxiters = 1000, -# odesolve_kwargs = (; abstol = 1e-8, reltol = 1e-5)) -# @test SciMLBase.successful_retcode(sol.retcode) - -# if prob.problem_type isa TwoPointBVProblem -# resida, residb = zeros(5), zeros(3) -# ray_tracing_bc_a!(resida, sol.u[1], p) -# ray_tracing_bc_b!(residb, sol.u[end], p) -# @test norm(vcat(resida, residb), Inf) < 1e-6 -# else -# resid = zeros(8) -# ray_tracing_bc!(resid, sol, p, sol.t) -# @test norm(resid, Inf) < 1e-6 -# end -# end -# end +@testset "Ray Tracing" begin + @inline v(x, y, z, p) = 1 / (4 + cos(p[1] * x) + sin(p[2] * y) - cos(p[3] * z)) + @inline ux(x, y, z, p) = -p[1] * sin(p[1] * x) + @inline uy(x, y, z, p) = p[2] * cos(p[2] * y) + @inline uz(x, y, z, p) = p[3] * sin(p[3] * z) + + function ray_tracing(u, p, t) + du = similar(u) + ray_tracing!(du, u, p, t) + return du + end + + function ray_tracing!(du, u, p, t) + x, y, z, ξ, η, ζ, T, S = u + + nu = v(x, y, z, p) # Velocity of a sound wave, function of space; + μx = ux(x, y, z, p) # ∂(slowness)/∂x, function of space + μy = uy(x, y, z, p) # ∂(slowness)/∂y, function of space + μz = uz(x, y, z, p) # ∂(slowness)/∂z, function of space + + du[1] = S * nu * ξ + du[2] = S * nu * η + du[3] = S * nu * ζ + + du[4] = S * μx + du[5] = S * μy + du[6] = S * μz + + du[7] = S / nu + du[8] = 0 + + return nothing + end + + function ray_tracing_bc(sol, p, t) + res = similar(first(sol)) + ray_tracing_bc!(res, sol, p, t) + return res + end + + function ray_tracing_bc!(res, sol, p, t) + ua = sol(0.0) + ub = sol(1.0) + nu = v(ua[1], ua[2], ua[3], p) # Velocity of a sound wave, function of space; + + res[1] = ua[1] - p[4] + res[2] = ua[2] - p[5] + res[3] = ua[3] - p[6] + res[4] = ua[7] # T(0) = 0 + res[5] = ua[4]^2 + ua[5]^2 + ua[6]^2 - 1 / nu^2 + res[6] = ub[1] - p[7] + res[7] = ub[2] - p[8] + res[8] = ub[3] - p[9] + return nothing + end + + function ray_tracing_bc_a(ua, p) + resa = similar(ua, 5) + ray_tracing_bc_a!(resa, ua, p) + return resa + end + + function ray_tracing_bc_a!(resa, ua, p) + nu = v(ua[1], ua[2], ua[3], p) # Velocity of a sound wave, function of space; + + resa[1] = ua[1] - p[4] + resa[2] = ua[2] - p[5] + resa[3] = ua[3] - p[5] + resa[4] = ua[7] + resa[5] = ua[4]^2 + ua[5]^2 + ua[6]^2 - 1 / nu^2 + + return nothing + end + + function ray_tracing_bc_b(ub, p) + resb = similar(ub, 3) + ray_tracing_bc_b!(resb, ub, p) + return resb + end + + function ray_tracing_bc_b!(resb, ub, p) + resb[1] = ub[1] - p[7] + resb[2] = ub[2] - p[8] + resb[3] = ub[3] - p[9] + return nothing + end + + p = [0, 1, 2, 0, 0, 0, 4, 3, 2.0] + + dx = p[7] - p[4] + dy = p[8] - p[5] + dz = p[9] - p[6] + + u0 = zeros(8) + u0[1:3] .= 0 # position + u0[4] = dx / v(p[4], p[5], p[6], p) + u0[5] = dy / v(p[4], p[5], p[6], p) + u0[6] = dz / v(p[4], p[5], p[6], p) + u0[8] = 1 + + tspan = (0.0, 1.0) + + prob_oop = BVProblem(BVPFunction{false}(ray_tracing, ray_tracing_bc), u0, tspan, p; + nlls = Val(false)) + prob_iip = BVProblem(BVPFunction{true}(ray_tracing!, ray_tracing_bc!), u0, tspan, p; + nlls = Val(true)) + prob_tp_oop = BVProblem(BVPFunction{false}(ray_tracing, + (ray_tracing_bc_a, ray_tracing_bc_b); twopoint = Val(true)), u0, tspan, p; + nlls = Val(true)) + prob_tp_iip = BVProblem(BVPFunction{true}(ray_tracing!, + (ray_tracing_bc_a!, ray_tracing_bc_b!); + bcresid_prototype = (zeros(5), zeros(3)), twopoint = Val(true)), u0, tspan, p; + nlls = Val(true)) + + @info "Ray Tracing: Multiple Shooting" + + alg_sp = MultipleShooting(10, AutoVern7(Rodas4P()); grid_coarsening = true, + nlsolve = TrustRegion(), + jac_alg = BVPJacobianAlgorithm(; bc_diffmode = AutoForwardDiff(; chunksize = 8), + nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 8))) + alg_dense = MultipleShooting(10, AutoVern7(Rodas4P()); grid_coarsening = true, + nlsolve = TrustRegion(), + jac_alg = BVPJacobianAlgorithm(; bc_diffmode = AutoForwardDiff(; chunksize = 8), + nonbc_diffmode = AutoForwardDiff(; chunksize = 8))) + alg_default = MultipleShooting(10, AutoVern7(Rodas4P()); grid_coarsening = true) + + for (prob, alg) in Iterators.product((prob_oop, prob_iip, prob_tp_oop, prob_tp_iip), + (alg_sp, alg_dense, alg_default)) + @info "Solver: $alg" + @time sol = solve(prob, alg; abstol = 1e-6, reltol = 1e-6, maxiters = 1000, + odesolve_kwargs = (; abstol = 1e-8, reltol = 1e-5)) + + @test SciMLBase.successful_retcode(sol.retcode) + @test norm(sol.resid, Inf) < 1e-6 + end +end diff --git a/test/shooting/nonlinear_least_squares.jl b/test/shooting/nonlinear_least_squares.jl index 304cc805..1b60d831 100644 --- a/test/shooting/nonlinear_least_squares.jl +++ b/test/shooting/nonlinear_least_squares.jl @@ -1,13 +1,22 @@ -using BoundaryValueDiffEq, LinearAlgebra, OrdinaryDiffEq, Test +using BoundaryValueDiffEq, LinearAlgebra, OrdinaryDiffEq, Test, JET @testset "Overconstrained BVP" begin SOLVERS = [ Shooting(Tsit5()), - Shooting(Tsit5(); nlsolve = GaussNewton()), - Shooting(Tsit5(); nlsolve = TrustRegion()), + Shooting(Tsit5(), GaussNewton(; autodiff = AutoForwardDiff(; chunksize = 2))), + Shooting(Tsit5(), GaussNewton(; autodiff = AutoFiniteDiff())), + Shooting(Tsit5(), TrustRegion(; autodiff = AutoForwardDiff(; chunksize = 2))), + Shooting(Tsit5(), TrustRegion(; autodiff = AutoFiniteDiff())), MultipleShooting(10, Tsit5()), - MultipleShooting(10, Tsit5(); nlsolve = GaussNewton()), - MultipleShooting(10, Tsit5(); nlsolve = TrustRegion())] + MultipleShooting(10, Tsit5(), + GaussNewton(; autodiff = AutoForwardDiff(; chunksize = 2))), + MultipleShooting(10, Tsit5(), GaussNewton(; autodiff = AutoFiniteDiff())), + MultipleShooting(10, Tsit5(), + TrustRegion(; autodiff = AutoForwardDiff(; chunksize = 2))), + MultipleShooting(10, Tsit5(), TrustRegion(; autodiff = AutoFiniteDiff())), + ] + JET_SKIP = [true, false, false, false, false, true, false, false, false, false] + JET_BROKEN = [false, false, false, false, false, false, false, false, false, false] @info "Solving Overconstrained BVPs" @@ -29,11 +38,19 @@ using BoundaryValueDiffEq, LinearAlgebra, OrdinaryDiffEq, Test bvp1 = BVProblem(BVPFunction{false}(f1, bc1; bcresid_prototype = zeros(4)), u0, tspan; nlls = Val(true)) - for solver in SOLVERS + for (i, solver) in enumerate(SOLVERS) @info "Testing $solver" sol = @time solve(bvp1, solver; verbose = false, abstol = 1e-6, reltol = 1e-6, odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) @test norm(sol.resid, Inf) < 0.005 + + JET_SKIP[i] && continue + @test_opt target_modules=(SciMLBase, DiffEqBase, NonlinearSolve, + BoundaryValueDiffEq) solve(bvp1, solver; verbose = false, abstol = 1e-6, + reltol = 1e-6, odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) broken=JET_BROKEN[i] + @test_call target_modules=(SciMLBase, DiffEqBase, NonlinearSolve, + BoundaryValueDiffEq) solve(bvp1, solver; verbose = false, abstol = 1e-6, + reltol = 1e-6, odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) broken=JET_BROKEN[i] end # IIP MP-BVP @@ -59,11 +76,19 @@ using BoundaryValueDiffEq, LinearAlgebra, OrdinaryDiffEq, Test bvp2 = BVProblem(BVPFunction{true}(f1!, bc1!; bcresid_prototype = zeros(4)), u0, tspan; nlls = Val(true)) - for solver in SOLVERS + for (i, solver) in enumerate(SOLVERS) @info "Testing $solver" sol = @time solve(bvp2, solver; verbose = false, abstol = 1e-6, reltol = 1e-6, odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) @test norm(sol.resid, Inf) < 0.005 + + JET_SKIP[i] && continue + @test_opt target_modules=(SciMLBase, DiffEqBase, NonlinearSolve, + BoundaryValueDiffEq) solve(bvp2, solver; verbose = false, abstol = 1e-6, + reltol = 1e-6, odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) broken=JET_BROKEN[i] + @test_call target_modules=(SciMLBase, DiffEqBase, NonlinearSolve, + BoundaryValueDiffEq) solve(bvp2, solver; verbose = false, abstol = 1e-6, + reltol = 1e-6, odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) broken=JET_BROKEN[i] end # OOP TP-BVP @@ -73,10 +98,19 @@ using BoundaryValueDiffEq, LinearAlgebra, OrdinaryDiffEq, Test bvp3 = BVProblem(BVPFunction{false}(f1, (bc1a, bc1b); twopoint = Val(true), bcresid_prototype = (zeros(1), zeros(2))), u0, tspan; nlls = Val(true)) - for solver in SOLVERS + for (i, solver) in enumerate(SOLVERS) @info "Testing $solver" - sol = @time solve(bvp3, solver; verbose = false) - @test norm(sol.resid, Inf) < 1e-4 + sol = @time solve(bvp3, solver; verbose = false, abstol = 1e-6, reltol = 1e-6, + odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) + @test norm(sol.resid, Inf) < 0.009 + + JET_SKIP[i] && continue + @test_opt target_modules=(SciMLBase, DiffEqBase, NonlinearSolve, + BoundaryValueDiffEq) solve(bvp3, solver; verbose = false, abstol = 1e-6, + reltol = 1e-6, odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) broken=JET_BROKEN[i] + @test_call target_modules=(SciMLBase, DiffEqBase, NonlinearSolve, + BoundaryValueDiffEq) solve(bvp3, solver; verbose = false, abstol = 1e-6, + reltol = 1e-6, odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) broken=JET_BROKEN[i] end # IIP TP-BVP @@ -86,9 +120,18 @@ using BoundaryValueDiffEq, LinearAlgebra, OrdinaryDiffEq, Test bvp4 = BVProblem(BVPFunction{true}(f1!, (bc1a!, bc1b!); twopoint = Val(true), bcresid_prototype = (zeros(1), zeros(2))), u0, tspan; nlls = Val(true)) - for solver in SOLVERS + for (i, solver) in enumerate(SOLVERS) @info "Testing $solver" - sol = @time solve(bvp4, solver; verbose = false) - @test norm(sol.resid, Inf) < 1e-4 + sol = @time solve(bvp4, solver; verbose = false, abstol = 1e-6, reltol = 1e-6, + odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) + @test norm(sol.resid, Inf) < 0.009 + + JET_SKIP[i] && continue + @test_opt target_modules=(SciMLBase, DiffEqBase, NonlinearSolve, + BoundaryValueDiffEq) solve(bvp4, solver; verbose = false, abstol = 1e-6, + reltol = 1e-6, odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) broken=JET_BROKEN[i] + @test_call target_modules=(SciMLBase, DiffEqBase, NonlinearSolve, + BoundaryValueDiffEq) solve(bvp4, solver; verbose = false, abstol = 1e-6, + reltol = 1e-6, odesolve_kwargs = (; abstol = 1e-6, reltol = 1e-6)) broken=JET_BROKEN[i] end end diff --git a/test/shooting/orbital.jl b/test/shooting/orbital.jl index 8e2d5067..f1f8899c 100644 --- a/test/shooting/orbital.jl +++ b/test/shooting/orbital.jl @@ -62,8 +62,6 @@ using BoundaryValueDiffEq, OrdinaryDiffEq, LinearAlgebra, Test cur_bc! = (resid, sol, p, t) -> bc!_generator(resid, sol, init_val) cur_bc_2point_a! = (resid, sol, p) -> bc!_generator_2p_a(resid, sol, init_val) 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)) @info "Solving Lambert's Problem in Multi Point BVP Form" @@ -72,14 +70,15 @@ using BoundaryValueDiffEq, OrdinaryDiffEq, LinearAlgebra, Test AutoFiniteDiff(; fdtype = Val(:central)), AutoSparseForwardDiff(; chunksize = 6), AutoFiniteDiff(; fdtype = Val(:forward)), AutoSparseFiniteDiff()) nlsolve = TrustRegion(; autodiff) + @info "Single Shooting Lambert's Problem: $(autodiff)" @time sol = solve(bvp, Shooting(DP5(); nlsolve); force_dtmin = true, abstol = 1e-6, reltol = 1e-6, verbose = false, odesolve_kwargs = (abstol = 1e-6, reltol = 1e-3)) - cur_bc!(resid_f, sol, nothing, sol.t) - @info "Single Shooting Lambert's Problem: $(norm(resid_f, Inf))" + + @info "Single Shooting Lambert's Problem: $(norm(sol.resid, Inf))" @test SciMLBase.successful_retcode(sol) - @test norm(resid_f, Inf) < 1e-6 + @test norm(sol.resid, Inf) < 1e-6 @info "Multiple Shooting Lambert's Problem: $(autodiff)" jac_alg = BVPJacobianAlgorithm(; nonbc_diffmode = autodiff, @@ -87,13 +86,12 @@ using BoundaryValueDiffEq, OrdinaryDiffEq, LinearAlgebra, Test @time sol = solve(bvp, MultipleShooting(10, DP5(); nlsolve, jac_alg); force_dtmin = true, abstol = 1e-6, reltol = 1e-6, verbose = false, odesolve_kwargs = (abstol = 1e-6, reltol = 1e-3)) - cur_bc!(resid_f, sol, nothing, sol.t) - @info "Multiple Shooting Lambert's Problem: $(norm(resid_f, Inf))" + + @info "Multiple Shooting Lambert's Problem: $(norm(sol.resid, Inf))" @test SciMLBase.successful_retcode(sol) - @test norm(resid_f, Inf) < 1e-6 + @test norm(sol.resid, Inf) < 1e-6 end - println() @info "Solving Lambert's Problem in Two Point BVP Form" bvp = TwoPointBVProblem(orbital!, (cur_bc_2point_a!, cur_bc_2point_b!), y0, tspan; @@ -103,25 +101,24 @@ using BoundaryValueDiffEq, OrdinaryDiffEq, LinearAlgebra, Test AutoFiniteDiff(; fdtype = Val(:central)), AutoFiniteDiff(; fdtype = Val(:forward)), AutoSparseForwardDiff(; chunksize = 6)) nlsolve = TrustRegion(; autodiff) + @info "Single Shooting Lambert's Problem: $(autodiff)" @time sol = solve(bvp, Shooting(DP5(); nlsolve); force_dtmin = true, abstol = 1e-6, reltol = 1e-6, verbose = false, odesolve_kwargs = (abstol = 1e-6, reltol = 1e-3)) - cur_bc_2point_a!(resid_f_2p[1], sol(t0), nothing) - cur_bc_2point_b!(resid_f_2p[2], sol(t1), nothing) - @info "Single Shooting Lambert's Problem: $(norm(reduce(vcat, resid_f_2p), Inf))" + + @info "Single Shooting Lambert's Problem: $(norm(sol.resid, Inf))" @test SciMLBase.successful_retcode(sol) - @test norm(reduce(vcat, resid_f_2p), Inf) < 1e-6 + @test norm(sol.resid, Inf) < 1e-6 @info "Multiple Shooting Lambert's Problem: $(autodiff)" jac_alg = BVPJacobianAlgorithm(; nonbc_diffmode = autodiff, bc_diffmode = autodiff) @time sol = solve(bvp, MultipleShooting(10, DP5(); nlsolve, jac_alg); force_dtmin = true, abstol = 1e-6, reltol = 1e-6, verbose = false, odesolve_kwargs = (abstol = 1e-6, reltol = 1e-3)) - cur_bc_2point_a!(resid_f_2p[1], sol(t0), nothing) - cur_bc_2point_b!(resid_f_2p[2], sol(t1), nothing) - @info "Multiple Shooting Lambert's Problem: $(norm(reduce(vcat, resid_f_2p), Inf))" + + @info "Multiple Shooting Lambert's Problem: $(norm(sol.resid, Inf))" @test SciMLBase.successful_retcode(sol) - @test norm(reduce(vcat, resid_f_2p), Inf) < 1e-6 + @test norm(sol.resid, Inf) < 1e-6 end end