BUG: populate surfaces_cp_to_cdm in flight_rocket()

flight_rocket() copies aerodynamic surfaces directly onto the
composed Rocket (bypassing add_surfaces() to avoid double-
transforming fin leading-edge positions - see the comment at the
copy site). That also skipped evaluate_surfaces_cp_to_cdm(), a
step add_surfaces() normally triggers and that
Flight.u_dot_generalized requires to apply aerodynamic forces
during a real 6DOF simulation.

Every existing test flew bare rockets with no surfaces, so this
never got exercised - only surfaced when actually running a stage
with real aerodynamic surfaces through Mission end-to-end (via the
new visualization notebook work). KeyError on
rocket.surfaces_cp_to_cdm[aero_surface] the moment such a Flight's
ODE integration starts.

ENH: add MultiStageRocket.draw() and Mission.all_flights

Visualization support, reusing RocketPy's existing plotting
entirely - no new drawing or comparison code:

- MultiStageRocket.draw() composes every stage via flight_rocket()
  and hands the result to Rocket's own plots.draw() unchanged. A
  real multi-diameter stack renders correctly for free, since
  Rocket's tube-drawing already handles a radius change between
  surfaces.
- Mission.all_flights collects every Flight Mission runs, in
  execution order, each exactly once (a Flight shared by several
  bodies, e.g. the full stack, isn't duplicated) - built
  specifically to feed CompareFlights directly, matching
  mission_multistage_design.md's stated design ("feeds
  CompareFlights").
- Every Flight Mission creates now gets a distinguishing name
  (e.g. "booster+sustainer", "booster", "sustainer") instead of
  Flight's own "Flight" default, so CompareFlights legends are
  actually readable with more than one flight in the comparison.

Also updates docs/notebooks/multistage_mission.ipynb with draw()
and trajectory-comparison demos, executed end-to-end. Departs from
this notebook's earlier convention of clearing outputs: plot image
outputs are now kept, matching docs/notebooks/getting_started.ipynb
rather than the print-only utilities_usage.ipynb, since the whole
point of these cells is the rendered image.

Author: aZira371

docs/notebooks/multistage_mission.ipynb

@@ -309,6 +309,11 @@ def flight_rocket(self, active_stages, carried_deployables=()):
         composed_rocket.evaluate_center_of_pressure()
         composed_rocket.evaluate_stability_margin()
         composed_rocket.evaluate_static_margin()
+        # add_motor() already populated surfaces_cp_to_cdm, but from
+        # before these surfaces were copied in above (it was empty at
+        # that point) - Flight.u_dot_generalized needs this dict to
+        # apply aerodynamic forces during a real 6DOF simulation.
+        composed_rocket.evaluate_surfaces_cp_to_cdm()
 
         return composed_rocket
 
@@ -377,3 +382,13 @@ def _derive_stack_drag(self, active_stages, attr_name, stack_radius):
             )
             combined = scaled if combined is None else combined + scaled
         return combined
+
+    def draw(self, vis_args=None, plane="xz", *, filename=None):
+        """Draw the stacked vehicle: every stage's aerodynamic surfaces,
+        combined exactly as flight_rocket() would compose them for a
+        flight with every stage attached. Reuses Rocket's own drawing
+        code unchanged - see :meth:`Rocket.plots.draw` for parameters
+        and the "at least one aerodynamic surface" requirement.
+        """
+        stack = self.flight_rocket(active_stages=tuple(self.stages))
+        stack.plots.draw(vis_args, plane, filename=filename)

rocketpy/rocket/multistage.py

@@ -114,8 +114,18 @@ def __init__(
 
         self.flights = {}
         self.timeline = []
+        self._all_flights = []
         self._simulate()
 
+    @property
+    def all_flights(self):
+        """Every Flight object, in execution order, each appearing once
+        (a Flight shared by several bodies, e.g. the full stack, is not
+        repeated). Feeds CompareFlights directly, e.g.
+        ``CompareFlights(mission.all_flights).trajectories_3d()``.
+        """
+        return list(self._all_flights)
+
     def _simulate(self):
         single_stage = len(self.vehicle.stages) == 1
         no_deployables = not self.vehicle.deployables
@@ -145,7 +155,7 @@ def _simulate_single_stage(self):
         self.timeline.append((0.0, f"ignition:{stage.name}"))
         self.timeline.append((0.0, "liftoff"))
 
-        flight = self._run_flight(rocket)
+        flight = self._run_flight(rocket, name=stage.name)
 
         self.timeline.append((flight.out_of_rail_time, "rail_departure"))
         self.timeline.append((flight.t_final, f"impact:{stage.name}"))
@@ -188,7 +198,11 @@ def _run_carrier_phase(self, stage, deployable):
         self.timeline.append((0.0, f"ignition:{stage.name}"))
         self.timeline.append((0.0, "liftoff"))
 
-        carrier_flight = self._run_flight(carrier_rocket, terminate_on_apogee=True)
+        carrier_flight = self._run_flight(
+            carrier_rocket,
+            name=f"{stage.name}+{deployable.name}",
+            terminate_on_apogee=True,
+        )
         self.timeline.append((carrier_flight.out_of_rail_time, "rail_departure"))
         self.timeline.append(
             (carrier_flight.apogee_time, f"ejection:{deployable.name}")
@@ -207,7 +221,7 @@ def _run_stage_after_ejection_phase(
             ending_state, carrier_rocket, stage_rocket, delta_v
         )
         stage_flight = self._run_flight(
-            stage_rocket, initial_solution=initial_solution
+            stage_rocket, name=stage.name, initial_solution=initial_solution
         )
         self.timeline.append((stage_flight.t_final, f"impact:{stage.name}"))
         self.flights[stage.name].append(stage_flight)
@@ -218,7 +232,9 @@ def _run_deployable_phase(self, deployable, carrier_rocket, ending_state, delta_
             ending_state, carrier_rocket, deployable.free_rocket, delta_v
         )
         deployable_flight = self._run_flight(
-            deployable.free_rocket, initial_solution=initial_solution
+            deployable.free_rocket,
+            name=deployable.name,
+            initial_solution=initial_solution,
         )
         self.timeline.append(
             (deployable_flight.t_final, f"impact:{deployable.name}")
@@ -264,7 +280,11 @@ def _run_stack_phase(self, booster, sustainer):
         self.timeline.append((0.0, f"ignition:{booster.name}"))
         self.timeline.append((0.0, "liftoff"))
 
-        stack_flight = self._run_flight(stack_rocket, max_time=separation_time)
+        stack_flight = self._run_flight(
+            stack_rocket,
+            name=f"{booster.name}+{sustainer.name}",
+            max_time=separation_time,
+        )
         self.timeline.append((stack_flight.out_of_rail_time, "rail_departure"))
         self.timeline.append((separation_time, f"separation:{booster.name}"))
 
@@ -279,7 +299,7 @@ def _run_booster_phase(self, booster, stack_rocket, ending_state, delta_v):
             ending_state, stack_rocket, booster_rocket, delta_v
         )
         booster_flight = self._run_flight(
-            booster_rocket, initial_solution=initial_solution
+            booster_rocket, name=booster.name, initial_solution=initial_solution
         )
         self.timeline.append((booster_flight.t_final, f"impact:{booster.name}"))
         self.flights[booster.name].append(booster_flight)
@@ -298,16 +318,28 @@ def _run_sustainer_phase(
             ending_state, stack_rocket, sustainer_rocket, delta_v
         )
         sustainer_flight = self._run_flight(
-            sustainer_rocket, initial_solution=initial_solution
+            sustainer_rocket, name=sustainer.name, initial_solution=initial_solution
         )
         self.timeline.append((sustainer_flight.t_final, f"impact:{sustainer.name}"))
         self.flights[sustainer.name].append(sustainer_flight)
 
     def _run_flight(
-        self, rocket, initial_solution=None, max_time=None, terminate_on_apogee=False
+        self,
+        rocket,
+        name,
+        initial_solution=None,
+        max_time=None,
+        terminate_on_apogee=False,
     ):
-        """Run one Flight in absolute mission time."""
-        return Flight(
+        """Run one Flight in absolute mission time.
+
+        The single choke point through which every Flight Mission creates
+        is constructed, so this is also where all_flights collects them,
+        in execution order, each exactly once. ``name`` distinguishes
+        each Flight in plots such as CompareFlights, which otherwise
+        labels every line "Flight" (Flight's own default).
+        """
+        flight = Flight(
             rocket=rocket,
             environment=self.environment,
             rail_length=self.rail_length,
@@ -321,7 +353,10 @@ def _run_flight(
             time_overshoot=self.time_overshoot,
             ode_solver=self.ode_solver,
             verbose=self.verbose,
+            name=name,
         )
+        self._all_flights.append(flight)
+        return flight
 
     @staticmethod
     def _momentum_split(mass_a, mass_b, delta_v):

rocketpy/simulation/mission.py

@@ -218,3 +218,42 @@ def test_flight_rocket_combines_surfaces_of_every_active_stage(calisto_nose_cone
 
     assert len(two_stage_rocket.aerodynamic_surfaces) == 1
     assert len(sustainer_alone_rocket.aerodynamic_surfaces) == 1
+
+
+def test_flight_rocket_populates_surfaces_cp_to_cdm(calisto, calisto_nose_cone):
+    # surfaces_cp_to_cdm is separate from center_of_pressure/stability -
+    # it's what Flight.u_dot_generalized looks up per surface to apply
+    # aerodynamic forces during a real 6DOF simulation. flight_rocket()
+    # copies surfaces directly (bypassing add_surfaces() to avoid
+    # double-transforming fin leading-edge positions), which must not
+    # skip populating this dict too.
+    stage = Stage(name="stage_1", rocket=calisto)
+    stage.rocket.add_surfaces(calisto_nose_cone, 1.0)
+    vehicle = MultiStageRocket(stages=[stage])
+
+    composed_rocket = vehicle.flight_rocket(active_stages=(stage,))
+
+    assert calisto_nose_cone in composed_rocket.surfaces_cp_to_cdm
+
+
+def test_draw_runs_for_a_stage_with_aerodynamic_surfaces(calisto_robust):
+    stage = Stage(name="stage_1", rocket=calisto_robust)
+    vehicle = MultiStageRocket(stages=[stage])
+
+    assert vehicle.draw(filename=None) is None
+
+
+def test_draw_combines_surfaces_of_every_stage(calisto_nose_cone):
+    booster, sustainer = _two_stage_vehicle()
+    booster.rocket.add_surfaces(calisto_nose_cone, 0.5)
+    vehicle = MultiStageRocket(stages=[booster, sustainer])
+
+    assert vehicle.draw(filename=None) is None
+
+
+def test_draw_raises_when_no_stage_has_aerodynamic_surfaces():
+    booster, sustainer = _two_stage_vehicle()
+    vehicle = MultiStageRocket(stages=[booster, sustainer])
+
+    with pytest.raises(ValueError):
+        vehicle.draw(filename=None)

tests/unit/rocket/test_multistage.py

@@ -3,6 +3,7 @@
 from rocketpy import Flight
 from rocketpy.mathutils.vector_matrix import Matrix, Vector
 from rocketpy.motors.point_mass_motor import PointMassMotor
+from rocketpy.plots.compare.compare_flights import CompareFlights
 from rocketpy.rocket.multistage import MultiStageRocket, Stage
 from rocketpy.rocket.rocket import Rocket
 from rocketpy.simulation.mission import Mission
@@ -377,3 +378,77 @@ def test_deployable_splits_separation_delta_v_by_momentum_conservation(
     assert payload_flight.solution[0][6] == pytest.approx(
         ending_vz + expected_payload_delta_v
     )
+
+
+def test_all_flights_is_the_single_flight_for_the_degenerate_case(
+    calisto, example_plain_env
+):
+    mission = Mission(
+        vehicle=calisto,
+        environment=example_plain_env,
+        rail_length=5.2,
+        inclination=85,
+        heading=0,
+    )
+
+    assert mission.all_flights == [mission.flights["stage_1"][0]]
+
+
+def test_all_flights_has_no_duplicate_for_the_shared_stack_flight(example_plain_env):
+    booster, sustainer = _two_stage_vehicle(booster_separation=0.5)
+    vehicle = MultiStageRocket(stages=[booster, sustainer])
+
+    mission = Mission(
+        vehicle=vehicle,
+        environment=example_plain_env,
+        rail_length=1.0,
+        inclination=90,
+        heading=0,
+        max_time=5,
+    )
+
+    assert len(mission.all_flights) == 3
+    assert len(set(id(flight) for flight in mission.all_flights)) == 3
+    assert mission.all_flights == [
+        mission.flights["booster"][0],
+        mission.flights["booster"][1],
+        mission.flights["sustainer"][1],
+    ]
+
+
+def test_all_flights_feeds_compare_flights(example_plain_env):
+    booster, sustainer = _two_stage_vehicle(booster_separation=0.5)
+    vehicle = MultiStageRocket(stages=[booster, sustainer])
+
+    mission = Mission(
+        vehicle=vehicle,
+        environment=example_plain_env,
+        rail_length=1.0,
+        inclination=90,
+        heading=0,
+        max_time=5,
+    )
+
+    assert CompareFlights(mission.all_flights).trajectories_3d(filename=None) is None
+
+
+def test_flight_names_distinguish_each_body_and_phase(example_plain_env):
+    # CompareFlights labels each line in a plot legend using flight.name;
+    # every Flight Mission creates must have a name that distinguishes it
+    # (Flight's own default "Flight" would make every legend entry
+    # identical and useless).
+    booster, sustainer = _two_stage_vehicle(booster_separation=0.5)
+    vehicle = MultiStageRocket(stages=[booster, sustainer])
+
+    mission = Mission(
+        vehicle=vehicle,
+        environment=example_plain_env,
+        rail_length=1.0,
+        inclination=90,
+        heading=0,
+        max_time=5,
+    )
+
+    names = [flight.name for flight in mission.all_flights]
+    assert len(names) == len(set(names))
+    assert all(name != "Flight" for name in names)