OptimalControl.jl icon indicating copy to clipboard operation
OptimalControl.jl copied to clipboard
verified publisher icon control-toolbox

[General] Code Tmax sun occult mission

Open alesiagr opened this issue 3 months ago • 2 comments

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)

alesiagr avatar Sep 05 '25 13:09 alesiagr

nice @alesiagr can you please

  • confirm that the 9 functions x/y/z_earth/moon/sun are the only external functions (that is, not basic Julia functions known to and overloaded by ExaModels)?
  • provide a simple way (directly from DataInterpolations.jl?) to retrieve their first and second order derivatives? (see https://exanauts.github.io/ExaModels.jl/stable/core/#ExaModels.@register_univariate-Tuple)

given the complexity of Fu, see also https://github.com/control-toolbox/OptimalControl.jl/issues/620

jbcaillau avatar Sep 06 '25 15:09 jbcaillau

@jbcaillau

  1. yes, these are the only external functions
  2. DataInterpolations.jl is compatible with autodiff. Can it be retrieved via ForwardDiff or something like that ?

alesiagr avatar Sep 08 '25 11:09 alesiagr