Update Rocket.draw() to visualize ClusterMotor configurations

1.  **Import `ClusterMotor`:**
    * The file now imports `ClusterMotor` (Line 5) to check the motor type.

2.  **Refactored `_draw_motor` Method (Line 234):**
    * This method is completely refactored. It now acts as a dispatcher,
      calling the new `_generate_motor_patches` helper function.
    * It checks `isinstance(self.rocket.motor, ClusterMotor)`.
    * If it's a cluster, it adds all patches generated by the helper.
    * If it's a simple motor, it uses the original logic (drawing the
      nozzle and chamber centered at y=0).
    * It also correctly calls `_draw_nozzle_tube` to connect the
      airframe to the start of the cluster or single motor.

3.  **New `_generate_motor_patches` Method (Line 259):**
    * This is an **entirely new** helper function.
    * It contains the core logic for drawing clusters.
    * It iterates through `cluster.motors` and `cluster.positions`.
    * For each sub-motor, it correctly calculates the 2D plot offset
      (using `sub_pos[0]` for the 'xz' plane or `sub_pos[1]` for 'yz')
      and the longitudinal position (`sub_pos[2]`).
    * It re-uses the plotting logic of the individual sub-motors
      (e.g., `_generate_combustion_chamber`, `_generate_grains`) but
      applies the correct `translate=(pos_z, offset)` transform.
    * This allows multiple motors to be drawn side-by-side.

4.  **Fix `_draw_center_of_mass_and_pressure` (Line 389):**
    * This method is updated to handle the new 3D `Vector` object
      returned by `self.rocket.center_of_mass(0)`.
    * It now accesses the Z-coordinate correctly using
      `cm_z = float(cm_vector.z.real)` instead of assuming the
      return value is a simple float.

Author: ayoubdsp

rocketpy/plots/rocket_plots.py

@@ -2,6 +2,7 @@
 import numpy as np
 
 from rocketpy.motors import EmptyMotor, HybridMotor, LiquidMotor, SolidMotor
+from rocketpy.motors.ClusterMotor import ClusterMotor
 from rocketpy.rocket.aero_surface import Fins, NoseCone, Tail
 from rocketpy.rocket.aero_surface.generic_surface import GenericSurface
 
@@ -201,7 +202,7 @@ def draw(self, vis_args=None, plane="xz", *, filename=None):
 
         drawn_surfaces = self._draw_aerodynamic_surfaces(ax, vis_args, plane)
         last_radius, last_x = self._draw_tubes(ax, drawn_surfaces, vis_args)
-        self._draw_motor(last_radius, last_x, ax, vis_args)
+        self._draw_motor(last_radius, last_x, ax, vis_args, plane)
         self._draw_rail_buttons(ax, vis_args)
         self._draw_center_of_mass_and_pressure(ax)
         self._draw_sensors(ax, self.rocket.sensors, plane)
@@ -410,45 +411,141 @@ def _draw_tubes(self, ax, drawn_surfaces, vis_args):
             )
         return radius, last_x
 
-    def _draw_motor(self, last_radius, last_x, ax, vis_args):
+    def _draw_motor(self, last_radius, last_x, ax, vis_args, plane="xz"):
         """Draws the motor from motor patches"""
-        total_csys = self.rocket._csys * self.rocket.motor._csys
-        nozzle_position = (
-            self.rocket.motor_position + self.rocket.motor.nozzle_position * total_csys
-        )
+        
+        # Obtenir les patches (logique cluster/simple est dans la fonction ci-dessous)
+        motor_patches = self._generate_motor_patches(ax, plane)
 
-        # Get motor patches translated to the correct position
-        motor_patches = self._generate_motor_patches(total_csys, ax)
-
-        # Draw patches
+        # Dessiner les patches
         if not isinstance(self.rocket.motor, EmptyMotor):
-            # Add nozzle last so it is in front of the other patches
-            nozzle = self.rocket.motor.plots._generate_nozzle(
-                translate=(nozzle_position, 0), csys=self.rocket._csys
-            )
-            motor_patches += [nozzle]
-
-            outline = self.rocket.motor.plots._generate_motor_region(
-                list_of_patches=motor_patches
-            )
-            # add outline first so it is behind the other patches
-            ax.add_patch(outline)
-            for patch in motor_patches:
-                ax.add_patch(patch)
+            
+            if isinstance(self.rocket.motor, ClusterMotor):
+                # Logique pour Cluster
+                for patch in motor_patches:
+                    ax.add_patch(patch)
+                
+                # Raccorder le tube au début (point z min) du cluster
+                if self.rocket.motor.positions:
+                    # Trouve le z le plus proche de la coiffe
+                    cluster_z_positions = [pos[2] for pos in self.rocket.motor.positions]
+                    z_connector = min(cluster_z_positions) if self.rocket._csys == 1 else max(cluster_z_positions)
+                    self._draw_nozzle_tube(last_radius, last_x, z_connector, ax, vis_args)
+            
+            else:
+                # Logique originale pour Moteur Simple
+                total_csys = self.rocket._csys * self.rocket.motor._csys
+                nozzle_position = (
+                    self.rocket.motor_position + self.rocket.motor.nozzle_position * total_csys
+                )
+                
+                # Ajouter la tuyère (logique originale)
+                nozzle = self.rocket.motor.plots._generate_nozzle(
+                    translate=(nozzle_position, 0), csys=self.rocket._csys
+                )
+                motor_patches += [nozzle]
 
-        self._draw_nozzle_tube(last_radius, last_x, nozzle_position, ax, vis_args)
+                outline = self.rocket.motor.plots._generate_motor_region(
+                    list_of_patches=motor_patches
+                )
+                # Ajouter l'outline en premier
+                ax.add_patch(outline)
+                for patch in motor_patches:
+                    ax.add_patch(patch)
 
-    def _generate_motor_patches(self, total_csys, ax):  # pylint: disable=unused-argument
+                self._draw_nozzle_tube(last_radius, last_x, nozzle_position, ax, vis_args)
+    def _generate_motor_patches(self, ax, plane="xz"):
         """Generates motor patches for drawing"""
         motor_patches = []
+        total_csys = self.rocket._csys * self.rocket.motor._csys
 
-        if isinstance(self.rocket.motor, SolidMotor):
+        # ################################################
+        # ## LOGIQUE D'AFFICHAGE POUR CLUSTER DE MOTEURS ##
+        # ################################################
+        if isinstance(self.rocket.motor, ClusterMotor):
+            cluster = self.rocket.motor
+            all_sub_patches = [] # Pour l'outline global
+            
+            for sub_motor, sub_pos in zip(cluster.motors, cluster.positions):
+                # sub_pos est [x, y, z]
+                motor_longitudinal_pos = sub_pos[2] # Position Z du moteur
+                
+                # Déterminer le décalage (offset)
+                offset = 0
+                if plane == "xz":
+                    offset = sub_pos[0]  # Coordonnée X
+                elif plane == "yz":
+                    offset = sub_pos[1]  # Coordonnée Y
+
+                # `sub_total_csys` convertit un déplacement relatif au moteur
+                # en un déplacement relatif à la fusée.
+                sub_total_csys = self.rocket._csys * sub_motor._csys
+                
+                # On réutilise la logique de SolidMotor, mais avec un offset
+                if isinstance(sub_motor, SolidMotor):
+                    current_motor_patches = [] # Patches pour CE moteur
+                    
+                    # Position Z du centre de masse des grains DANS LE REPERE FUSÉE
+                    grains_cm_pos = (
+                        motor_longitudinal_pos
+                        + sub_motor.grains_center_of_mass_position * sub_total_csys
+                    )
+                    ax.scatter(
+                        grains_cm_pos.real, # Utiliser .real pour éviter ComplexWarning
+                        offset, 
+                        color="brown",
+                        label=f"Grains CM ({sub_pos[0]:.2f}, {sub_pos[1]:.2f})", 
+                        s=8,
+                        zorder=10,
+                    )
+
+                    # `translate` prend (pos_z, pos_y_offset)
+                    # Ces fonctions utilisent le _csys interne du sub_motor (qui est 1)
+                    chamber = sub_motor.plots._generate_combustion_chamber(
+                        translate=(grains_cm_pos.real, offset), label=None
+                    )
+                    grains = sub_motor.plots._generate_grains(
+                        translate=(grains_cm_pos.real, offset)
+                    )
+                    
+                    # Position Z de la tuyère DANS LE REPERE FUSÉE
+                    nozzle_position = (
+                        motor_longitudinal_pos
+                        + sub_motor.nozzle_position * sub_total_csys
+                    )
+                    
+                    # ***** CORRECTION ICI *****
+                    # On ne passe PAS de 'csys' !
+                    # On laisse la fonction _generate_nozzle utiliser son
+                    # propre _csys interne (qui est 1, tail_to_nose),
+                    # car 'nozzle_position' est déjà la coordonnée absolue.
+                    nozzle = sub_motor.plots._generate_nozzle(
+                        translate=(nozzle_position.real, offset)
+                    )
+                    # **************************
+                    
+                    current_motor_patches += [chamber, *grains, nozzle]
+                    all_sub_patches.extend(current_motor_patches)
+
+            # Créer un outline global pour tous les moteurs
+            if all_sub_patches:
+                # Utiliser .plots du premier moteur pour la méthode
+                outline = self.rocket.motor.motors[0].plots._generate_motor_region(
+                    list_of_patches=all_sub_patches
+                )
+                motor_patches.append(outline) # Mettre l'outline en premier
+                motor_patches.extend(all_sub_patches)
+
+        # #####################################
+        # ## LOGIQUE ORIGINALE (MOTEUR SIMPLE) ##
+        # #####################################
+        elif isinstance(self.rocket.motor, SolidMotor):
             grains_cm_position = (
                 self.rocket.motor_position
                 + self.rocket.motor.grains_center_of_mass_position * total_csys
             )
             ax.scatter(
-                grains_cm_position,
+                grains_cm_position.real, # .real
                 0,
                 color="brown",
                 label="Grains Center of Mass",
@@ -457,10 +554,10 @@ def _generate_motor_patches(self, total_csys, ax):  # pylint: disable=unused-arg
             )
 
             chamber = self.rocket.motor.plots._generate_combustion_chamber(
-                translate=(grains_cm_position, 0), label=None
+                translate=(grains_cm_position.real, 0), label=None # .real
             )
             grains = self.rocket.motor.plots._generate_grains(
-                translate=(grains_cm_position, 0)
+                translate=(grains_cm_position.real, 0) # .real
             )
 
             motor_patches += [chamber, *grains]
@@ -471,7 +568,7 @@ def _generate_motor_patches(self, total_csys, ax):  # pylint: disable=unused-arg
                 + self.rocket.motor.grains_center_of_mass_position * total_csys
             )
             ax.scatter(
-                grains_cm_position,
+                grains_cm_position.real, # .real
                 0,
                 color="brown",
                 label="Grains Center of Mass",
@@ -483,16 +580,16 @@ def _generate_motor_patches(self, total_csys, ax):  # pylint: disable=unused-arg
                 translate=(self.rocket.motor_position, 0), csys=total_csys
             )
             chamber = self.rocket.motor.plots._generate_combustion_chamber(
-                translate=(grains_cm_position, 0), label=None
+                translate=(grains_cm_position.real, 0), label=None # .real
             )
             grains = self.rocket.motor.plots._generate_grains(
-                translate=(grains_cm_position, 0)
+                translate=(grains_cm_position.real, 0) # .real
             )
             motor_patches += [chamber, *grains]
             for tank, center in tanks_and_centers:
                 ax.scatter(
-                    center[0],
-                    center[1],
+                    center[0].real, # .real
+                    center[1].real, # .real
                     color="black",
                     alpha=0.2,
                     s=5,
@@ -506,8 +603,8 @@ def _generate_motor_patches(self, total_csys, ax):  # pylint: disable=unused-arg
             )
             for tank, center in tanks_and_centers:
                 ax.scatter(
-                    center[0],
-                    center[1],
+                    center[0].real, # .real
+                    center[1].real, # .real
                     color="black",
                     alpha=0.2,
                     s=4,
@@ -576,12 +673,19 @@ def _draw_rail_buttons(self, ax, vis_args):
     def _draw_center_of_mass_and_pressure(self, ax):
         """Draws the center of mass and center of pressure of the rocket."""
         # Draw center of mass and center of pressure
-        cm = self.rocket.center_of_mass(0)
-        ax.scatter(cm, 0, color="#1565c0", label="Center of Mass", s=10)
-
-        cp = self.rocket.cp_position(0)
+        
+        # MODIFICATION 1: Récupérer le vecteur CM et extraire .z
+        cm_vector = self.rocket.center_of_mass(0)
+        # On prend la partie réelle pour éviter les warnings
+        cm_z = float(cm_vector.z.real) 
+        
+        # On suppose que le CM est sur l'axe pour le dessin 2D
+        ax.scatter(cm_z, 0, color="#1565c0", label="Center of Mass", s=10)
+
+        # MODIFICATION 2: Prendre la partie réelle du CP aussi
+        cp_z = float(self.rocket.cp_position(0).real)
         ax.scatter(
-            cp, 0, label="Static Center of Pressure", color="red", s=10, zorder=10
+            cp_z, 0, label="Static Center of Pressure", color="red", s=10, zorder=10
         )
 
     def _draw_sensors(self, ax, sensors, plane):