Multi-Phase Two-Stage Rocket

Background

Rocket staging is a hybrid event: altitude and velocity remain continuous at separation, but vehicle mass changes instantaneously when empty hardware is discarded. Representing each powered stage as a separate phase makes both kinds of transition explicit and avoids approximating the mass jump with an artificial, extremely fast continuous process.

This example optimizes a vertical ascent to a prescribed apogee. All quantities are nondimensional so that the dynamics and nonlinear program remain well scaled.

Problem formulation

For stage \(i\in\{1,2\}\), the state is altitude \(h_i\), vertical velocity \(v_i\), and mass \(m_i\). The control \(\alpha_i\) is a throttle fraction. The dynamics are

\[ \dot h_i=v_i, \qquad \dot v_i=\frac{T_i\alpha_i}{m_i}-g, \qquad \dot m_i=-\beta_i\alpha_i, \]

with

\[ g=1, \qquad (T_1,\beta_1)=(2.40,0.080), \qquad (T_2,\beta_2)=(1.60,0.045). \]

The first phase spans \([0,t_s]\), and the second spans \([t_s,t_f]\). The objective balances final time and control effort:

\[ \min\quad t_f+0.04\left(\int_0^{t_s}\alpha_1^2\,\mathrm{d}t +\int_{t_s}^{t_f}\alpha_2^2\,\mathrm{d}t\right). \]

The mission starts from rest with unit mass and ends at a prescribed apogee:

\[ [h_1(0),v_1(0),m_1(0)]=[0,0,1], \]

\[ [h_2(t_f),v_2(t_f),m_2(t_f)]=[2,0,m_f]. \]

The mass inventory assigns \(0.06\) to first-stage propellant, \(0.20\) to discarded first-stage hardware, \(0.12\) to second-stage propellant, and \(0.62\) to second-stage dry mass. The first stage separates at burnout, so

\[ m_s^-=1-0.06=0.94, \qquad m_s^+=m_s^- -0.20=0.74. \]

At the shared separation time, altitude and velocity are continuous while mass follows this discrete reset:

\[ \begin{aligned} h_1(t_s)&=h_s=h_2(t_s), \\ v_1(t_s)&=v_s=v_2(t_s), \\ m_1(t_s)&=m_s^-, \\ m_2(t_s)&=m_s^+. \end{aligned} \]

The path bounds are

\[ 0\le\alpha_i\le1, \qquad 0\le h_i\le2, \qquad 0\le v_i\le3, \]

\[ 0.94\le m_1\le1, \qquad 0.62\le m_2\le0.74. \]

The remaining system bounds are

\[ \begin{aligned} 0.10&\le h_s\le1.80, & 0.05&\le v_s\le2.50, \\ 0.20&\le t_s\le3.00, & 0.62&\le m_f\le0.74, \\ 0.40&\le t_f-t_s\le5.00, & 1.00&\le t_f\le7.00. \end{aligned} \]

Variables and units

All values are nondimensional; each dash in the unit column denotes a scaled model quantity.

Symbol	Meaning	Unit
\(h_i\)	Altitude in stage \(i\)	-
\(v_i\)	Vertical velocity in stage \(i\)	-
\(m_i\)	Vehicle mass in stage \(i\)	-
\(\alpha_i\)	Throttle fraction in stage \(i\)	-
\(T_i\)	Full-throttle thrust coefficient	-
\(\beta_i\)	Full-throttle mass-flow coefficient	-
\(t_s,t_f\)	Separation and final times	-
\(m_s^-,m_s^+\)	Mass immediately before and after separation	-

Modeling choices

The shared static parameters \(h_s\), \(v_s\), and \(t_s\) connect the two Radau phases exactly. Separate parameters \(m_s^-\) and \(m_s^+\), joined by a system equality, encode the physical reset without incorrectly forcing mass continuity.

The first-stage burnout mass and second-stage dry mass come from the stated propellant inventory. Consequently, both stages contribute because the first stage alone does not have enough propellant to reach the target apogee; no artificial lower bound is placed on second-stage throttle. A small throttle-squared term regularizes the minimum-time objective.

Each phase uses 96 Radau mesh intervals with one collocation point per interval. This gives a piecewise-constant throttle and piecewise-linear state reconstruction. Consequently, the nodal bounds also hold between nodes instead of being vulnerable to high-order interpolation overshoot. The initial guess includes the exact 0.20 mass drop and continuous altitude and velocity. After optimization, the implementation independently checks every path bound at 2,001 uniformly spaced times in each phase.

Run the example

From the repository root, run:

python -m examples.multiphase_two_stage_rocket

To save the figure without opening a window:

python -m examples.multiphase_two_stage_rocket --save two-stage-rocket.png --no-show

Key implementation

The phase endpoints use shared system symbols for continuous states and distinct symbols for mass on either side of separation:

system = System(
    [
        "h_separation", "v_separation",
        "m_before_drop", "m_after_drop",
        "t_separation", "m_final", "t_final",
    ]
)
h_s, v_s, m_before, m_after, t_s, m_f, t_f = system.s

first_stage.set_boundary_condition(
    [0.0, 0.0, 1.0], [h_s, v_s, m_before], 0.0, t_s
)
second_stage.set_boundary_condition(
    [h_s, v_s, m_after], [2.0, 0.0, m_f], t_s, t_f
)
system.set_phase([first_stage, second_stage])

The burnout mass and discrete reset are system-level equalities. The final-mass lower bound represents second-stage dry mass:

system.set_system_constraint(
    [m_before, m_after - m_before, m_f, t_f - t_s],
    [0.94, -0.20, 0.62, 0.40],
    [0.94, -0.20, 0.74, 5.00],
)
system.set_objective(
    t_f + 0.04 * (first_stage.I[0] + second_stage.I[0])
)

The complete source includes all path and static-variable bounds, a dynamically plausible initial guess, solver status checks, and numerical validation of phase continuity, propellant limits, and the mass reset.

Verified result

The verified solution has objective \(2.64500359\), separation time \(t_s=0.750000\), and final time \(t_f=2.601708\). Separation occurs at \(h_s=0.403132\) and \(v_s=1.105664\). The second stage uses \(0.015374\) mass units of propellant and then coasts to the target apogee with final mass \(m_f=0.724626\).

Altitude and velocity match exactly at the reported precision, and the computed separation mass drop is \(0.200000\). A dense reconstruction confirms zero path-bound violation at the reported precision; the maximum velocity is \(1.495408<3\), and the minimum second-stage mass remains \(0.104626\) above its dry-mass limit. In the figure, the dashed line marks separation and the dotted segment shows the instantaneous mass reset.

Source code

See the complete runnable example: examples/multiphase_two_stage_rocket.py.