ENH : Multistage mission architecture implementation

docs/user/index.rst

@@ -30,6 +30,7 @@ RocketPy's User Guide
    Deployable Payload <deployable.rst>
    Controllers <controllers.rst>
    Air Brakes Example <airbrakes.rst>
+   Multistage Mission <../notebooks/multistage_mission.ipynb>
    ../notebooks/sensors.ipynb
    ../matlab/matlab.rst
 

rocketpy/plots/rocket_plots.py

@@ -9,6 +9,24 @@
 from .plot_helpers import show_or_save_plot
 
 
+def _default_vis_args():
+    """The same default ``vis_args`` draw() builds when none is given,
+    as a fresh dict each call - shared so other callers that draw onto
+    an existing Axes (e.g. MultiStageRocket.draw_motor()) match draw()'s
+    own defaults instead of drifting from them independently.
+    """
+    return {
+        "background": "#EEEEEE",
+        "tail": "black",
+        "nose": "black",
+        "body": "black",
+        "fins": "black",
+        "motor": "black",
+        "buttons": "black",
+        "line_width": 1.0,
+    }
+
+
 class _RocketPlots:
     """Class that holds plot methods for Rocket class.
 
@@ -150,7 +168,7 @@ def thrust_to_weight(self):
             lower=0, upper=self.rocket.motor.burn_out_time
         )
 
-    def draw(self, vis_args=None, plane="xz", *, filename=None):
+    def draw(self, vis_args=None, plane="xz", *, filename=None, return_axes=False):
         """Draws the rocket in a matplotlib figure.
 
         Parameters
@@ -182,21 +200,17 @@ def draw(self, vis_args=None, plane="xz", *, filename=None):
             the plot will be shown instead of saved. Supported file endings are:
             eps, jpg, jpeg, pdf, pgf, png, ps, raw, rgba, svg, svgz, tif, tiff
             and webp (these are the formats supported by matplotlib).
+        return_axes : bool, optional
+            If ``True``, skip showing/saving the plot and return the
+            matplotlib ``Axes`` instead, so a caller can add its own
+            annotations before showing/saving it. Default ``False``
+            (existing behavior, unchanged).
         """
 
         self.__validate_aerodynamic_surfaces(plane)
 
         if vis_args is None:
-            vis_args = {
-                "background": "#EEEEEE",
-                "tail": "black",
-                "nose": "black",
-                "body": "black",
-                "fins": "black",
-                "motor": "black",
-                "buttons": "black",
-                "line_width": 1.0,
-            }
+            vis_args = _default_vis_args()
 
         _, ax = plt.subplots(figsize=(8, 6), facecolor=vis_args["background"])
         ax.set_aspect("equal")
@@ -220,7 +234,30 @@ def draw(self, vis_args=None, plane="xz", *, filename=None):
         plt.ylabel("Radius (m)")
         plt.legend(bbox_to_anchor=(1.05, 1), loc="upper left")
         plt.tight_layout()
+
+        if return_axes:
+            return ax
         show_or_save_plot(filename)
+        return None
+
+    def draw_motor(self, ax, vis_args=None):
+        """Draw just this rocket's own motor (grains/chamber/nozzle)
+        onto an existing Axes - no aerodynamic surfaces, no connecting
+        body-tube segment down to it.
+
+        For a caller that already has its own Axes (e.g. via
+        ``draw(..., return_axes=True)``) and only needs one more motor
+        added to it - MultiStageRocket.draw() is exactly this case: a
+        composed multi-stage Rocket can only ever carry ONE active
+        motor (RocketPy's Rocket supports a single motor), so every
+        stage past the currently-firing one has its own motor riding
+        along inert, invisible to the composed Rocket's own draw().
+        Reuses the same per-motor-type patch generation draw() itself
+        uses, rather than a second implementation of it.
+        """
+        if vis_args is None:
+            vis_args = _default_vis_args()
+        self._draw_motor(self.rocket.radius, self.rocket.motor_position, ax, vis_args)
 
     def __validate_aerodynamic_surfaces(self, plane):
         if not self.rocket.aerodynamic_surfaces:

rocketpy/rocket/multistage/__init__.py

@@ -0,0 +1,7 @@
+"""Vehicle composition layer for multistage rockets and deployable payloads."""
+
+from rocketpy.rocket.multistage.deployable import Deployable
+from rocketpy.rocket.multistage.geometry import axial_extent
+from rocketpy.rocket.multistage.multistage_rocket import MultiStageRocket
+from rocketpy.rocket.multistage.separable_body import SeparableBody
+from rocketpy.rocket.multistage.stage import Stage

rocketpy/rocket/multistage/deployable.py

@@ -0,0 +1,80 @@
+"""An inert body the vehicle carries and ejects in flight."""
+
+from rocketpy.rocket.multistage.separable_body import SeparableBody
+
+
+class Deployable(SeparableBody):
+    """An inert body carried by the vehicle and ejected in flight.
+
+    No aerodynamic identity while attached: contributes only mass and
+    inertia at a position inside the carrying stage.
+
+    Free-flight aerodynamics come from one of two sources, mutually
+    exclusive:
+    - ``free_rocket``: a fully built Rocket or PointMassRocket (full
+      control);
+    - surfaces added with ``add_surface()``: Mission assembles the
+      free-flight Rocket from the deployable's mass, inertia, radius and
+      the added surfaces.
+
+    Parameters
+    ----------
+    name : str
+        Unique body name; used to group Mission results.
+    mass : float
+        Carried mass in kg.
+    inertia : tuple of float
+        Inertia (I11, I22, I33) about the deployable's own center of
+        mass, in kg*m^2.
+    position : float
+        Position of the deployable's center of mass in the carrying
+        stage's rocket coordinate system, in meters.
+    radius : float, optional
+        The deployable's largest radius in meters. Required only when
+        defining its free-flight aerodynamics via add_surface().
+    free_rocket : Rocket, PointMassRocket, optional
+        Free-flight configuration after ejection. Mutually exclusive
+        with add_surface().
+    ejection : Event, optional
+        When the deployable is released, e.g. Event(trigger="apogee").
+    separation_delta_v : float, optional
+        See SeparableBody.
+    """
+
+    def __init__(
+        self,
+        name,
+        mass,
+        inertia,
+        position,
+        radius=None,
+        free_rocket=None,
+        ejection=None,
+        separation_delta_v=0.0,
+    ):
+        super().__init__(name=name, separation_delta_v=separation_delta_v)
+        self.mass = mass
+        self.inertia = inertia
+        self.position = position
+        self.radius = radius
+        self.free_rocket = free_rocket
+        self.ejection = ejection
+        self.surfaces = []
+
+    def add_surface(self, surface, position):
+        """Add an aerodynamic surface to the deployable's free flight.
+
+        Takes effect only after ejection; while attached the deployable
+        still contributes only mass and inertia. ``position`` is in the
+        deployable's own coordinate system, in meters. Requires
+        ``radius`` to be set and is mutually exclusive with
+        ``free_rocket``.
+        """
+        if self.free_rocket is not None:
+            raise ValueError(
+                "add_surface is mutually exclusive with free_rocket; "
+                "a free_rocket was already provided for this deployable."
+            )
+        if self.radius is None:
+            raise ValueError("add_surface requires radius to be set on the deployable.")
+        self.surfaces.append((surface, position))

rocketpy/rocket/multistage/geometry.py

@@ -0,0 +1,43 @@
+"""Axial geometry of a rocket, as the stacking of stages needs it."""
+
+from rocketpy.rocket.aero_surface.fins.fins import Fins
+from rocketpy.rocket.aero_surface.nose_cone import NoseCone
+from rocketpy.rocket.aero_surface.tail import Tail
+
+
+def axial_extent(rocket):
+    """Axial extent (bottom, top) spanned by a rocket's aerodynamic
+    surfaces, in the rocket's own coordinate system.
+
+    NoseCone, Tail and Fins each occupy a span (their own ``length`` or
+    ``root_chord``, starting from their own reference position); every
+    other surface type (GenericSurface, RailButtons, individual Fin,
+    TubeFins, ...) is treated as a single point at its own position. This
+    is an approximation for those types, not an exact geometric fit.
+
+    Parameters
+    ----------
+    rocket : Rocket
+        Must have at least one aerodynamic surface.
+
+    Returns
+    -------
+    tuple of float
+        (bottom, top) - the lowest and highest axial coordinates spanned.
+    """
+    if not rocket.aerodynamic_surfaces:
+        raise ValueError(
+            "Rocket must have at least one aerodynamic surface to compute "
+            "its axial extent."
+        )
+    bounds = []
+    for surface, position, *_ in rocket.aerodynamic_surfaces:
+        z = position.z
+        bounds.append(z)
+        if isinstance(surface, NoseCone):
+            bounds.append(z - rocket._csys * surface.length)
+        elif isinstance(surface, Tail):
+            bounds.append(z - rocket._csys * surface.length)
+        elif isinstance(surface, Fins):
+            bounds.append(z - rocket._csys * surface.root_chord)
+    return min(bounds), max(bounds)