powered_descent.jl

# Açıkmeşe & Ploen (2007), J. Guidance, Control, and Dynamics 30(5):1353–1366
# "Convex Programming Approach to Powered Descent Guidance for Mars Landing"
#
# A Julia/JuMP implementation of Problem 4 (the discretised second-order cone
# program).  Algorithm 1 of the paper line-searches the integer number of
# time-steps N to find the optimal time-of-flight; the cost is unimodal in N,
# so we scan upward from N_min and stop once the objective has been
# increasing for a few consecutive steps.
#
# Run with:
#   julia --project=. powered_descent.jl
# (or just:  julia powered_descent.jl)

using JuMP
using ECOS
using Optim
using LinearAlgebra
using Printf
ENV["GKSwstype"] = "100"            # offscreen GR — no gksqt popup on macOS
using Plots

# ─────────────────── Sec. V example parameters (eq. 59) ────────────────────

const g_vec = [-3.7114, 0.0, 0.0]   # Mars gravity, PSF frame [m/s²]
const g_e   = 9.807                 # Earth gravity magnitude [m/s²]

const m_dry = 1505.0                # dry mass [kg]
const m_wet = 1905.0                # wet (initial) mass [kg]
const I_sp  = 225.0                 # specific impulse [s]
const T_bar = 3100.0                # max thrust per engine [N]
const n_th  = 6                     # number of engines
const φ     = deg2rad(27.0)         # cant angle [rad]

const T_1 = 0.3 * T_bar             # min throttle per engine
const T_2 = 0.8 * T_bar             # max throttle per engine
const ρ_1 = n_th * T_1 * cos(φ)     # min net thrust [N]
const ρ_2 = n_th * T_2 * cos(φ)     # max net thrust [N]
const α   = 1 / (I_sp * g_e * cos(φ))  # mass-flow coefficient [s/m]

# Initial state for Figs. 5–7 of the paper
const r0 = [1500.0, 0.0, 2000.0]    # [m]
const v0 = [-75.0,  0.0,  100.0]    # [m/s]

# Final thrust pointing direction (eq. 12) — opposite gravity (here, +x)
const n_f = [1.0, 0.0, 0.0]

# Glide-slope angle (eq. 11).  Set ≥ 90° to disable.
const γ_gs_deg = 86.0

# ───────────────────────────────────────────────────────────────────────────
# build_and_solve(N, Δt)
#
# Builds and solves Problem 4 for a fixed number of intervals N (so t_f =
# N·Δt).  Uses zero-order-hold dynamics — A_c² = 0 for the [r;v;z] system,
# so exp(A_c Δt) = I + A_c Δt and the closed-form ZOH update is exact.
# ───────────────────────────────────────────────────────────────────────────
function build_and_solve(N::Int, Δt::Float64;
                         γ_gs_deg::Float64 = γ_gs_deg,
                         enforce_final_pointing::Bool = true)

    model = Model(ECOS.Optimizer)
    set_silent(model)

    @variable(model, r[1:3, 0:N])    # position
    @variable(model, v[1:3, 0:N])    # velocity
    @variable(model, z[0:N])         # log(mass)
    @variable(model, u[1:3, 0:N])    # u  ≜ T_c / m   (eq. 25)
    @variable(model, σ[0:N])         # σ  ≜ Γ   / m

    # Boundary conditions (eqs. 4–5)
    @constraint(model, r[:, 0] .== r0)
    @constraint(model, v[:, 0] .== v0)
    @constraint(model, z[0]     == log(m_wet))
    @constraint(model, r[:, N] .== 0)
    @constraint(model, v[:, N] .== 0)
    if enforce_final_pointing
        @constraint(model, u[:, N] .== σ[N] * n_f)   # eq. 37
    end

    # Closed-form ZOH dynamics (eqs. 26, 31)
    for k = 0:N-1
        @constraint(model, r[:, k+1] .== r[:, k] + Δt*v[:, k] +
                                          (Δt^2/2) * (u[:, k] + g_vec))
        @constraint(model, v[:, k+1] .== v[:, k] + Δt * (u[:, k] + g_vec))
        @constraint(model, z[k+1]    == z[k]     - α * Δt * σ[k])
    end

    for k = 0:N
        t_k = k * Δt

        # ‖u_k‖ ≤ σ_k                                              (eq. 50)
        @constraint(model, [σ[k]; u[:, k]] in SecondOrderCone())

        # Convex relaxation of σ ∈ [ρ_1/m, ρ_2/m]                  (eq. 51)
        z0_k = log(m_wet - α * ρ_2 * t_k)
        μ1_k = ρ_1 * exp(-z0_k)            # = ρ_1 / (m_wet − αρ_2 t_k)
        μ2_k = ρ_2 * exp(-z0_k)
        δ    = z[k] - z0_k

        #   σ_k ≥ μ1_k (1 − δ + δ²/2)
        # ⇔ 2 (σ_k − μ1_k(1−δ)) · 1  ≥  (√μ1_k δ)²        (rotated SOC)
        @constraint(model, [σ[k] - μ1_k*(1 - δ); 1.0; sqrt(μ1_k) * δ]
                            in MOI.RotatedSecondOrderCone(3))
        #   σ_k ≤ μ2_k (1 − δ)                                     (linear)
        @constraint(model, σ[k] ≤ μ2_k * (1 - δ))

        # Mass physical bounds (eq. 52); ignored at k=0 since z(0)=ln m_wet
        if k ≥ 1
            @constraint(model, z[k] ≥ log(m_wet - α * ρ_2 * t_k))
            @constraint(model, z[k] ≤ log(m_wet - α * ρ_1 * t_k))
        end

        # Glide-slope SOC: ‖[r₂; r₃]‖ ≤ tan(γ) r₁                   (eq. 11)
        # When the glide slope is on, r₁ ≥ 0 is implied by the SOC; otherwise
        # we enforce the no-subsurface-flight constraint (eq. 7) explicitly.
        if γ_gs_deg < 90.0
            tg = tan(deg2rad(γ_gs_deg))
            @constraint(model, [tg * r[1, k]; r[2, k]; r[3, k]] in SecondOrderCone())
        else
            @constraint(model, r[1, k] ≥ 0)
        end
    end

    # Cost ≈ Δt Σ_{k=0}^{N-1} σ_k   — left Riemann sum (eq. 47)
    @objective(model, Min, Δt * sum(σ[k] for k = 0:N-1))

    optimize!(model)

    status = termination_status(model)
    ok = status ∈ (MOI.OPTIMAL, MOI.LOCALLY_SOLVED, MOI.ALMOST_OPTIMAL)
    if !ok
        return (status = status, ok = false, obj = NaN,
                r = nothing, v = nothing, z = nothing,
                u = nothing, σ = nothing)
    end

    asmat(x) = [value(x[i, k]) for i = 1:3, k = 0:N]
    asvec(x) = [value(x[k])     for k = 0:N]

    return (status = status, ok = true,
            obj = objective_value(model),
            r   = asmat(r),
            v   = asmat(v),
            z   = asvec(z),
            u   = asmat(u),
            σ   = asvec(σ))
end


# ─────────────────────── line search over N (Algorithm 1) ───────────────────
#
# The fuel cost is unimodal in t_f (paper, Remark 10), so we use Brent's
# method (Optim.jl) instead of a brute-force scan.  Each evaluation rounds
# to the nearest integer N and is cached so the parabolic-interpolation
# steps never re-solve the same SOCP.  Infeasible N's return a large finite
# penalty rather than +∞, which keeps Brent's interpolation well-behaved.

const Δt = 1.0

# Conservative bounds on t_f from physical limits (eq. 55)
const m_fuel = m_wet - m_dry
const t_min  = m_fuel * norm(v0) / ρ_2
const t_max  = m_fuel / (α * ρ_1)
const N_min  = max(2, Int(floor(t_min / Δt)) + 1)
const N_max  = Int(floor(t_max / Δt))

@printf("\nMars pinpoint landing — Açıkmeşe & Ploen (2007)\n")
@printf("Search range: N ∈ [%d, %d]   (t_f ∈ [%.1f, %.1f] s)\n\n",
        N_min, N_max, N_min*Δt, N_max*Δt)

const _cache = Dict{Int, Any}()

function eval_N(N_real::Real)
    N   = clamp(round(Int, N_real), N_min, N_max)
    res = get!(() -> build_and_solve(N, Δt), _cache, N)
    if res.ok
        fuel = m_wet - exp(res.z[end])
        @printf("  Brent: N=%3d  obj=%9.4f  fuel=%6.2f kg\n", N, res.obj, fuel)
        return res.obj
    else
        @printf("  Brent: N=%3d  %s\n", N, res.status)
        return 1.0e10                    # large but finite — keeps Brent stable
    end
end

result = optimize(eval_N, float(N_min), float(N_max), Brent();
                  abs_tol = 0.5)        # integer-grain tolerance

# Pick the best feasible N from the cache (defensive — usually == minimizer)
best = (obj = Inf, N = 0, sol = nothing)
for (N, res) in _cache
    res.ok || continue
    res.obj < best.obj && (global best = (obj = res.obj, N = N, sol = res))
end
best.sol === nothing && error("Brent failed to find a feasible N")

@printf("\nBrent took %d SOCP solves (vs ~%d for a brute-force scan)\n",
        length(_cache), N_max - N_min + 1)
@printf("Optimal: N* = %d   t_f* = %.1f s   fuel = %.2f kg\n",
        best.N, best.N * Δt, m_wet - exp(best.sol.z[end]))


# ───────────────────────────────── plotting ─────────────────────────────────

N_opt = best.N
sol   = best.sol
ts    = collect(0:N_opt) .* Δt
ms    = exp.(sol.z)
Tnet  = sol.u .* permutedims(ms)                # net thrust force [N]
Tmag  = [norm(Tnet[:, k+1]) for k = 0:N_opt]
# Per-thruster throttle level T/T_bar = Tmag / (n cos φ T_bar) ∈ [0.3, 0.8]
throttle = Tmag ./ (n_th * T_bar * cos(φ))
θ_deg = [Tmag[k+1] > 1e-6 ?
         acosd(clamp(dot(Tnet[:, k+1], -g_vec) /
                     (Tmag[k+1] * norm(g_vec)), -1.0, 1.0)) : 0.0
         for k = 0:N_opt]

p1 = plot(ts, sol.r', label = ["r₁" "r₂" "r₃"],
          xlabel = "t (s)", ylabel = "position (m)")
p2 = plot(ts, sol.v', label = ["v₁" "v₂" "v₃"],
          xlabel = "t (s)", ylabel = "velocity (m/s)")
# Net acceleration r̈ = u + g (eq. 1).  sol.u alone is the thrust contribution.
accel = sol.u .+ g_vec
p3 = plot(ts, accel', label = ["a₁" "a₂" "a₃"],
          xlabel = "t (s)", ylabel = "acceleration (m/s²)")
p4 = plot(ts, Tnet'/1000, label = ["T₁" "T₂" "T₃"],
          xlabel = "t (s)", ylabel = "net force (kN)")
p5 = plot(ts, throttle, legend = false,
          xlabel = "t (s)", ylabel = "throttle level")
p6 = plot(ts, θ_deg, legend = false,
          xlabel = "t (s)", ylabel = "θ (deg)")
plt = plot(p1, p2, p3, p4, p5, p6,
           layout = (3, 2), size = (900, 700),
           plot_title = @sprintf("t_f* = %.0f s,  fuel = %.1f kg,  γ_gs = %s",
                                  N_opt*Δt,  m_wet - ms[end],
                                  γ_gs_deg < 90.0 ? "$(γ_gs_deg)°" : "off"))
savefig(plt, "trajectory.png")
@printf("\nSaved time-history plot to trajectory.png\n")

# Vertical-vs-horizontal trajectory (Fig. 8 of the paper)
horiz = sqrt.(sol.r[2, :].^2 .+ sol.r[3, :].^2)
plt2 = plot(horiz, sol.r[1, :], legend = false,
            xlabel = "horizontal distance (m)",
            ylabel = "altitude (m)",
            title = "Trajectory" *
                    (γ_gs_deg < 90.0 ?
                     "  (glide slope $(γ_gs_deg)°)" : ""))
savefig(plt2, "trajectory_xy.png")
@printf("Saved trajectory plot to trajectory_xy.png\n")