|
1 | 1 | from typing import Self |
2 | 2 |
|
3 | 3 | import dill |
| 4 | +import numpy as np |
4 | 5 |
|
5 | | -from rocketpy.motors.motor import GenericMotor, Motor as RocketPyMotor |
| 6 | +from rocketpy.motors.motor import Motor as RocketPyMotor |
6 | 7 | from rocketpy.motors.solid_motor import SolidMotor |
7 | 8 | from rocketpy.motors.liquid_motor import LiquidMotor |
8 | 9 | from rocketpy.motors.hybrid_motor import HybridMotor |
9 | | -from rocketpy import ( |
| 10 | +from rocketpy.motors import ( |
| 11 | + EmptyMotor, |
| 12 | + GenericMotor, |
10 | 13 | LevelBasedTank, |
11 | 14 | MassBasedTank, |
12 | 15 | MassFlowRateBasedTank, |
|
16 | 19 |
|
17 | 20 | from fastapi import HTTPException, status |
18 | 21 |
|
| 22 | +from src import logger |
19 | 23 | from src.models.sub.tanks import TankKinds |
20 | 24 | from src.models.motor import MotorKinds, MotorModel |
21 | | -from src.views.motor import MotorSimulation |
| 25 | +from src.views.motor import MotorSimulation, MotorDrawingGeometryView |
| 26 | +from src.views.drawing import ( |
| 27 | + DrawingBounds, |
| 28 | + MotorDrawingGeometry, |
| 29 | + MotorPatch, |
| 30 | +) |
22 | 31 | from src.utils import collect_attributes |
23 | 32 |
|
24 | 33 |
|
| 34 | +# --------------------------------------------------------------------------- |
| 35 | +# Drawing-geometry helpers (module-level, shared by the motor and rocket |
| 36 | +# drawing paths). Kept private with the `_` prefix; consumers should go |
| 37 | +# through `MotorService.get_drawing_geometry` for the motor-only response |
| 38 | +# or through `RocketService.get_drawing_geometry` for the composed rocket |
| 39 | +# view (which delegates motor assembly here). |
| 40 | +# --------------------------------------------------------------------------- |
| 41 | +def _polygon_xy(patch) -> dict: |
| 42 | + """Extract (x, y) coordinate lists from a matplotlib Polygon patch. |
| 43 | +
|
| 44 | + Rocketpy's `_MotorPlots` generator helpers return matplotlib `Polygon` |
| 45 | + objects; we only ever read `patch.xy` (an Nx2 numpy array) as a data |
| 46 | + carrier, never for rendering. |
| 47 | + """ |
| 48 | + xy = np.asarray(patch.xy) |
| 49 | + return {"x": xy[:, 0].tolist(), "y": xy[:, 1].tolist()} |
| 50 | + |
| 51 | + |
| 52 | +def _rebuild_polygon(x: list[float], y: list[float]): |
| 53 | + """Rebuild a matplotlib `Polygon` from coordinate lists. |
| 54 | +
|
| 55 | + Used only so `_MotorPlots._generate_motor_region` can read `patch.xy` |
| 56 | + bounds when we assemble the motor outline. |
| 57 | + """ |
| 58 | + from matplotlib.patches import ( |
| 59 | + Polygon, |
| 60 | + ) # local import keeps service cold-start lean |
| 61 | + |
| 62 | + return Polygon(np.column_stack([np.asarray(x), np.asarray(y)])) |
| 63 | + |
| 64 | + |
| 65 | +def _build_generic_chamber_patch( |
| 66 | + center_x: float, chamber_height: float, chamber_radius: float |
| 67 | +): |
| 68 | + """Build a rectangular combustion-chamber polygon for a GenericMotor. |
| 69 | +
|
| 70 | + Mirrors the vertex order of |
| 71 | + `rocketpy.plots.motor_plots._generate_combustion_chamber` so the patch |
| 72 | + can flow through `_generate_motor_region` for outline computation |
| 73 | + identically to a SolidMotor chamber. |
| 74 | + """ |
| 75 | + from matplotlib.patches import ( |
| 76 | + Polygon, |
| 77 | + ) # local import keeps service cold-start lean |
| 78 | + |
| 79 | + half_len = chamber_height / 2.0 |
| 80 | + x = np.array([-half_len, half_len, half_len, -half_len]) |
| 81 | + y = np.array( |
| 82 | + [chamber_radius, chamber_radius, -chamber_radius, -chamber_radius] |
| 83 | + ) |
| 84 | + x = x + center_x |
| 85 | + return Polygon(np.column_stack([x, y])) |
| 86 | + |
| 87 | + |
25 | 88 | class MotorService: |
26 | 89 | _motor: RocketPyMotor |
27 | 90 |
|
@@ -184,3 +247,193 @@ def get_motor_binary(self) -> bytes: |
184 | 247 | bytes |
185 | 248 | """ |
186 | 249 | return dill.dumps(self.motor) |
| 250 | + |
| 251 | + # -------------------------------------------------------------------- |
| 252 | + # Drawing geometry |
| 253 | + # -------------------------------------------------------------------- |
| 254 | + def get_drawing_geometry( |
| 255 | + self, |
| 256 | + motor_position: float = 0.0, |
| 257 | + parent_csys: int = 1, |
| 258 | + ) -> MotorDrawingGeometryView: |
| 259 | + """Build a motor-only drawing-geometry response. |
| 260 | +
|
| 261 | + Defaults render the motor at its own coordinate origin |
| 262 | + (`motor_position=0`, `parent_csys=1`) so the payload can be used |
| 263 | + standalone in the playground's "motor-only" view. Callers that |
| 264 | + embed the motor into a rocket (see ``RocketService``) pass the |
| 265 | + rocket-level position + csys to align the motor inside the rocket |
| 266 | + frame. |
| 267 | +
|
| 268 | + Raises: |
| 269 | + HTTPException 422 if the motor has no drawable patches. |
| 270 | + """ |
| 271 | + motor_geometry, _nozzle_position = self.build_drawing_geometry( |
| 272 | + motor_position=motor_position, parent_csys=parent_csys |
| 273 | + ) |
| 274 | + if motor_geometry is None: |
| 275 | + raise HTTPException( |
| 276 | + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, |
| 277 | + detail="Motor has no drawable geometry.", |
| 278 | + ) |
| 279 | + # Fall back to the nozzle radius when the motor has no drawable |
| 280 | + # patches (e.g. EmptyMotor) so the bounds aren't a zero-height |
| 281 | + # line. Every real rocketpy motor class exposes nozzle_radius. |
| 282 | + fallback_radius = float(getattr(self._motor, "nozzle_radius", 0.0) or 0.0) |
| 283 | + return MotorDrawingGeometryView( |
| 284 | + motor=motor_geometry, |
| 285 | + coordinate_system_orientation=str( |
| 286 | + self._motor.coordinate_system_orientation |
| 287 | + ), |
| 288 | + bounds=_compute_motor_bounds(motor_geometry, fallback_radius), |
| 289 | + ) |
| 290 | + |
| 291 | + def build_drawing_geometry( |
| 292 | + self, |
| 293 | + motor_position: float, |
| 294 | + parent_csys: int, |
| 295 | + ) -> tuple[MotorDrawingGeometry | None, float]: |
| 296 | + """Construct motor patches + the absolute nozzle x-position. |
| 297 | +
|
| 298 | + Returned as a tuple so the rocket service can use the nozzle |
| 299 | + position when extending body tubes to meet the motor. Standalone |
| 300 | + motor rendering can discard the second element. |
| 301 | + """ |
| 302 | + motor = self._motor |
| 303 | + total_csys = parent_csys * motor._csys |
| 304 | + nozzle_position = motor_position + motor.nozzle_position * total_csys |
| 305 | + |
| 306 | + if isinstance(motor, EmptyMotor): |
| 307 | + return ( |
| 308 | + MotorDrawingGeometry( |
| 309 | + type="empty", |
| 310 | + position=float(motor_position), |
| 311 | + nozzle_position=float(nozzle_position), |
| 312 | + patches=[], |
| 313 | + ), |
| 314 | + float(nozzle_position), |
| 315 | + ) |
| 316 | + |
| 317 | + patches: list[MotorPatch] = [] |
| 318 | + grains_cm_position: float | None = None |
| 319 | + motor_type = "generic" |
| 320 | + |
| 321 | + if isinstance(motor, SolidMotor): |
| 322 | + motor_type = "solid" |
| 323 | + grains_cm_position = ( |
| 324 | + motor_position |
| 325 | + + motor.grains_center_of_mass_position * total_csys |
| 326 | + ) |
| 327 | + chamber = motor.plots._generate_combustion_chamber( |
| 328 | + translate=(grains_cm_position, 0), label=None |
| 329 | + ) |
| 330 | + patches.append(MotorPatch(role="chamber", **_polygon_xy(chamber))) |
| 331 | + for grain in motor.plots._generate_grains( |
| 332 | + translate=(grains_cm_position, 0) |
| 333 | + ): |
| 334 | + patches.append(MotorPatch(role="grain", **_polygon_xy(grain))) |
| 335 | + elif isinstance(motor, HybridMotor): |
| 336 | + motor_type = "hybrid" |
| 337 | + grains_cm_position = ( |
| 338 | + motor_position |
| 339 | + + motor.grains_center_of_mass_position * total_csys |
| 340 | + ) |
| 341 | + chamber = motor.plots._generate_combustion_chamber( |
| 342 | + translate=(grains_cm_position, 0), label=None |
| 343 | + ) |
| 344 | + patches.append(MotorPatch(role="chamber", **_polygon_xy(chamber))) |
| 345 | + for grain in motor.plots._generate_grains( |
| 346 | + translate=(grains_cm_position, 0) |
| 347 | + ): |
| 348 | + patches.append(MotorPatch(role="grain", **_polygon_xy(grain))) |
| 349 | + for tank, _center in motor.plots._generate_positioned_tanks( |
| 350 | + translate=(motor_position, 0), csys=total_csys |
| 351 | + ): |
| 352 | + patches.append(MotorPatch(role="tank", **_polygon_xy(tank))) |
| 353 | + elif isinstance(motor, LiquidMotor): |
| 354 | + motor_type = "liquid" |
| 355 | + for tank, _center in motor.plots._generate_positioned_tanks( |
| 356 | + translate=(motor_position, 0), csys=total_csys |
| 357 | + ): |
| 358 | + patches.append(MotorPatch(role="tank", **_polygon_xy(tank))) |
| 359 | + elif isinstance(motor, GenericMotor): |
| 360 | + # RocketPy's Rocket.draw() does not render a chamber for |
| 361 | + # GenericMotor — `_generate_combustion_chamber` depends on |
| 362 | + # grain fields GenericMotor lacks. We build an equivalent |
| 363 | + # rectangular chamber from the GenericMotor fields so users |
| 364 | + # see their chamber geometry in the playground. |
| 365 | + motor_type = "generic" |
| 366 | + chamber_center_x = ( |
| 367 | + motor_position + motor.chamber_position * total_csys |
| 368 | + ) |
| 369 | + chamber_patch = _build_generic_chamber_patch( |
| 370 | + center_x=chamber_center_x, |
| 371 | + chamber_height=motor.chamber_height, |
| 372 | + chamber_radius=motor.chamber_radius, |
| 373 | + ) |
| 374 | + patches.append( |
| 375 | + MotorPatch(role="chamber", **_polygon_xy(chamber_patch)) |
| 376 | + ) |
| 377 | + |
| 378 | + # Nozzle is always appended after the body so the motor-region |
| 379 | + # outline encompasses it, matching rocketpy. |
| 380 | + nozzle_patch = motor.plots._generate_nozzle( |
| 381 | + translate=(nozzle_position, 0), csys=parent_csys |
| 382 | + ) |
| 383 | + patches.append(MotorPatch(role="nozzle", **_polygon_xy(nozzle_patch))) |
| 384 | + |
| 385 | + # Motor-region outline. `_generate_motor_region` reads patch.xy |
| 386 | + # arrays, so we rebuild matplotlib Polygons once from our |
| 387 | + # coordinate copies. Any failure here is logged and dropped; the |
| 388 | + # outline is advisory, not load-bearing. |
| 389 | + try: |
| 390 | + mpl_patches = [_rebuild_polygon(p.x, p.y) for p in patches] |
| 391 | + outline_patch = motor.plots._generate_motor_region( |
| 392 | + list_of_patches=mpl_patches |
| 393 | + ) |
| 394 | + patches.insert( |
| 395 | + 0, MotorPatch(role="outline", **_polygon_xy(outline_patch)) |
| 396 | + ) |
| 397 | + except Exception as exc: # pragma: no cover - defensive |
| 398 | + logger.warning("Failed to generate motor outline patch: %s", exc) |
| 399 | + |
| 400 | + return ( |
| 401 | + MotorDrawingGeometry( |
| 402 | + type=motor_type, |
| 403 | + position=float(motor_position), |
| 404 | + nozzle_position=float(nozzle_position), |
| 405 | + grains_center_of_mass_position=( |
| 406 | + float(grains_cm_position) |
| 407 | + if grains_cm_position is not None |
| 408 | + else None |
| 409 | + ), |
| 410 | + patches=patches, |
| 411 | + ), |
| 412 | + float(nozzle_position), |
| 413 | + ) |
| 414 | + |
| 415 | + |
| 416 | +def _compute_motor_bounds( |
| 417 | + motor: MotorDrawingGeometry, radius: float |
| 418 | +) -> DrawingBounds: |
| 419 | + """Compute a tight bounding box for a motor-only drawing payload. |
| 420 | +
|
| 421 | + Mirrors the rocket-side bounds helper but limited to motor patches — |
| 422 | + callers who want rocket-wide bounds use the rocket service's own |
| 423 | + computation. |
| 424 | + """ |
| 425 | + xs: list[float] = [] |
| 426 | + ys: list[float] = [] |
| 427 | + for patch in motor.patches: |
| 428 | + xs += patch.x |
| 429 | + ys += patch.y |
| 430 | + if not xs: |
| 431 | + xs = [float(motor.position)] |
| 432 | + if not ys: |
| 433 | + ys = [-float(radius), float(radius)] |
| 434 | + return DrawingBounds( |
| 435 | + x_min=float(min(xs)), |
| 436 | + x_max=float(max(xs)), |
| 437 | + y_min=float(min(ys)), |
| 438 | + y_max=float(max(ys)), |
| 439 | + ) |
0 commit comments