Got MBH working! It can find better-than-basic-spiral trajectories

julia/src/find_closest.jl

@@ -1,6 +1,15 @@
-using JuMP, Ipopt
+using NLsolve
 
-export nlp_solve
+export nlp_solve, mass_est
+
+function mass_est(T)
+  ans = 0
+  n = Int(length(T)/3)
+  for i in 1:n
+    ans += norm(T[i,:])
+  end
+  return ans/n
+end
 
 function nlp_solve(start::Vector{Float64},
                    final::Vector{Float64},
@@ -8,77 +17,14 @@ function nlp_solve(start::Vector{Float64},
                    μ::Float64,
                    t0::Float64,
                    tf::Float64,
-                   Txi::Vector{Float64},
-                   Tyi::Vector{Float64},
-                   Tzi::Vector{Float64};
-                   solver_options=(),
+                   x0::Matrix{Float64};
                    tol=1e-6)
 
-  n::Int = length(Txi)
-  if length(Tyi) != n
-    throw("Bad number of Tys")
-  elseif length(Tzi) != n
-    throw("Bad number of Tzs")
+  function f!(F,x)
+    F .= 0.0
+    F[1:6, 1] .= prop_nlsolve(tanh.(x), start, craft, μ, tf-t0) .- final
   end
 
-  # First propagate the initial guess 
-  path, masses, final_state = prop(hcat(Txi, Tyi, Tzi),
-                                   start,
-                                   craft,
-                                   μ,
-                                   tf-t0)
-  @assert final ≈ final_state
-  model = Model(optimizer_with_attributes(Ipopt.Optimizer, solver_options...))
-
-  @variables(model, begin
-    Tx[i=1:n], (start=Txi[i])
-    Ty[i=1:n], (start=Tyi[i])
-    Tz[i=1:n], (start=Tzi[i])
-    x[i=1:n], (start=path[1][i])
-    y[i=1:n], (start=path[2][i])
-    z[i=1:n], (start=path[3][i])
-    dx[i=1:n], (start=path[4][i])
-    dy[i=1:n], (start=path[5][i])
-    dz[i=1:n], (start=path[6][i])
-    mass[i=1:n], (start=masses[i])
-  end)
-
-  # Fix initial conditions
-  fix(x[1], start[1], force = true)
-  fix(y[1], start[2], force = true)
-  fix(z[1], start[3], force = true)
-  fix(dx[1], start[4], force = true)
-  fix(dy[1], start[5], force = true)
-  fix(dz[1], start[6], force = true)
-  fix(mass[1], craft.mass, force = true)
-
-  # Fix final conditions
-  fix(x[n], final[1], force = true)
-  fix(y[n], final[2], force = true)
-  fix(z[n], final[3], force = true)
-  fix(dx[n], final[4], force = true)
-  fix(dy[n], final[5], force = true)
-  fix(dz[n], final[6], force = true)
-
-  @NLexpression(model, r_mag[i = 1:n], sqrt(x[i]^2 + y[i]^2 + z[i]^2))
-  @NLexpression(model, v_mag[i = 1:n], sqrt(dx[i]^2 + dy[i]^2 + dz[i]^2))
-  @NLexpression(model, σ0[i = 1:n], (x[i]*dx[i] + y[i]*dy[i] + z[i]*dz[i])/sqrt(μ))
-  @NLexpression(model, a[i = 1:n], 1/(2/r_mag[i] - v_mag[i]^2/μ))
-
-  # Elliptical Specific
-  @NLexpression(model, ΔM_ell[i = 1:n], sqrt(μ) / sqrt(a[i]^3) * (tf-t0)/(2n))
-
-  # Parabolic/Hyperbolic Specific
-  @NLexpression(model, ΔN_hyp[i = 1:n], sqrt(μ) / sqrt(-a[i]^3) * (tf-t0)/(2n))
-
-  for i in 2:n
-    @NLconstraint(model, x[i] == a[i-1])
-    @NLconstraint(model, mass[i] == mass[i-1] - craft.mass_flow_rate*√(Tx[i-1]^2 +
-                                                                       Ty[i-1]^2 +
-                                                                       Tz[i-1]^2)*(tf-t0)/n)
-  end
-
-  optimize!(model)
-  return model
+  return nlsolve(f!, atanh.(x0), ftol=tol, autodiff=:forward, iterations=1_000)
 
 end