add wingtip ribbons

peterdsharpe · Dec 24, 2023 · 52418e2 · 52418e2
1 parent ecbf32a
commit 52418e2
Showing 1 changed file with 99 additions and 50 deletions.
diff --git a/aerosandbox/dynamics/point_mass/common_point_mass.py b/aerosandbox/dynamics/point_mass/common_point_mass.py
@@ -407,7 +407,7 @@ def op_point(self):
         )
 
     def draw(self,
-             vehicle_model: Airplane = None,
+             vehicle_model: Union[Airplane, "PolyData"] = None,
              backend: str = "pyvista",
              plotter=None,
              draw_axes: bool = True,
@@ -416,10 +416,13 @@ def draw(self,
              scale_vehicle_model: Union[float, None] = None,
              n_vehicles_to_draw: int = 10,
              cg_axes: str = "geometry",
+             draw_trajectory_line: bool = True,
+             trajectory_line_color=None,
              draw_altitude_drape: bool = True,
              draw_ground_plane: bool = True,
+             draw_wingtip_ribbon: bool = True,
+             set_sky_background: bool = True,
              vehicle_color=None,
-             trajectory_line_color=None,
              show: bool = True,
              ):
         if backend == "pyvista":
@@ -432,7 +435,7 @@ def draw(self,
             elif isinstance(vehicle_model, pv.PolyData):
                 pass
             elif isinstance(vehicle_model, Airplane):
-                vehicle_model = vehicle_model.draw(
+                vehicle_model: pv.PolyData = vehicle_model.draw(
                     backend="pyvista",
                     show=False
                 )
@@ -461,18 +464,31 @@ def draw(self,
                 [z_e.min(), z_e.max()],
             ])
 
-            if scale_vehicle_model is None:  # Compute an auto-scaling factor
-                trajectory_size = np.max(np.diff(trajectory_bounds, axis=1))
+            vehicle_bounds = np.array(vehicle_model.bounds).reshape((3, 2))
 
-                vehicle_bounds = np.array(vehicle_model.bounds).reshape((3, 2))
-                vehicle_size = np.max(np.diff(vehicle_bounds, axis=1))
+            # trajectory_size = np.max(np.diff(trajectory_bounds, axis=1))  # Max dimension
+            # vehicle_size = np.max(np.diff(vehicle_bounds, axis=1))  # Max dimension
 
-                scale_vehicle_model = 0.8 * trajectory_size / vehicle_size / n_vehicles_to_draw
+            if scale_vehicle_model is None:  # Compute an auto-scaling factor
+                if len(self) == 1:
+                    scale_vehicle_model = 1
+                else:
+                    path_length = np.sum(
+                        (np.diff(x_e) ** 2 + np.diff(y_e) ** 2 + np.diff(z_e) ** 2) ** 0.5
+                    )
+                    vehicle_length = np.diff(vehicle_bounds[0, :])
+                    scale_vehicle_model = 0.5 * path_length / vehicle_length / n_vehicles_to_draw
 
             ### Initialize the plotter
             if plotter is None:
                 plotter = pv.Plotter()
 
+            if set_sky_background:
+                plotter.set_background(
+                    color="#FFFFFF",
+                    top="#A5B8D7",
+                )
+
             # Set the window title
             title = "ASB Dynamics"
             addenda = []
@@ -597,53 +613,86 @@ def draw(self,
                 z_e,
             ], axis=1)
 
-            polyline = pv.Spline(path)
-            plotter.add_mesh(
-                polyline,
-                color=(
-                    p.adjust_lightness(p.palettes["categorical"][0], 1.3)
-                    if trajectory_line_color is None
-                    else trajectory_line_color
-                ),
-                line_width=3,
-            )
+            if len(self) > 1:
+                if draw_trajectory_line:
+                    polyline = pv.Spline(path)
+                    plotter.add_mesh(
+                        polyline,
+                        color=(
+                            p.adjust_lightness(p.palettes["categorical"][0], 1.3)
+                            if trajectory_line_color is None
+                            else trajectory_line_color
+                        ),
+                        line_width=3,
+                    )
 
-            if draw_altitude_drape:
-                ### Drape
+                if draw_wingtip_ribbon:
 
-                points = np.concatenate([
-                    path,
-                    path * np.array([[1, 1, 0]])
-                ], axis=0)
+                    left_wingtip_points = np.array(self.convert_axes(
+                        0, scale_vehicle_model * vehicle_bounds[1, 0], 0,
+                        from_axes="body",
+                        to_axes="earth"
+                    )).T + path
+                    plotter.add_mesh(
+                        pv.Spline(left_wingtip_points),
+                        color="pink",
+                    )
+                    right_wingtip_points = np.array(self.convert_axes(
+                        0, scale_vehicle_model * vehicle_bounds[1, 1], 0,
+                        from_axes="body",
+                        to_axes="earth"
+                    )).T + path
+                    plotter.add_mesh(
+                        pv.Spline(right_wingtip_points),
+                        color="pink",
+                    )
 
-                grid = pv.StructuredGrid()
-                grid.points = points
-                grid.dimensions = len(path), 2, 1
+                    grid = pv.StructuredGrid()
+                    grid.points = np.concatenate([
+                        left_wingtip_points,
+                        right_wingtip_points,
+                    ], axis=0)
+                    grid.dimensions = len(left_wingtip_points), 2, 1
+
+                    plotter.add_mesh(
+                        grid,
+                        color="pink",
+                        opacity=0.5,
+                    )
 
-                plotter.add_mesh(
-                    grid,
-                    color="black",
-                    opacity=0.5,
-                )
+                if draw_altitude_drape:
+                    ### Drape
+                    grid = pv.StructuredGrid()
+                    grid.points = np.concatenate([
+                        path,
+                        path * np.array([[1, 1, 0]])
+                    ], axis=0)
+                    grid.dimensions = len(path), 2, 1
+
+                    plotter.add_mesh(
+                        grid,
+                        color="black",
+                        opacity=0.5,
+                    )
 
-            if draw_ground_plane:
-                ### Plane
-                grid = pv.StructuredGrid()
-                xlim = (x_e.min(), x_e.max())
-                ylim = (y_e.min(), y_e.max())
-
-                grid.points = np.array([
-                    [xlim[0], ylim[0], 0],
-                    [xlim[1], ylim[0], 0],
-                    [xlim[0], ylim[1], 0],
-                    [xlim[1], ylim[1], 0]
-                ])
-                grid.dimensions = 2, 2, 1
-                plotter.add_mesh(
-                    grid,
-                    color="darkkhaki",
-                    opacity=0.5
-                )
+                if draw_ground_plane:
+                    ### Plane
+                    grid = pv.StructuredGrid()
+                    xlim = (x_e.min(), x_e.max())
+                    ylim = (y_e.min(), y_e.max())
+
+                    grid.points = np.array([
+                        [xlim[0], ylim[0], 0],
+                        [xlim[1], ylim[0], 0],
+                        [xlim[0], ylim[1], 0],
+                        [xlim[1], ylim[1], 0]
+                    ])
+                    grid.dimensions = 2, 2, 1
+                    plotter.add_mesh(
+                        grid,
+                        color="darkkhaki",
+                        opacity=0.5
+                    )
 
             ### Finalize the plotter
             plotter.camera.up = (0, 0, -1)