[General] Code Tmax sun occult mission

The problem is formulated in terms of "what is the minimum amount of thrusters needed for relations these two points fixed in time and state (position & velocity), mass is free"
@jbcaillau

using SPICE
using OptimalControl
using NLPModelsIpopt
using JLD2
using OrdinaryDiffEq
using DataInterpolations
using CSV
using JLD2

using Plots
using Plots.PlotMeasures # for leftmargin, bottommargin
using LaTeXStrings

# SPICE Downloads
using Downloads: download


# SPICE files
const LSK = "https://naif.jpl.nasa.gov/pub/naif/generic_kernels/lsk/naif0012.tls"
const SPK = "https://naif.jpl.nasa.gov/pub/naif/generic_kernels/spk/planets/de440.bsp"
const PCK = "https://naif.jpl.nasa.gov/pub/naif/generic_kernels/pck/pck00010.tpc"

# Download kernels
download(LSK, "naif0012.tls")
download(SPK, "de440.bsp")
download(PCK, "pck00010.tpc")

furnsh("naif0012.tls") # Leap seconds kernel
furnsh("de430.bsp")
furnsh("pck00010.tpc")  # PCK kernel for planetary constants

# Astrodynamical parameters
G = 6.6743e-11 / 1000^3 # km^3 kg^-1 s^-2
M_moon = 7.346e22 # kg
M_earth = 5.972168e24 # kg
M_sun = 1.989e30 # kg
AU = 1.495978707e11 / 1e3 # km
LU = 384399 # km
TU = sqrt(LU^3 / G / (M_moon + M_earth))
MUnit =  M_moon + M_earth
mu_Earth = 3.9860044188e14 / 1000^3 / LU^3 * TU^2
mu_Moon = 4.90486959e12 / 1000^3 / LU^3 * TU^2
mu_Sun = 1.327124400189e20  / 1000^3 / LU^3 * TU^2
mass_moon = M_moon / MUnit
mass_earth = M_earth / MUnit
mass_sun = M_sun / MUnit

# Spacecraft parameters
g0 = 9.80665 / 1000 / LU * TU^2
Isp = 1660 / TU 
wet_mass = 400 / MUnit 
fuel_mass = 90 / MUnit 
dry_mass = wet_mass - fuel_mass
Tmax =  0.090 / 1000 / MUnit / LU * TU^2 #thrust force in N
Q = Tmax / Isp / g0

mass_scaling_factor = dry_mass + fuel_mass

# COASTING parameter
deltat = 2 * 24 * 3600 / TU

# Plot functions 
function plot_trajNV_multiarc_multidyn_withcoasting_1arc(sol, nb_arcs, tvec, moon_arcs, uncontrolled_arcs)

    sol_time = sol.time_grid.value 

    Nstate = length(sol.state.value(sol.time_grid.value[1]))


    Nsol = length(sol_time)
    sol_state = zeros(Nstate * nb_arcs, Nsol)
    sol_control = zeros(3 * (nb_arcs - length(uncontrolled_arcs)), Nsol)
    # sol_control = zeros(3 * nb_arcs, Nsol)
    # k = 1
    for i in range(1, Nsol)
        sol_state[:, i] = sol.state.value(sol.time_grid.value[i])
        sol_control[:, i] = sol.control.value(sol.time_grid.value[i])
        # if i ∉ uncontrolled_arcs
            # k += 1
        # end
    end

    for m in moon_arcs
        for i in range(1, Nsol)
            Nbegin = 1 + Nstate * (m - 1)
            Nend = 7 + Nstate * (m - 1)
            sol_state[Nbegin : Nend, i] = moon2emb(phi(sol_time[i], tvec[m], tvec[m+1]), sol_state[Nbegin : Nend, i])
        end
    end

    x_vals_opti = zeros(nb_arcs, Nsol)
    y_vals_opti = zeros(nb_arcs, Nsol)
    z_vals_opti = zeros(nb_arcs, Nsol)
    vx_vals_opti = zeros(nb_arcs, Nsol)
    vy_vals_opti = zeros(nb_arcs, Nsol)
    vz_vals_opti = zeros(nb_arcs, Nsol)
    m_vals_opti = zeros(nb_arcs, Nsol)

    ux_vals_opti = zeros(nb_arcs, Nsol)
    uy_vals_opti = zeros(nb_arcs, Nsol)
    uz_vals_opti = zeros(nb_arcs, Nsol)
    u_norm = zeros(nb_arcs, Nsol)

    k = 1
    for i in range(1, nb_arcs)
        x_vals_opti[i, :] = sol_state[1 + Nstate * (i - 1), :]
        y_vals_opti[i, :] = sol_state[2 + Nstate * (i - 1), :]
        z_vals_opti[i, :] = sol_state[3 + Nstate * (i - 1), :]
        vx_vals_opti[i, :] = sol_state[4 + Nstate * (i - 1), :]
        vy_vals_opti[i, :] = sol_state[5 + Nstate * (i - 1), :]
        vz_vals_opti[i, :] = sol_state[6 + Nstate * (i - 1), :]
        m_vals_opti[i, :] = sol_state[7 + Nstate * (i - 1), :]
        if i ∉ uncontrolled_arcs
            ux_vals_opti[i, :] = sol_control[1 + 3 * (k - 1), :]
            uy_vals_opti[i, :] = sol_control[2 + 3 * (k - 1), :]
            uz_vals_opti[i, :] = sol_control[3 + 3 * (k - 1), :]
            u_norm[i, :] = (ux_vals_opti[i, :].^2 + uy_vals_opti[i, :].^2 + uz_vals_opti[i, :].^2).^(1/2)
            k += 1
        end
    end

    sol_time_plot = LinRange(tvec[1], tvec[end], N)
    x_moon_pos = [x_moon(t) for t in sol_time_plot]
    y_moon_pos = [y_moon(t) for t in sol_time_plot]
    z_moon_pos = [z_moon(t) for t in sol_time_plot]
    x_earth_pos = [x_earth(t) for t in sol_time_plot]
    y_earth_pos = [y_earth(t) for t in sol_time_plot]
    z_earth_pos = [z_earth(t) for t in sol_time_plot]

    p = plot(x_moon_pos, y_moon_pos, z_moon_pos, label = "moon orbit")
    plot!(x_earth_pos, y_earth_pos, z_earth_pos, label = "earth orbit")

    scatter!([x_vals_opti[1, 1]], [y_vals_opti[1, 1]], [z_vals_opti[1, 1]], label = "previous occultation", color = :blue)

    labels = ["next occultation"]

    for i in range(1, nb_arcs)
        rand_color = RGB(rand(), rand(), rand()) 
        plot!(x_vals_opti[i, :], y_vals_opti[i, :], z_vals_opti[i, :], size=(600, 600), camera = (10,30), xlabel="x", ylabel="y", label = false, zlabel="z", linewidth = 2, color = "black", legend =:topleft)
        scatter!([x_vals_opti[i, end]], [y_vals_opti[i, end]], [z_vals_opti[i, end]], label = labels[i], color = :red)
        arrow_factor = 1e-1
        inds = findall(u_norm[i, :] .> 0.001)[1:2:end]
        quiver!(x_vals_opti[i, inds], y_vals_opti[i, inds], z_vals_opti[i, inds], quiver=(ux_vals_opti[i, inds] * arrow_factor, uy_vals_opti[i, inds] * arrow_factor, uz_vals_opti[i, inds] * arrow_factor), label="thrust directions", arrow=true, color = "red")
            
        # if i ∉ uncontrolled_arcs
            # arrow_factor = 1e-1
            # inds = findall(u_norm[k, :] .> 0.01)[1:2:end]
            # quiver!(x_vals_opti[i, inds], y_vals_opti[i, inds], z_vals_opti[i, inds], quiver=(ux_vals_opti[k, inds] * arrow_factor, uy_vals_opti[k, inds] * arrow_factor, uz_vals_opti[k, inds] * arrow_factor), label="thrust directions", arrow=true, color = "red")
            # k += 1
        # end
    end


    xlims!(-2, 2)
    ylims!(-2, 2)
    zlims!(-2, 2)
    
    display(p)
    

end


function plot_trajNV_multiarc_multidyn_synodic_1arc(sol, nb_arcs, tvec, moon_arcs)

    sol_time = sol.time_grid.value 

    Nstate = length(sol.state.value(sol.time_grid.value[1]))

    Nsol = length(sol_time)
    sol_state = zeros(Nstate * nb_arcs, Nsol)
    for i in range(1, Nsol)
        sol_state[:, i] = sol.state.value(sol.time_grid.value[i])
    end

    for m in moon_arcs
        for i in range(1, Nsol)
            Nbegin = 1 + Nstate * (m - 1)
            Nend = Nstate + Nstate * (m - 1)
            sol_state[Nbegin : Nend, i] = moon2emb(sol_time[i], sol_state[Nbegin : Nend, i])
        end
    end

    x_vals_opti = zeros(Nsol)
    y_vals_opti = zeros(Nsol)
    z_vals_opti = zeros(Nsol)
    vx_vals_opti = zeros(Nsol)
    vy_vals_opti = zeros(Nsol)
    vz_vals_opti = zeros(Nsol)
    m_vals_opti = zeros(Nsol)

    for i in range(1, nb_arcs)
        x_vals_opti = sol_state[1 , :]
        y_vals_opti = sol_state[2 , :]
        z_vals_opti = sol_state[3 , :]
        vx_vals_opti = sol_state[4 , :]
        vy_vals_opti = sol_state[5 , :]
        vz_vals_opti = sol_state[6 , :]
        m_vals_opti = sol_state[7 , :]
    end

    x_syn = zeros( Nsol)
    y_syn = zeros(Nsol)
    z_syn = zeros( Nsol)

    for ii = 1:Nsol
        
        x_sol_syn = ECLIPJ2000toSynodic(sol_time[ii] * TU,  
            sol_state[1 : 3, ii])
        x_syn[ii] = x_sol_syn[1]
        y_syn[ii] = x_sol_syn[2]
        z_syn[ii] = x_sol_syn[3]

    end

    p = plot(x_syn, y_syn, z_syn, legend = false, linewidth = 2, color = :black)
    scatter!([x_syn[end]], [y_syn[end]], [z_syn[end]], color = :red, xlabel = "x", ylabel = "y", zlabel = "z")
    scatter!([x_syn[1]], [y_syn[1]], [z_syn[1]], color = :red, xlabel = "x", ylabel = "y", zlabel = "z")
    
    scatter!([norm([x_moon(t0), y_moon(t0), z_moon(t0)])], [0], [0], color = :grey, label = "moon", markersize=7)
    scatter!([norm([x_earth(t0), y_earth(t0), z_earth(t0)])], [0], [0], color = :blue, label = "earth", markersize=10)

    xlims!(-2, 2)
    ylims!(-2, 2)
    zlims!(-2, 2)
    
    display(p)
    

end


function ECLIPJ2000toSynodic(et, x)

    xm    = spkezr("MOON", et, "ECLIPJ2000", "NONE", "EARTH")
        
    r     = xm[1][1:3]
    v     = xm[1][4:6]

    normr = norm(r)

    h     = cross(r, v)
    normh = norm(h)

    er    = r/normr
    eh    = h/normh
    et    = cross(eh, er)/norm(cross(eh, er))

    IS    = [er[1] er[2] er[3];
             et[1] et[2] et[3];
             eh[1] eh[2] eh[3]]

    xs    = IS * x[1:3]

    return xs
end




# Dynamics functions

function propagate!(du, u, p, t)
    r = u[1:3]
    v = u[4:6]
    m = u[7] * mass_scaling_factor

    x_e = [x_earth(t), y_earth(t), z_earth(t)]
    x_m = [x_moon(t), y_moon(t), z_moon(t)]
    x_s = [x_sun(t), y_sun(t), z_sun(t)]

    normrE = sqrt(sum((x_e - r).^2))
    normrM = sqrt(sum((x_m - r).^2))
    normrS = sqrt(sum((x_s - r).^2))

    normrSE = sqrt(sum((x_e - x_s).^2))
    normrMS = sqrt(sum((x_m - x_s).^2))
    
    dv = mu_Earth / normrE^3 * (x_e - r) + mu_Moon / normrM^3 * (x_m - r) + mu_Sun / normrS^3 * (x_s - r) - (1 / (mass_earth + mass_moon)) * ( x_s * (mu_Sun * mass_earth / normrSE^3 + mu_Moon * mass_sun / normrMS^3) - x_m * (mu_Moon * mass_sun / normrMS^3) - x_e * (mu_Sun * mass_earth / normrSE^3)) 
    dvu = Tmax * p / m

    dm = - Q * sqrt(sum(p.^2)) / mass_scaling_factor

    du[1:3] = v
    du[4:6] = dv + dvu
    du[7]   = dm
end

function F0(t, x, u)
    # Earth & Moon & Sun gravity
    
    r     = x[1:3]
    v     = x[4:6]
    m     = x[7] * mass_scaling_factor
    N_th  = x[8]

    x_e = [x_earth(t), y_earth(t), z_earth(t)]
    x_m = [x_moon(t), y_moon(t), z_moon(t)]
    x_s = [x_sun(t), y_sun(t), z_sun(t)]

    normrE = sqrt(sum((x_e - r).^2))
    normrM = sqrt(sum((x_m - r).^2))
    normrS = sqrt(sum((x_s - r).^2))

    normrSE = sqrt(sum((x_e - x_s).^2))
    normrMS = sqrt(sum((x_m - x_s).^2))

    dv = mu_Earth / normrE^3 * (x_e - r) + mu_Moon / normrM^3 * (x_m - r) + mu_Sun / normrS^3 * (x_s - r) - (1 / (mass_earth + mass_moon)) * ( x_s * (mu_Sun * mass_earth / normrSE^3 + mu_Moon * mass_sun / normrMS^3) - x_m * (mu_Moon * mass_sun / normrMS^3) - x_e * (mu_Sun * mass_earth / normrSE^3)) 
    dvu = N_th * Tmax * u / m
    dm = - N_th * Q * sqrt(sum(u.^2)) / mass_scaling_factor

    dx = [v; dv + dvu; dm; 0]
    return dx
end

function Fu(t, x, u)
    # Earth & Moon & Sun gravity
    
    r     = x[1:3]
    v     = x[4:6]
    m     = x[7] * mass_scaling_factor
    N_th  = x[8]

    x_e = [x_earth(t), y_earth(t), z_earth(t)]
    x_m = [x_moon(t), y_moon(t), z_moon(t)]
    x_s = [x_sun(t), y_sun(t), z_sun(t)]

    normrE = sqrt(sum((x_e - r).^2))
    normrM = sqrt(sum((x_m - r).^2))
    normrS = sqrt(sum((x_s - r).^2))

    normrSE = sqrt(sum((x_e - x_s).^2))
    normrMS = sqrt(sum((x_m - x_s).^2))

    dvg = mu_Earth / normrE^3 * (x_e - r) + mu_Moon / normrM^3 * (x_m - r) + mu_Sun / normrS^3 * (x_s - r) - (1 / (mass_earth + mass_moon)) * ( x_s * (mu_Sun * mass_earth / normrSE^3 + mu_Moon * mass_sun / normrMS^3) - x_m * (mu_Moon * mass_sun / normrMS^3) - x_e * (mu_Sun * mass_earth / normrSE^3)) 
    dvu = N_th * Tmax * u / m
    dv = dvg + dvu
    dm = - N_th * Q * sqrt(sum(u.^2)) / mass_scaling_factor

    dx = [ v; dv; dm; 0 ]
    return dx
end

t0_J200 = 1.107747244e9
tf_J200 = 1.110817407e9
N = 50000 
t_values = LinRange(t0_J200, tf_J200, N-3)
t_values = vcat(-5e9, 0.0, collect(t_values), tf_J200 * 5)
states_moon = zeros(6, N)
states_earth = zeros(6, N)
states_sun = zeros(6, N)
states_moon_wrt_sun = zeros(6, N)
states_earth_wrt_moon = zeros(6, N)
t0 = t0_J200 / TU
tf = tf_J200 / TU

for i = 1:N
    states_moon[:,i], _             = spkezr("MOON", t_values[i], "ECLIPJ2000", "NONE", "EARTH MOON BARYCENTER")
    states_earth[:,i], _            = spkezr("EARTH", t_values[i], "ECLIPJ2000", "NONE", "EARTH MOON BARYCENTER")
    states_sun[:,i], _              = spkezr("SUN", t_values[i], "ECLIPJ2000", "NONE", "EARTH MOON BARYCENTER")
    states_moon_wrt_sun[:,i], _     = spkezr("MOON", t_values[i], "ECLIPJ2000", "NONE", "SUN")
    states_earth_wrt_moon[:,i], _   = spkezr("EARTH", t_values[i], "ECLIPJ2000", "NONE", "MOON")
end


x_moon           = DataInterpolations.LinearInterpolation(states_moon[1,:] / LU,           t_values / TU) 
y_moon           = DataInterpolations.LinearInterpolation(states_moon[2,:] / LU,           t_values / TU) 
z_moon           = DataInterpolations.LinearInterpolation(states_moon[3,:] / LU,           t_values / TU) 

vx_moon          = DataInterpolations.LinearInterpolation(states_moon[4,:] / LU * TU,      t_values / TU) 
vy_moon          = DataInterpolations.LinearInterpolation(states_moon[5,:] / LU * TU,      t_values / TU) 
vz_moon          = DataInterpolations.LinearInterpolation(states_moon[6,:] / LU * TU,      t_values / TU) 

x_earth          = DataInterpolations.LinearInterpolation(states_earth[1,:] / LU,          t_values / TU) 
y_earth          = DataInterpolations.LinearInterpolation(states_earth[2,:] / LU,          t_values / TU) 
z_earth          = DataInterpolations.LinearInterpolation(states_earth[3,:] / LU,          t_values / TU) 

vx_earth         = DataInterpolations.LinearInterpolation(states_earth[4,:] / LU * TU,     t_values / TU) 
vy_earth         = DataInterpolations.LinearInterpolation(states_earth[5,:] / LU * TU,     t_values / TU) 
vz_earth         = DataInterpolations.LinearInterpolation(states_earth[6,:] / LU * TU,     t_values / TU) 

x_sun            = DataInterpolations.LinearInterpolation(states_sun[1,:] / LU,            t_values / TU) 
y_sun            = DataInterpolations.LinearInterpolation(states_sun[2,:] / LU,            t_values / TU) 
z_sun            = DataInterpolations.LinearInterpolation(states_sun[3,:] / LU,            t_values / TU) 

x_earth_wrt_moon = DataInterpolations.LinearInterpolation(states_earth_wrt_moon[1,:] / LU, t_values / TU)
y_earth_wrt_moon = DataInterpolations.LinearInterpolation(states_earth_wrt_moon[2,:] / LU, t_values / TU)
z_earth_wrt_moon = DataInterpolations.LinearInterpolation(states_earth_wrt_moon[3,:] / LU, t_values / TU)

x_moon_wrt_sun   = DataInterpolations.LinearInterpolation(states_moon_wrt_sun[1,:] / LU,   t_values / TU)
y_moon_wrt_sun   = DataInterpolations.LinearInterpolation(states_moon_wrt_sun[2,:] / LU,   t_values / TU)
z_moon_wrt_sun   = DataInterpolations.LinearInterpolation(states_moon_wrt_sun[3,:] / LU,   t_values / TU)


t1_J200      = 1.10829103e9
t1_moon, _   = spkezr("MOON", t1_J200, "ECLIPJ2000", "NONE", "EARTH MOON BARYCENTER")
t1           = t1_J200 / TU
occult1      = zeros(6)
occult1_ECI  = [-297664.1158..., 212871.4885..., -76.96308643..., 1.453863208..., -1.136754567..., -0.032783807...]
occult1[1:3] = (occult1_ECI[1:3] + t1_moon[1:3]) / LU
occult1[4:6] = (occult1_ECI[4:6] + t1_moon[4:6]) / LU * TU

t6_J200      = 1.110817407e9
t6_moon, _   = spkezr("MOON", t6_J200, "ECLIPJ2000", "NONE", "EARTH MOON BARYCENTER")
t6           = t6_J200 / TU
occult6      = zeros(6)
occult6_ECI  = [-365732.5404..., 39738.61049..., -83.77279923..., 1.854101596..., -0.310457164..., 0.010142901...]
occult6[1:3] = (occult6_ECI[1:3] + t6_moon[1:3]) / LU
occult6[4:6] = (occult6_ECI[4:6] + t6_moon[4:6]) / LU * TU




@def ocp begin
    t ∈ [ t1, t6 ], time
    x ∈ R^8, state
    u ∈ R³, control

    thrust = x[8]
    mass = x[7]

    -10 ≤ x₁(t) ≤ 10
    -10 ≤ x₂(t) ≤ 10
    -10 ≤ x₃(t) ≤ 10
    -10 ≤ x₄(t) ≤ 10
    -10 ≤ x₅(t) ≤ 10
    -10 ≤ x₆(t) ≤ 10
    0   ≤ x₇(t) ≤ 1 # careful here ! dry mass > 0   &&   dry_mass ≤ x₇(t) ≤ dry_mass + fuel_mass
    1e-5   ≤ x₈(t) ≤ 50

    0 ≤ sum(u(t).^2) ≤ 1

    x[1:6](t1) == occult1
    x₇(t1) == mass_init 
    x[1:6](t6) == occult6

    ẋ(t) == Fu(t, x(t), u(t)) 
    
    x₈(t6) → min

end


# Works but is slower:
# sol = solve(ocp, grid_size = 100, max_iter = 100, tol = 1e-6, acceptable_tol = 1e-5, disc_method=:gauss_legendre_2)


# Creating initial guess 
ppr = 1e-16 * ones(3)
x0 = [occult1..., 1...,  1..., ]
tspan = (t1, t6)
time_points = LinRange(t1, t6, 500)
prob = ODEProblem(propagate!, x0, tspan, ppr)
solpr = solve(prob, Tsit5(), saveat=time_points)

Nsolu = length(solpr.u)
plot([solpr.u[i][1] for i in range(1, Nsolu)], [solpr.u[i][2] for i in range(1, Nsolu)], [solpr.u[i][3] for i in range(1, Nsolu)])
xlims!(-2, 2)
ylims!(-2, 2)
zlims!(-2, 2)


sol = solve(ocp, 
        init = (time=solpr.t, state=solpr.u, control=ppr), 
        grid_size = 100, max_iter = 1000, tol = 1e-6, acceptable_tol = 1e-5,
        disc_method=:gauss_legendre_2)
    
initial_guess = sol

sol = solve(ocp, init = initial_guess, 
        grid_size = 150, max_iter = 1000, tol = 1e-6, disc_method=:gauss_legendre_2)

    
time_intervals = [t1, t6]

plot(sol,size = (600,1500))

control_matrix = zeros(3, length(sol.time_grid.value))
for i in range(1, length(sol.time_grid.value))
    control_matrix[:,i] =  sol.control.value(sol.time_grid.value[i])
end
control_norm = [norm(col) for col in eachcol(control_matrix)]
plot(sol.time_grid.value, control_norm, linewidth = 2, label = false)
ylims!(0,1)

moon_arcs = []
uncontrolled_arcs = []
plot_trajNV_multiarc_multidyn_withcoasting_1arc(sol, 1, time_intervals, moon_arcs, uncontrolled_arcs)

plot_trajNV_multiarc_multidyn_synodic_1arc(sol, 1, time_intervals, moon_arcs)