Fix some dimension issues

Author: avik-pal

src/algorithms.jl

@@ -39,6 +39,10 @@ function concretize_jacobian_algorithm(alg::Shooting, prob)
     return Shooting(alg.ode_alg, alg.nlsolve, BVPJacobianAlgorithm(diffmode))
 end
 
+function Shooting(ode_alg, nlsolve; jac_alg = BVPJacobianAlgorithm())
+    return Shooting(ode_alg, nlsolve, jac_alg)
+end
+
 function Shooting(ode_alg; nlsolve = NewtonRaphson(), jac_alg = BVPJacobianAlgorithm())
     return Shooting(ode_alg, nlsolve, jac_alg)
 end

src/solve/multiple_shooting.jl

@@ -204,8 +204,7 @@ function __multiple_shooting_init_jacobian_odecache(ensemblealg, prob, jac_cache
 end
 
 function __multiple_shooting_init_jacobian_odecache(ensemblealg, prob, jac_cache,
-        ::Union{AutoForwardDiff, AutoSparseForwardDiff}, alg, nshoots, u;
-        kwargs...)
+        ::Union{AutoForwardDiff, AutoSparseForwardDiff}, alg, nshoots, u; kwargs...)
     cache = jac_cache.cache
     if cache isa ForwardDiff.JacobianConfig
         xduals = reshape(cache.duals[2][1:length(u)], size(u))
@@ -390,9 +389,9 @@ end
     sol = solve!(odecache)
 
     if SciMLBase.successful_retcode(sol)
-        u_at_nodes[1:N] .= vec(sol.u[1])
-        for i in 2:(nshoots + 1)
-            u_at_nodes[(N + (i - 2) * N) .+ (1:N)] .= vec(sol.u[i])
+        res = sol(nodes).u
+        for i in 1:length(nodes)
+            u_at_nodes[(i - 1) * N .+ (1:N)] .= vec(res[i])
         end
     else
         @warn "Initialization using odesolve failed. Initializing using 0s. It is \

src/solve/single_shooting.jl

@@ -38,7 +38,7 @@ function __solve(prob::BVProblem, alg_::Shooting; odesolve_kwargs = (;),
     end
 
     ode_cache_jac_fn = __single_shooting_jacobian_ode_cache(internal_prob, jac_cache,
-        alg.jac_alg.diffmode, vec(u0), alg.ode_alg; ode_kwargs...)
+        alg.jac_alg.diffmode, u0, alg.ode_alg; ode_kwargs...)
 
     jac_prototype = init_jacobian(jac_cache)
 
@@ -131,6 +131,6 @@ function __single_shooting_jacobian_ode_cache(prob, jac_cache,
         xduals = cache.t
     end
     fill!(xduals, 0)
-    prob_ = remake(prob; u0 = xduals)
+    prob_ = remake(prob; u0 = reshape(xduals, size(u0)))
     return SciMLBase.__init(prob_, ode_alg; kwargs...)
 end