Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Porting over Multiple Shooting and Single Shooting from NeuralBVP #110

Merged
merged 29 commits into from
Oct 12, 2023
Merged

Porting over Multiple Shooting and Single Shooting from NeuralBVP #110

merged 29 commits into from
Oct 12, 2023

Conversation

avik-pal
Copy link
Member

@avik-pal avik-pal commented Sep 13, 2023 •
edited
Loading

Don't merge before #109

  • Single shooting now takes odesolve_kwargs and nlsolve_kwargs to allow different kwargs. repetition in kwargs overrides these choices
  • Single shooting can now handle non-vector inputs.
  • Working Multiple Shooting Code with Grid Coarsening

TODOs

  • Port tests from NeuralBVP.jl
  • Translate NeuralBVP.jl to use BoundaryValueDiffEq.jl (Upstream issue)
  • Incorporate using Initial Guess in Shooting Codes
  • Use Jacobian sparsity to speed up Newton Steps!
    • Sparse Jacobian Prototype Constructor
    • Hardcoded version using ForwardDiff
      • Threading causes Segfaults
    • Use Coloring
    • Generalize
  • Two Point BVP Fast Path
  • Test BVP with NLS (Will do it in a separate PR)

@avik-pal avik-pal force-pushed the ap/multiple-shooting branch 6 times, most recently from 6f5da29 to a1ba675 Compare September 28, 2023 18:08
@avik-pal avik-pal force-pushed the ap/multiple-shooting branch from 8d69044 to ee7e240 Compare October 1, 2023 02:08
@avik-pal avik-pal force-pushed the ap/multiple-shooting branch 2 times, most recently from c44d95b to 060116a Compare October 6, 2023 18:21
@avik-pal avik-pal force-pushed the ap/multiple-shooting branch from 060116a to a4694e8 Compare October 9, 2023 14:37
@avik-pal avik-pal force-pushed the ap/multiple-shooting branch from b432e16 to df3909c Compare October 9, 2023 19:42
@avik-pal avik-pal force-pushed the ap/multiple-shooting branch from bdcef85 to 8eedf1f Compare October 10, 2023 03:07
@avik-pal avik-pal changed the title [WIP] Porting over Multiple Shooting and Single Shooting from NeuralBVP Porting over Multiple Shooting and Single Shooting from NeuralBVP Oct 10, 2023
@avik-pal avik-pal force-pushed the ap/multiple-shooting branch from cfbeba1 to a5b10cc Compare October 10, 2023 21:07
@avik-pal avik-pal force-pushed the ap/multiple-shooting branch from 20b966e to 6f40df9 Compare October 11, 2023 03:25
avik-pal
avik-pal commented Oct 11, 2023
View reviewed changes
src/solve/multiple_shooting.jl Outdated
Comment on lines 62 to 114
compute_bc_residual! = if prob.problem_type isa TwoPointBVProblem
@views function compute_bc_residual_tp!(resid_bc, us::ArrayPartition, p,
cur_nshoots, nodes, resid_nodes::Union{Nothing, MaybeDiffCache} = nothing)
ua, ub0 = us.x
# 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,
kwargs..., save_everystep = false, saveat = (), save_end = true)
ub = vec(sol_ode_last.u[end])

resid_bc_a, resid_bc_b = if resid_bc isa ArrayPartition
resid_bc.x
else
resid_bc[1:resida_len], resid_bc[(resida_len + 1):end]
end

if iip
bc[1](resid_bc_a, ua, p)
bc[2](resid_bc_b, ub, p)
else
resid_bc_a .= bc[1](ua, p)
resid_bc_b .= bc[2](ub, p)
end

return resid_bc
end
else
@views function compute_bc_residual_mp!(resid_bc, us, p, cur_nshoots, nodes,
resid_nodes::Union{Nothing, MaybeDiffCache} = nothing)
if resid_nodes === nothing
_resid_nodes = similar(us, cur_nshoots * N) # This might be Dual based on `us`
else
_resid_nodes = get_tmp(resid_nodes, us)
end

# NOTE: We need to recompute this to correctly propagate the dual numbers / gradients
_us, _ts = solve_internal_odes!(_resid_nodes, us, p, cur_nshoots, nodes)

# Boundary conditions
# Builds an ODESolution object to keep the framework for bc(,,) consistent
odeprob = ODEProblem{iip}(f, reshape(us[1:N], u0_size), tspan, p)
total_solution = SciMLBase.build_solution(odeprob, alg.ode_alg, _ts, _us)

if iip
eval_bc_residual!(resid_bc, prob.problem_type, bc, total_solution, p)
else
resid_bc .= eval_bc_residual(prob.problem_type, bc, total_solution, p)
end

return resid_bc
end
end
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@ChrisRackauckas can you take a look here? For the correct BC jacobian part we do need to forward solve again?

@avik-pal avik-pal force-pushed the ap/multiple-shooting branch from 4c135e2 to d4326d6 Compare October 11, 2023 16:57
@avik-pal avik-pal force-pushed the ap/multiple-shooting branch from d4326d6 to b9074b6 Compare October 11, 2023 18:10
@avik-pal
Copy link
Member Author

This PR got more bloated than I anticipated. But it should be done now. Takeaways:

  • For MS using the two-point form is significantly faster
  • For the Orbital problem, we need to use TrustRegion for proper convergence
  • Solving an Ensemble Problem for the MS gives the most gains.
  • There was a bug in the interpolation code, which has been patched

TODOs for later:

  • MS can't handle Non-Vector like Inputs. I will deal with it together with the handling for matrix inputs to handle Neural Networks better!

@ChrisRackauckas ChrisRackauckas merged commit 9611bb7 into SciML:master Oct 12, 2023
3 of 5 checks passed
@ChrisRackauckas
Copy link
Member

🎉 It was a bit hard to read since there were multiple things mixed in there, but generally fine so merging and will deal with anything else in the future.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ErikQQY ErikQQY ErikQQY left review comments

@ChrisRackauckas ChrisRackauckas ChrisRackauckas approved these changes

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@avik-pal @ChrisRackauckas @ErikQQY