diff --git a/requirements.txt b/requirements.txt index 7274c7a..80bc759 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,7 +9,7 @@ pymongo>=4.15 jsonpickle gunicorn uvicorn -git+https://github.com/RocketPy-Team/RocketPy.git@develop +rocketpy uptrace opentelemetry.instrumentation.fastapi opentelemetry.instrumentation.requests diff --git a/src/controllers/motor.py b/src/controllers/motor.py index 3483c1b..1341e91 100644 --- a/src/controllers/motor.py +++ b/src/controllers/motor.py @@ -2,7 +2,7 @@ ControllerBase, controller_exception_handler, ) -from src.views.motor import MotorSimulation +from src.views.motor import MotorSimulation, MotorDrawingGeometryView from src.models.motor import MotorModel from src.services.motor import MotorService @@ -57,3 +57,27 @@ async def get_motor_simulation(self, motor_id: str) -> MotorSimulation: motor = await self.get_motor_by_id(motor_id) motor_service = MotorService.from_motor_model(motor.motor) return motor_service.get_motor_simulation() + + @controller_exception_handler + async def get_motor_drawing_geometry( + self, motor_id: str + ) -> MotorDrawingGeometryView: + """ + Build the motor-only drawing-geometry payload for a persisted motor. + + Renders the motor at its own coordinate origin (motor_position=0, + parent_csys=1) so the playground can show a motor in isolation. + + Args: + motor_id: str + + Returns: + views.MotorDrawingGeometryView + + Raises: + HTTP 404 Not Found: If the motor does not exist in the database. + HTTP 422: If the motor has no drawable geometry. + """ + motor = await self.get_motor_by_id(motor_id) + motor_service = MotorService.from_motor_model(motor.motor) + return motor_service.get_drawing_geometry() diff --git a/src/controllers/rocket.py b/src/controllers/rocket.py index ce1c36c..9c71db5 100644 --- a/src/controllers/rocket.py +++ b/src/controllers/rocket.py @@ -4,7 +4,11 @@ ControllerBase, controller_exception_handler, ) -from src.views.rocket import RocketSimulation, RocketCreated +from src.views.rocket import ( + RocketSimulation, + RocketCreated, + RocketDrawingGeometry, +) from src.models.motor import MotorModel from src.models.rocket import ( RocketModel, @@ -75,6 +79,27 @@ async def get_rocketpy_rocket_binary(self, rocket_id: str) -> bytes: rocket_service = RocketService.from_rocket_model(rocket.rocket) return rocket_service.get_rocket_binary() + @controller_exception_handler + async def get_rocket_drawing_geometry( + self, rocket_id: str + ) -> RocketDrawingGeometry: + """ + Build the drawing geometry payload for a persisted rocket. + + Args: + rocket_id: str + + Returns: + views.RocketDrawingGeometry + + Raises: + HTTP 404 Not Found: If the rocket does not exist in the database. + HTTP 422: If the rocket has no aerodynamic surfaces to draw. + """ + rocket = await self.get_rocket_by_id(rocket_id) + rocket_service = RocketService.from_rocket_model(rocket.rocket) + return rocket_service.get_drawing_geometry() + @controller_exception_handler async def get_rocket_simulation( self, diff --git a/src/models/motor.py b/src/models/motor.py index 93eb0b4..9e0c049 100644 --- a/src/models/motor.py +++ b/src/models/motor.py @@ -82,6 +82,22 @@ def validate_motor_kind(self): ) return self + @model_validator(mode='after') + def validate_dry_inertia_for_kind(self): + # RocketPy's SolidMotor/LiquidMotor/HybridMotor require dry_inertia with no default. + # Only GenericMotor accepts (0, 0, 0). Surface a clear error at the API boundary + # instead of letting RocketPy crash deep in construction. + if self.motor_kind != MotorKinds.GENERIC and self.dry_inertia == ( + 0, + 0, + 0, + ): + raise ValueError( + f"dry_inertia is required for {self.motor_kind} motors " + f"and must be explicitly provided (cannot be (0, 0, 0))." + ) + return self + @staticmethod def UPDATED(): return diff --git a/src/models/sub/tanks.py b/src/models/sub/tanks.py index 0873264..a2daaf0 100644 --- a/src/models/sub/tanks.py +++ b/src/models/sub/tanks.py @@ -22,7 +22,8 @@ class MotorTank(BaseModel): liquid: TankFluids flux_time: Tuple[float, float] position: float - discretize: int + # discretize is optional in RocketPy's Tank classes (defaults to 100). + discretize: int = 100 # Level based tank parameters liquid_height: Optional[float] = None diff --git a/src/routes/motor.py b/src/routes/motor.py index 0673708..e85f713 100644 --- a/src/routes/motor.py +++ b/src/routes/motor.py @@ -9,6 +9,7 @@ MotorSimulation, MotorCreated, MotorRetrieved, + MotorDrawingGeometryView, ) from src.models.motor import MotorModel from src.dependencies import MotorControllerDep @@ -138,3 +139,24 @@ async def get_motor_simulation( """ with tracer.start_as_current_span("get_motor_simulation"): return await controller.get_motor_simulation(motor_id) + + +@router.get("/{motor_id}/drawing-geometry") +async def get_motor_drawing_geometry( + motor_id: str, + controller: MotorControllerDep, +) -> MotorDrawingGeometryView: + """ + Returns motor-only drawing geometry so a frontend can render the + motor in isolation. The payload mirrors what the motor portion of + `rocketpy.Rocket.draw()` would produce, but at the motor's own + coordinate origin rather than embedded inside a rocket. + + Use `GET /rockets/{rocket_id}/drawing-geometry` instead when the + motor should be shown inside a complete rocket. + + ## Args + ``` motor_id: Motor ID ``` + """ + with tracer.start_as_current_span("get_motor_drawing_geometry"): + return await controller.get_motor_drawing_geometry(motor_id) diff --git a/src/routes/rocket.py b/src/routes/rocket.py index e65c846..c3cae33 100644 --- a/src/routes/rocket.py +++ b/src/routes/rocket.py @@ -9,6 +9,7 @@ RocketSimulation, RocketCreated, RocketRetrieved, + RocketDrawingGeometry, ) from src.models.rocket import ( RocketModel, @@ -181,3 +182,23 @@ async def simulate_rocket( """ with tracer.start_as_current_span("get_rocket_simulation"): return await controller.get_rocket_simulation(rocket_id) + + +@router.get("/{rocket_id}/drawing-geometry") +async def get_rocket_drawing_geometry( + rocket_id: str, + controller: RocketControllerDep, +) -> RocketDrawingGeometry: + """ + Returns structured drawing geometry for the rocket so that a frontend + can redraw exactly what rocketpy.Rocket.draw() would render. + + Response contains shape coordinate arrays for each aerodynamic surface, + tube segments, motor polygons (nozzle, chamber, grains, tanks, outline), + rail-button positions, CG/CP at t=0, sensors, and overall drawing bounds. + + ## Args + ``` rocket_id: Rocket ID ``` + """ + with tracer.start_as_current_span("get_rocket_drawing_geometry"): + return await controller.get_rocket_drawing_geometry(rocket_id) diff --git a/src/services/motor.py b/src/services/motor.py index 7275920..41493b8 100644 --- a/src/services/motor.py +++ b/src/services/motor.py @@ -1,12 +1,15 @@ from typing import Self import dill +import numpy as np -from rocketpy.motors.motor import GenericMotor, Motor as RocketPyMotor +from rocketpy.motors.motor import Motor as RocketPyMotor from rocketpy.motors.solid_motor import SolidMotor from rocketpy.motors.liquid_motor import LiquidMotor from rocketpy.motors.hybrid_motor import HybridMotor -from rocketpy import ( +from rocketpy.motors import ( + EmptyMotor, + GenericMotor, LevelBasedTank, MassBasedTank, MassFlowRateBasedTank, @@ -16,12 +19,72 @@ from fastapi import HTTPException, status +from src import logger from src.models.sub.tanks import TankKinds from src.models.motor import MotorKinds, MotorModel -from src.views.motor import MotorSimulation +from src.views.motor import MotorSimulation, MotorDrawingGeometryView +from src.views.drawing import ( + DrawingBounds, + MotorDrawingGeometry, + MotorPatch, +) from src.utils import collect_attributes +# --------------------------------------------------------------------------- +# Drawing-geometry helpers (module-level, shared by the motor and rocket +# drawing paths). Kept private with the `_` prefix; consumers should go +# through `MotorService.get_drawing_geometry` for the motor-only response +# or through `RocketService.get_drawing_geometry` for the composed rocket +# view (which delegates motor assembly here). +# --------------------------------------------------------------------------- +def _polygon_xy(patch) -> dict: + """Extract (x, y) coordinate lists from a matplotlib Polygon patch. + + Rocketpy's `_MotorPlots` generator helpers return matplotlib `Polygon` + objects; we only ever read `patch.xy` (an Nx2 numpy array) as a data + carrier, never for rendering. + """ + xy = np.asarray(patch.xy) + return {"x": xy[:, 0].tolist(), "y": xy[:, 1].tolist()} + + +def _rebuild_polygon(x: list[float], y: list[float]): + """Rebuild a matplotlib `Polygon` from coordinate lists. + + Used only so `_MotorPlots._generate_motor_region` can read `patch.xy` + bounds when we assemble the motor outline. + """ + from matplotlib.patches import ( + Polygon, + ) # local import keeps service cold-start lean + + return Polygon(np.column_stack([np.asarray(x), np.asarray(y)])) + + +def _build_generic_chamber_patch( + center_x: float, chamber_height: float, chamber_radius: float +): + """Build a rectangular combustion-chamber polygon for a GenericMotor. + + Mirrors the vertex order of + `rocketpy.plots.motor_plots._generate_combustion_chamber` so the patch + can flow through `_generate_motor_region` for outline computation + identically to a SolidMotor chamber. + """ + from matplotlib.patches import ( + Polygon, + ) # local import keeps service cold-start lean + + half_len = chamber_height / 2.0 + x = np.array([-half_len, half_len, half_len, -half_len]) + y = np.array( + [chamber_radius, chamber_radius, -chamber_radius, -chamber_radius] + ) + x = x + center_x + return Polygon(np.column_stack([x, y])) + + class MotorService: _motor: RocketPyMotor @@ -184,3 +247,193 @@ def get_motor_binary(self) -> bytes: bytes """ return dill.dumps(self.motor) + + # -------------------------------------------------------------------- + # Drawing geometry + # -------------------------------------------------------------------- + def get_drawing_geometry( + self, + motor_position: float = 0.0, + parent_csys: int = 1, + ) -> MotorDrawingGeometryView: + """Build a motor-only drawing-geometry response. + + Defaults render the motor at its own coordinate origin + (`motor_position=0`, `parent_csys=1`) so the payload can be used + standalone in the playground's "motor-only" view. Callers that + embed the motor into a rocket (see ``RocketService``) pass the + rocket-level position + csys to align the motor inside the rocket + frame. + + Raises: + HTTPException 422 if the motor has no drawable patches. + """ + motor_geometry, _nozzle_position = self.build_drawing_geometry( + motor_position=motor_position, parent_csys=parent_csys + ) + if motor_geometry is None: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail="Motor has no drawable geometry.", + ) + # Fall back to the nozzle radius when the motor has no drawable + # patches (e.g. EmptyMotor) so the bounds aren't a zero-height + # line. Every real rocketpy motor class exposes nozzle_radius. + fallback_radius = float(getattr(self._motor, "nozzle_radius", 0.0) or 0.0) + return MotorDrawingGeometryView( + motor=motor_geometry, + coordinate_system_orientation=str( + self._motor.coordinate_system_orientation + ), + bounds=_compute_motor_bounds(motor_geometry, fallback_radius), + ) + + def build_drawing_geometry( + self, + motor_position: float, + parent_csys: int, + ) -> tuple[MotorDrawingGeometry | None, float]: + """Construct motor patches + the absolute nozzle x-position. + + Returned as a tuple so the rocket service can use the nozzle + position when extending body tubes to meet the motor. Standalone + motor rendering can discard the second element. + """ + motor = self._motor + total_csys = parent_csys * motor._csys + nozzle_position = motor_position + motor.nozzle_position * total_csys + + if isinstance(motor, EmptyMotor): + return ( + MotorDrawingGeometry( + type="empty", + position=float(motor_position), + nozzle_position=float(nozzle_position), + patches=[], + ), + float(nozzle_position), + ) + + patches: list[MotorPatch] = [] + grains_cm_position: float | None = None + motor_type = "generic" + + if isinstance(motor, SolidMotor): + motor_type = "solid" + grains_cm_position = ( + motor_position + + motor.grains_center_of_mass_position * total_csys + ) + chamber = motor.plots._generate_combustion_chamber( + translate=(grains_cm_position, 0), label=None + ) + patches.append(MotorPatch(role="chamber", **_polygon_xy(chamber))) + for grain in motor.plots._generate_grains( + translate=(grains_cm_position, 0) + ): + patches.append(MotorPatch(role="grain", **_polygon_xy(grain))) + elif isinstance(motor, HybridMotor): + motor_type = "hybrid" + grains_cm_position = ( + motor_position + + motor.grains_center_of_mass_position * total_csys + ) + chamber = motor.plots._generate_combustion_chamber( + translate=(grains_cm_position, 0), label=None + ) + patches.append(MotorPatch(role="chamber", **_polygon_xy(chamber))) + for grain in motor.plots._generate_grains( + translate=(grains_cm_position, 0) + ): + patches.append(MotorPatch(role="grain", **_polygon_xy(grain))) + for tank, _center in motor.plots._generate_positioned_tanks( + translate=(motor_position, 0), csys=total_csys + ): + patches.append(MotorPatch(role="tank", **_polygon_xy(tank))) + elif isinstance(motor, LiquidMotor): + motor_type = "liquid" + for tank, _center in motor.plots._generate_positioned_tanks( + translate=(motor_position, 0), csys=total_csys + ): + patches.append(MotorPatch(role="tank", **_polygon_xy(tank))) + elif isinstance(motor, GenericMotor): + # RocketPy's Rocket.draw() does not render a chamber for + # GenericMotor — `_generate_combustion_chamber` depends on + # grain fields GenericMotor lacks. We build an equivalent + # rectangular chamber from the GenericMotor fields so users + # see their chamber geometry in the playground. + motor_type = "generic" + chamber_center_x = ( + motor_position + motor.chamber_position * total_csys + ) + chamber_patch = _build_generic_chamber_patch( + center_x=chamber_center_x, + chamber_height=motor.chamber_height, + chamber_radius=motor.chamber_radius, + ) + patches.append( + MotorPatch(role="chamber", **_polygon_xy(chamber_patch)) + ) + + # Nozzle is always appended after the body so the motor-region + # outline encompasses it, matching rocketpy. + nozzle_patch = motor.plots._generate_nozzle( + translate=(nozzle_position, 0), csys=parent_csys + ) + patches.append(MotorPatch(role="nozzle", **_polygon_xy(nozzle_patch))) + + # Motor-region outline. `_generate_motor_region` reads patch.xy + # arrays, so we rebuild matplotlib Polygons once from our + # coordinate copies. Any failure here is logged and dropped; the + # outline is advisory, not load-bearing. + try: + mpl_patches = [_rebuild_polygon(p.x, p.y) for p in patches] + outline_patch = motor.plots._generate_motor_region( + list_of_patches=mpl_patches + ) + patches.insert( + 0, MotorPatch(role="outline", **_polygon_xy(outline_patch)) + ) + except Exception as exc: # pragma: no cover - defensive + logger.warning("Failed to generate motor outline patch: %s", exc) + + return ( + MotorDrawingGeometry( + type=motor_type, + position=float(motor_position), + nozzle_position=float(nozzle_position), + grains_center_of_mass_position=( + float(grains_cm_position) + if grains_cm_position is not None + else None + ), + patches=patches, + ), + float(nozzle_position), + ) + + +def _compute_motor_bounds( + motor: MotorDrawingGeometry, radius: float +) -> DrawingBounds: + """Compute a tight bounding box for a motor-only drawing payload. + + Mirrors the rocket-side bounds helper but limited to motor patches — + callers who want rocket-wide bounds use the rocket service's own + computation. + """ + xs: list[float] = [] + ys: list[float] = [] + for patch in motor.patches: + xs += patch.x + ys += patch.y + if not xs: + xs = [float(motor.position)] + if not ys: + ys = [-float(radius), float(radius)] + return DrawingBounds( + x_min=float(min(xs)), + x_max=float(max(xs)), + y_min=float(min(ys)), + y_max=float(max(ys)), + ) diff --git a/src/services/rocket.py b/src/services/rocket.py index a22a2d4..2ed66db 100644 --- a/src/services/rocket.py +++ b/src/services/rocket.py @@ -1,6 +1,7 @@ from typing import Self, List import dill +import numpy as np from rocketpy.rocket.rocket import Rocket as RocketPyRocket from rocketpy.rocket.parachute import Parachute as RocketPyParachute @@ -11,6 +12,7 @@ Fins as RocketPyFins, Tail as RocketPyTail, ) +from rocketpy.rocket.aero_surface.generic_surface import GenericSurface from fastapi import HTTPException, status @@ -18,7 +20,20 @@ from src.models.rocket import RocketModel, Parachute from src.models.sub.aerosurfaces import NoseCone, Tail, Fins from src.services.motor import MotorService -from src.views.rocket import RocketSimulation +from src.views.rocket import ( + RocketSimulation, + RocketDrawingGeometry, + NoseConeGeometry, + TailGeometry, + FinsGeometry, + FinOutline, + TubeGeometry, + MotorDrawingGeometry, + MotorPatch, + RailButtonsGeometry, + SensorGeometry, + DrawingBounds, +) from src.views.motor import MotorSimulation from src.utils import collect_attributes @@ -125,6 +140,322 @@ def get_rocket_binary(self) -> bytes: """ return dill.dumps(self.rocket) + def get_drawing_geometry(self) -> RocketDrawingGeometry: + """ + Build the drawing-geometry payload that mirrors rocketpy.Rocket.draw(). + + Coordinates are emitted in the draw frame used by _RocketPlots: the + axial axis is x (applying rocket._csys), the radial axis is y with the + caller expected to mirror (x, y) and (x, -y) for the two halves of + nose/tail/body, and each fin outline is a closed polyline in that + same frame. + + Returns: + RocketDrawingGeometry + + Raises: + HTTP 422: if the rocket has no aerodynamic surfaces to draw. + """ + rocket = self._rocket + + if not rocket.aerodynamic_surfaces: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail=( + "Rocket must have at least one aerodynamic surface " + "before a drawing can be produced." + ), + ) + + csys = rocket._csys + rocket.aerodynamic_surfaces.sort_by_position(reverse=csys == 1) + + nose_cones: list[NoseConeGeometry] = [] + tails: list[TailGeometry] = [] + fins: list[FinsGeometry] = [] + # drawn_surfaces mirrors the tuples that _RocketPlots._draw_tubes + # consumes: (surface, reference_x, radius_at_end, last_x). + drawn_surfaces: list[tuple] = [] + + for surface, position in rocket.aerodynamic_surfaces: + position_z = position.z + if isinstance(surface, RocketPyNoseCone): + x_vals = -csys * np.asarray(surface.shape_vec[0]) + position_z + y_vals = np.asarray(surface.shape_vec[1]) + nose_cones.append( + NoseConeGeometry( + name=getattr(surface, "name", None), + kind=getattr(surface, "kind", None), + position=float(position_z), + x=x_vals.tolist(), + y=y_vals.tolist(), + ) + ) + drawn_surfaces.append( + ( + surface, + float(x_vals[-1]), + float(surface.rocket_radius), + float(x_vals[-1]), + ) + ) + elif isinstance(surface, RocketPyTail): + x_vals = -csys * np.asarray(surface.shape_vec[0]) + position_z + y_vals = np.asarray(surface.shape_vec[1]) + tails.append( + TailGeometry( + name=getattr(surface, "name", None), + position=float(position_z), + x=x_vals.tolist(), + y=y_vals.tolist(), + ) + ) + drawn_surfaces.append( + ( + surface, + float(position_z), + float(surface.bottom_radius), + float(x_vals[-1]), + ) + ) + elif isinstance(surface, RocketPyFins): + num_fins = surface.n + x_fin = -csys * np.asarray(surface.shape_vec[0]) + position_z + y_fin = ( + np.asarray(surface.shape_vec[1]) + surface.rocket_radius + ) + outlines: list[FinOutline] = [] + last_x_rotated = float(x_fin[-1]) + for i in range(num_fins): + angle = 2 * np.pi * i / num_fins + rotation_matrix = np.array([[1, 0], [0, np.cos(angle)]]) + rotated = rotation_matrix @ np.vstack((x_fin, y_fin)) + outlines.append( + FinOutline( + x=rotated[0].tolist(), + y=rotated[1].tolist(), + ) + ) + last_x_rotated = float(rotated[0][-1]) + kind = ( + "trapezoidal" + if isinstance(surface, RocketPyTrapezoidalFins) + else ( + "elliptical" + if isinstance(surface, RocketPyEllipticalFins) + else "free_form" + ) + ) + fins.append( + FinsGeometry( + name=getattr(surface, "name", None), + kind=kind, + n=int(num_fins), + cant_angle_deg=float( + getattr(surface, "cant_angle", 0.0) or 0.0 + ), + position=float(position_z), + outlines=outlines, + ) + ) + drawn_surfaces.append( + ( + surface, + float(position_z), + float(surface.rocket_radius), + last_x_rotated, + ) + ) + elif isinstance(surface, GenericSurface): + # Generic surfaces aren't part of the rendered rocket shell; + # they contribute a reference point for tube continuity. + drawn_surfaces.append( + ( + surface, + float(position_z), + float(rocket.radius), + float(position_z), + ) + ) + + tubes = self._build_tubes(drawn_surfaces) + motor_geometry, nozzle_position = self._build_motor_geometry(csys) + tubes += self._build_nozzle_tube(drawn_surfaces, nozzle_position, csys) + rail_buttons = self._build_rail_buttons(csys) + sensors = self._build_sensors() + + try: + center_of_mass = float(rocket.center_of_mass(0)) + except ( + Exception + ): # pragma: no cover - defensive; rocket may not be fully built + center_of_mass = None + try: + cp_position = float(rocket.cp_position(0)) + except Exception: # pragma: no cover + cp_position = None + + bounds = self._compute_bounds( + nose_cones, tails, fins, tubes, motor_geometry, rocket.radius + ) + + return RocketDrawingGeometry( + radius=float(rocket.radius), + csys=int(csys), + coordinate_system_orientation=str( + rocket.coordinate_system_orientation + ), + nose_cones=nose_cones, + tails=tails, + fins=fins, + tubes=tubes, + motor=motor_geometry, + rail_buttons=rail_buttons, + center_of_mass=center_of_mass, + cp_position=cp_position, + sensors=sensors, + bounds=bounds, + ) + + @staticmethod + def _build_tubes(drawn_surfaces: list) -> list[TubeGeometry]: + tubes: list[TubeGeometry] = [] + for i, d_surface in enumerate(drawn_surfaces): + surface, position, radius, last_x = d_surface + if i == len(drawn_surfaces) - 1: + if isinstance(surface, RocketPyTail): + continue + x_start, x_end = position, last_x + else: + next_position = drawn_surfaces[i + 1][1] + x_start, x_end = last_x, next_position + tubes.append( + TubeGeometry( + x_start=float(x_start), + x_end=float(x_end), + radius=float(radius), + ) + ) + return tubes + + def _build_motor_geometry( + self, csys: int + ) -> tuple[MotorDrawingGeometry | None, float]: + # Delegate to the motor service so motor-only and rocket-embedded + # drawings share exactly one code path. The motor service carries + # all the isinstance branches, private rocketpy plot calls, and + # matplotlib Polygon plumbing; here we just bind the rocket-level + # position + csys and forward. + motor_service = MotorService(self._rocket.motor) + return motor_service.build_drawing_geometry( + motor_position=self._rocket.motor_position, + parent_csys=csys, + ) + + def _build_nozzle_tube( + self, + drawn_surfaces: list, + nozzle_position: float, + csys: int, + ) -> list[TubeGeometry]: + if not drawn_surfaces: + return [] + last_surface, _, last_radius, last_x = drawn_surfaces[-1] + if isinstance(last_surface, RocketPyTail): + return [] + if csys == 1 and nozzle_position < last_x: + extra_x = nozzle_position + elif csys == -1 and nozzle_position > last_x: + extra_x = nozzle_position + else: + return [] + return [ + TubeGeometry( + x_start=float(last_x), + x_end=float(extra_x), + radius=float(last_radius), + ) + ] + + def _build_rail_buttons(self, csys: int) -> RailButtonsGeometry | None: + rocket = self._rocket + try: + buttons, pos = rocket.rail_buttons[0] + except IndexError: + return None + lower = float(pos.z) + upper = lower + float(buttons.buttons_distance) * csys + return RailButtonsGeometry( + lower_x=lower, + upper_x=upper, + y=-float(rocket.radius), + angular_position_deg=float( + getattr(buttons, "angular_position", 0.0) or 0.0 + ), + ) + + def _build_sensors(self) -> list[SensorGeometry]: + rocket = self._rocket + sensors: list[SensorGeometry] = [] + for sensor_pos in getattr(rocket, "sensors", []) or []: + sensor = sensor_pos[0] + pos = sensor_pos[1] + normal = getattr(sensor, "normal_vector", None) + normal_tuple = ( + (float(normal.x), float(normal.y), float(normal.z)) + if normal is not None + else (0.0, 0.0, 0.0) + ) + sensors.append( + SensorGeometry( + name=getattr(sensor, "name", None), + position=(float(pos[0]), float(pos[1]), float(pos[2])), + normal=normal_tuple, + ) + ) + return sensors + + @staticmethod + def _compute_bounds( + nose_cones: list[NoseConeGeometry], + tails: list[TailGeometry], + fins: list[FinsGeometry], + tubes: list[TubeGeometry], + motor: MotorDrawingGeometry | None, + radius: float, + ) -> DrawingBounds: + xs: list[float] = [] + ys: list[float] = [] + for nc in nose_cones: + xs += nc.x + ys += nc.y + ys += [-v for v in nc.y] + for tail in tails: + xs += tail.x + ys += tail.y + ys += [-v for v in tail.y] + for finset in fins: + for outline in finset.outlines: + xs += outline.x + ys += outline.y + for tube in tubes: + xs += [tube.x_start, tube.x_end] + ys += [tube.radius, -tube.radius] + if motor is not None: + for patch in motor.patches: + xs += patch.x + ys += patch.y + if not xs: + xs = [0.0] + if not ys: + ys = [-float(radius), float(radius)] + return DrawingBounds( + x_min=float(min(xs)), + x_max=float(max(xs)), + y_min=float(min(ys)), + y_max=float(max(ys)), + ) + @staticmethod def get_rocketpy_nose(nose: NoseCone) -> RocketPyNoseCone: """ @@ -264,3 +595,11 @@ def check_parachute_trigger(trigger) -> bool: if isinstance(trigger, (int, float)): return True return False + + + +# Drawing helpers (_polygon_xy, _rebuild_polygon, _build_generic_chamber_patch) +# moved to src/services/motor.py so the motor-only drawing endpoint and the +# rocket drawing endpoint share a single implementation. This module now +# imports nothing motor-specific for drawing — RocketService._build_motor_geometry +# delegates to MotorService.build_drawing_geometry. diff --git a/src/views/drawing.py b/src/views/drawing.py new file mode 100644 index 0000000..dd74273 --- /dev/null +++ b/src/views/drawing.py @@ -0,0 +1,86 @@ +"""Shared drawing-geometry view types. + +Lives as its own module so both the rocket views (`src/views/rocket.py`) +and the motor views (`src/views/motor.py`) can depend on it without +creating a circular import. Nothing here is rocket- or motor-specific; +these are the raw coordinate-carrier types that rocketpy's drawing +surface produces. +""" + +from typing import Literal, Optional + +from pydantic import BaseModel, ConfigDict + + +class NoseConeGeometry(BaseModel): + model_config = ConfigDict(ser_json_exclude_none=True) + name: Optional[str] = None + kind: Optional[str] = None + position: float + x: list[float] + y: list[float] + + +class TailGeometry(BaseModel): + model_config = ConfigDict(ser_json_exclude_none=True) + name: Optional[str] = None + position: float + x: list[float] + y: list[float] + + +class FinOutline(BaseModel): + x: list[float] + y: list[float] + + +class FinsGeometry(BaseModel): + model_config = ConfigDict(ser_json_exclude_none=True) + name: Optional[str] = None + kind: str + n: int + cant_angle_deg: Optional[float] = None + position: float + outlines: list[FinOutline] + + +class TubeGeometry(BaseModel): + x_start: float + x_end: float + radius: float + + +class MotorPatch(BaseModel): + role: Literal["nozzle", "chamber", "grain", "tank", "outline"] + x: list[float] + y: list[float] + + +class MotorDrawingGeometry(BaseModel): + model_config = ConfigDict(ser_json_exclude_none=True) + type: Literal["solid", "hybrid", "liquid", "empty", "generic"] + position: float + nozzle_position: float + grains_center_of_mass_position: Optional[float] = None + patches: list[MotorPatch] + + +class RailButtonsGeometry(BaseModel): + lower_x: float + upper_x: float + y: float + angular_position_deg: float + + +class SensorGeometry(BaseModel): + model_config = ConfigDict(ser_json_exclude_none=True) + name: Optional[str] = None + position: tuple[float, float, float] + normal: tuple[float, float, float] + + +class DrawingBounds(BaseModel): + x_min: float + x_max: float + y_min: float + y_max: float diff --git a/src/views/motor.py b/src/views/motor.py index 1ca4aea..2b3aaba 100644 --- a/src/views/motor.py +++ b/src/views/motor.py @@ -1,6 +1,7 @@ from typing import Optional, Any from pydantic import ConfigDict from src.views.interface import ApiBaseView +from src.views.drawing import DrawingBounds, MotorDrawingGeometry from src.models.motor import MotorModel @@ -85,3 +86,21 @@ class MotorCreated(ApiBaseView): class MotorRetrieved(ApiBaseView): message: str = "Motor successfully retrieved" motor: MotorView + + +class MotorDrawingGeometryView(ApiBaseView): + """Motor-only drawing-geometry response. + + Returns exactly the subset of the rocket drawing payload that belongs + to the motor — the patches rocketpy's ``Rocket.draw()`` would feed to + matplotlib for this motor, plus a tight bounding box. When the client + needs the motor embedded in a full rocket render it should call + ``GET /rockets/{rocket_id}/drawing-geometry`` instead. + """ + + model_config = ConfigDict(ser_json_exclude_none=True) + + message: str = "Motor drawing geometry retrieved" + motor: MotorDrawingGeometry + coordinate_system_orientation: str + bounds: DrawingBounds diff --git a/src/views/rocket.py b/src/views/rocket.py index 25d08fe..6b07111 100644 --- a/src/views/rocket.py +++ b/src/views/rocket.py @@ -3,6 +3,38 @@ from src.models.rocket import RocketModel from src.views.interface import ApiBaseView from src.views.motor import MotorView, MotorSimulation +# Re-export drawing types from src.views.drawing so callers that previously +# imported them from src.views.rocket (the original home) keep working. +from src.views.drawing import ( + NoseConeGeometry, + TailGeometry, + FinOutline, + FinsGeometry, + TubeGeometry, + MotorPatch, + MotorDrawingGeometry, + RailButtonsGeometry, + SensorGeometry, + DrawingBounds, +) + +__all__ = [ + "RocketSimulation", + "RocketView", + "RocketCreated", + "RocketRetrieved", + "RocketDrawingGeometry", + "NoseConeGeometry", + "TailGeometry", + "FinOutline", + "FinsGeometry", + "TubeGeometry", + "MotorPatch", + "MotorDrawingGeometry", + "RailButtonsGeometry", + "SensorGeometry", + "DrawingBounds", +] class RocketSimulation(MotorSimulation): @@ -64,6 +96,32 @@ class RocketView(RocketModel): motor: MotorView +class RocketDrawingGeometry(ApiBaseView): + """ + Geometry payload that mirrors what ``rocketpy.Rocket.draw()`` feeds to + matplotlib, but as raw coordinate arrays instead of a rendered figure. + All x/y values are already in the rocket drawing frame (the csys-applied + axial direction matches what ``_RocketPlots`` would plot). + """ + + model_config = ConfigDict(ser_json_exclude_none=True) + + message: str = "Rocket drawing geometry retrieved" + radius: float + csys: int + coordinate_system_orientation: str + nose_cones: list[NoseConeGeometry] = [] + tails: list[TailGeometry] = [] + fins: list[FinsGeometry] = [] + tubes: list[TubeGeometry] = [] + motor: Optional[MotorDrawingGeometry] = None + rail_buttons: Optional[RailButtonsGeometry] = None + center_of_mass: Optional[float] = None + cp_position: Optional[float] = None + sensors: list[SensorGeometry] = [] + bounds: DrawingBounds + + class RocketCreated(ApiBaseView): message: str = "Rocket successfully created" rocket_id: str diff --git a/tests/unit/test_routes/conftest.py b/tests/unit/test_routes/conftest.py index 74c63f9..9eb0b94 100644 --- a/tests/unit/test_routes/conftest.py +++ b/tests/unit/test_routes/conftest.py @@ -17,12 +17,16 @@ def stub_environment_dump(): @pytest.fixture def stub_motor_dump(): + # Non-zero dry_inertia so tests that override motor_kind to SOLID/LIQUID/HYBRID + # pass the validate_dry_inertia_for_kind guard. GENERIC still accepts (0, 0, 0) + # at the model level, but we use a non-default value here to keep the stub + # compatible with every motor_kind. motor = MotorModel( thrust_source=[[0, 0]], burn_time=0, nozzle_radius=0, dry_mass=0, - dry_inertia=[0, 0, 0], + dry_inertia=[0.1, 0.1, 0.1], center_of_dry_mass_position=0, motor_kind='GENERIC', )