Refactor constants and simplify function returns

docs/src/index.md

@@ -49,18 +49,20 @@ using LinearAlgebra
 ## Problem definition
 
 ```@example main
-Tmax = 60                                  # Maximum thrust in Newtons
-cTmax = 3600^2 / 1e6; T = Tmax * cTmax     # Conversion from Newtons to kg x Mm / h²
-mass0 = 1500                               # Initial mass of the spacecraft
-β = 1.42e-02                               # Engine specific impulsion
-μ = 5165.8620912                           # Earth gravitation constant
-P0 = 11.625                                # Initial semilatus rectum
-ex0, ey0 = 0.75, 0                         # Initial eccentricity
-hx0, hy0 = 6.12e-2, 0                      # Initial ascending node and inclination
-L0 = π                                     # Initial longitude
-Pf = 42.165                                # Final semilatus rectum
-exf, eyf = 0, 0                            # Final eccentricity
-hxf, hyf = 0, 0                            # Final ascending node and inclination
+
+const mass0 = 1500                               # Initial mass of the spacecraft
+const β = 1.42e-02                               # Engine specific impulsion
+const μ = 5165.8620912                           # Earth gravitation constant
+const P0 = 11.625                                # Initial semilatus rectum
+const ex0, ey0 = 0.75, 0                         # Initial eccentricity
+const hx0, hy0 = 6.12e-2, 0                      # Initial ascending node and inclination
+const L0 = π                                     # Initial longitude
+const Pf = 42.165                                # Final semilatus rectum
+const exf, eyf = 0, 0                            # Final eccentricity
+const hxf, hyf = 0, 0                            # Final ascending node and inclination
+const Lf = 3π                                    # Estimation of final longitude for Tmax = 60
+const x0 = [P0, ex0, ey0, hx0, hy0, L0]          # Initial state
+const xf = [Pf, exf, eyf, hxf, hyf, Lf]          # Final state
 
 asqrt(x; ε=1e-9) = sqrt(sqrt(x^2+ε^2))     # Avoid issues with AD
 
@@ -70,20 +72,15 @@ function F0(x)
     cl = cos(L)
     sl = sin(L)
     w = 1 + ex * cl + ey * sl
-    F = zeros(eltype(x), 6) # Use eltype to allow overloading for AD
-    F[6] = w^2 / (P * pdm)
-    return F
+    return [0, 0, 0, 0, 0, w^2 / (P * pdm)]
 end
 
 function F1(x)
     P, ex, ey, hx, hy, L = x
     pdm = asqrt(P / μ)
     cl = cos(L)
     sl = sin(L)
-    F = zeros(eltype(x), 6)
-    F[2] = pdm *   sl
-    F[3] = pdm * (-cl)
-    return F
+    return pdm * [0, sl, -cl, 0, 0, 0]
 end
 
 function F2(x)
@@ -92,11 +89,7 @@ function F2(x)
     cl = cos(L)
     sl = sin(L)
     w = 1 + ex * cl + ey * sl
-    F = zeros(eltype(x), 6)
-    F[1] = pdm * 2 * P / w
-    F[2] = pdm * (cl + (ex + cl) / w)
-    F[3] = pdm * (sl + (ey + sl) / w)
-    return F
+    return pdm * [2P / w, cl + (ex + cl) / w, sl + (ey + sl) / w, 0, 0, 0]
 end
 
 function F3(x)
@@ -108,13 +101,7 @@ function F3(x)
     pdmw = pdm / w
     zz = hx * sl - hy * cl
     uh = (1 + hx^2 + hy^2) / 2
-    F = zeros(eltype(x), 6)
-    F[2] = pdmw * (-zz * ey)
-    F[3] = pdmw *   zz * ex
-    F[4] = pdmw *   uh * cl
-    F[5] = pdmw *   uh * sl
-    F[6] = pdmw *   zz
-    return F
+    return pdmw * [0, -zz * ey, zz * ex, uh * cl, uh * sl, zz]
 end
 
 nothing # hide
@@ -123,10 +110,9 @@ nothing # hide
 ## Direct solve
 
 ```@example main
+Tmax = 60                                    # Maximum thrust in Newtons
+cTmax = 3600^2 / 1e6; T = Tmax * cTmax       # Conversion from Newtons to kg x Mm / h²
 tf = 15                                      # Estimation of final time
-Lf = 3π                                      # Estimation of final longitude
-x0 = [P0, ex0, ey0, hx0, hy0, L0]            # Initial state
-xf = [Pf, exf, eyf, hxf, hyf, Lf]            # Final state
 x(t) = x0 + (xf - x0) * t / tf               # Linear interpolation
 u = [0.1, 0.5, 0.]                           # Initial guess for the control
 nlp_init = (state=x, control=u, variable=tf) # Initial guess for the NLP
@@ -146,7 +132,7 @@ end
 ```
 
 ```@example main
-nlp_sol = OptimalControl.solve(ocp; init=nlp_init, grid_size=100)
+nlp_sol = solve(ocp; init=nlp_init, grid_size=100)
 nothing # hide
 ```
 
@@ -193,8 +179,8 @@ nothing # hide
 ## Shooting (2/2), Tmax = 0.7 Newtons
 
 ```@example main
-hr = (t, x, p) -> begin # Regular maximised Hamiltonian (more efficient)
-    H0 = p' * F0(x)
+hr = (t, x, p) -> begin # Regular maximised Hamiltonian (more efficient...
+    H0 = p' * F0(x) # ... some computations could be factored)
     H1 = p' * F1(x)
     H2 = p' * F2(x)
     H3 = p' * F3(x)