diff --git a/.gitignore b/.gitignore index ef1b4cc..a7b5be3 100644 --- a/.gitignore +++ b/.gitignore @@ -160,4 +160,10 @@ cython_debug/ #.idea/ # VSCode config -.vscode/ \ No newline at end of file +.vscode/ + +# context specific ignores +.context + +# lint +.pylint.d/ \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 9d6f505..85041b3 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/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..113576a 100644 --- a/src/models/motor.py +++ b/src/models/motor.py @@ -20,7 +20,11 @@ class MotorModel(ApiBaseModel): # Required parameters thrust_source: List[List[float]] - burn_time: float + # burn_time is optional for Liquid/Hybrid/Solid motors — rocketpy + # auto-detects the burn window from the thrust_source array span. + # GenericMotor still requires it; the motor service re-raises an + # explicit error when the GENERIC path receives None. + burn_time: Optional[float] = None nozzle_radius: float dry_mass: float dry_inertia: Tuple[float, float, float] = (0, 0, 0) @@ -82,6 +86,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..2dda7c9 100644 --- a/src/models/sub/tanks.py +++ b/src/models/sub/tanks.py @@ -1,6 +1,7 @@ from enum import Enum -from typing import Optional, Tuple, List -from pydantic import BaseModel +from typing import Annotated, List, Literal, Optional, Tuple, Union + +from pydantic import BaseModel, Field, model_validator class TankKinds(str, Enum): @@ -10,19 +11,81 @@ class TankKinds(str, Enum): ULLAGE: str = "ULLAGE" +# Scalar density keeps the legacy behaviour (constant kg/m^3). +# A list of (temperature_K, density_kg_per_m3) samples enables +# temperature-dependent density — required for realistic LOX / N2O +# modelling. Pressure dependence is out of scope for this iteration. +DensityInput = Union[float, List[Tuple[float, float]]] + + class TankFluids(BaseModel): name: str - density: float + density: DensityInput + + +# --- Tank geometry discriminated union ---------------------------------- +# RocketPy ships three concrete geometry classes. We mirror them as a +# discriminated Pydantic union keyed on `geometry_kind`. `custom` is the +# generic piecewise form (original API shape); `cylindrical` and +# `spherical` map to `rocketpy.motors.CylindricalTank` and +# `SphericalTank` respectively. + + +class CustomTankGeometry(BaseModel): + geometry_kind: Literal["custom"] = "custom" + geometry: List[Tuple[Tuple[float, float], float]] + + +class CylindricalTankGeometry(BaseModel): + geometry_kind: Literal["cylindrical"] = "cylindrical" + radius: float + height: float + spherical_caps: bool = False + + +class SphericalTankGeometry(BaseModel): + geometry_kind: Literal["spherical"] = "spherical" + radius: float + + +TankGeometryInput = Annotated[ + Union[ + CustomTankGeometry, + CylindricalTankGeometry, + SphericalTankGeometry, + ], + Field(discriminator="geometry_kind"), +] + + +# Map tank_kind → tuple of MotorTank field names that rocketpy's +# corresponding Tank subclass requires. The validator below rejects +# payloads that omit any of them so the API returns 422 instead of +# letting rocketpy crash during motor construction. +_REQUIRED_FIELDS_BY_TANK_KIND = { + TankKinds.MASS_FLOW: ( + "initial_liquid_mass", + "initial_gas_mass", + "liquid_mass_flow_rate_in", + "liquid_mass_flow_rate_out", + "gas_mass_flow_rate_in", + "gas_mass_flow_rate_out", + ), + TankKinds.LEVEL: ("liquid_height",), + TankKinds.ULLAGE: ("ullage",), + TankKinds.MASS: ("liquid_mass", "gas_mass"), +} class MotorTank(BaseModel): # Required parameters - geometry: List[Tuple[Tuple[float, float], float]] + geometry: TankGeometryInput gas: TankFluids 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 @@ -47,3 +110,20 @@ class MotorTank(BaseModel): # Computed parameters tank_kind: TankKinds = TankKinds.MASS_FLOW + + @model_validator(mode='after') + def validate_tank_kind_fields(self): + # Mirrors the validate_dry_inertia_for_kind pattern used on + # MotorModel: reject incoherent payloads at the API boundary + # instead of letting rocketpy crash during Tank construction. + missing = [ + field + for field in _REQUIRED_FIELDS_BY_TANK_KIND[self.tank_kind] + if getattr(self, field) is None + ] + if missing: + raise ValueError( + f"tank_kind={self.tank_kind.value} requires: " + f"{', '.join(missing)}" + ) + return self diff --git a/src/routes/motor.py b/src/routes/motor.py index 0673708..192d6ad 100644 --- a/src/routes/motor.py +++ b/src/routes/motor.py @@ -36,6 +36,31 @@ async def create_motor( ## Args ``` models.Motor JSON ``` + + For liquid/hybrid motors the `tanks` field supports three geometry + kinds via the `geometry_kind` discriminator and a scalar-or-sampled + fluid `density`: + + ``` + { + "motor_kind": "LIQUID", + ... + "tanks": [{ + "geometry": { + "geometry_kind": "cylindrical", // or "spherical", "custom" + "radius": 0.1, "height": 0.5 + }, + "liquid": { + "name": "LOX", + "density": [[90.0, 1141.0], [120.0, 1091.0]] // or scalar + }, + "gas": {"name": "N2", "density": 1.2}, + "tank_kind": "LEVEL", + "liquid_height": 0.25, + ... + }] + } + ``` """ with tracer.start_as_current_span("create_motor"): return await controller.post_motor(motor) 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/flight.py b/src/services/flight.py index c976862..b1eb21b 100644 --- a/src/services/flight.py +++ b/src/services/flight.py @@ -224,6 +224,22 @@ def _to_float(value) -> float: case _: return float(value) + @staticmethod + def _extract_fluid_density(fluid): + """Project a rocketpy Fluid's density back onto the API schema. + + The API accepts either a scalar or a list of (T_K, density) + samples. Rocketpy may store density as either a raw scalar or a + ``Function`` wrapping a 2D ``(T, P) -> density`` callable. A + full sample round-trip is not supported in this iteration; + Function-valued densities are collapsed to a scalar evaluated + at rocketpy's default reference (273.15 K, 101325 Pa). + """ + density = fluid.density + if isinstance(density, Function): + return float(density(273.15, 101325)) + return density + @staticmethod def _extract_tanks(motor) -> list[MotorTank]: tanks: list[MotorTank] = [] @@ -240,20 +256,29 @@ def _extract_tanks(motor) -> list[MotorTank]: case _: tank_kind = TankKinds.MASS_FLOW - geometry = [ + # Geometry round-trip is lossy: even if the client originally + # sent a cylindrical/spherical geometry, we discretise it back + # to the generic piecewise form on read. Every rocketpy tank + # geometry exposes its internal piecewise dict via + # `tank.geometry.geometry`, so this path covers all three + # geometry subclasses uniformly. + geometry_segments = [ (bounds, float(func(0))) for bounds, func in tank.geometry.geometry.items() ] data: dict = { - "geometry": geometry, + "geometry": { + "geometry_kind": "custom", + "geometry": geometry_segments, + }, "gas": TankFluids( name=tank.gas.name, - density=tank.gas.density, + density=FlightService._extract_fluid_density(tank.gas), ), "liquid": TankFluids( name=tank.liquid.name, - density=tank.liquid.density, + density=FlightService._extract_fluid_density(tank.liquid), ), "flux_time": tank.flux_time, "position": position, diff --git a/src/services/motor.py b/src/services/motor.py index 7275920..32ca40a 100644 --- a/src/services/motor.py +++ b/src/services/motor.py @@ -7,21 +7,83 @@ from rocketpy.motors.liquid_motor import LiquidMotor from rocketpy.motors.hybrid_motor import HybridMotor from rocketpy import ( + CylindricalTank, + Fluid, + Function, LevelBasedTank, MassBasedTank, MassFlowRateBasedTank, - UllageBasedTank, + SphericalTank, TankGeometry, + UllageBasedTank, ) from fastapi import HTTPException, status -from src.models.sub.tanks import TankKinds +from src.models.sub.tanks import ( + CustomTankGeometry, + CylindricalTankGeometry, + SphericalTankGeometry, + TankFluids, + TankKinds, +) from src.models.motor import MotorKinds, MotorModel from src.views.motor import MotorSimulation from src.utils import collect_attributes +def _build_rocketpy_tank_geometry(geometry): + """Convert an API geometry model into a rocketpy geometry object. + + Dispatch mirrors the discriminated union in + ``src.models.sub.tanks.TankGeometryInput``. + """ + if isinstance(geometry, CylindricalTankGeometry): + return CylindricalTank( + radius=geometry.radius, + height=geometry.height, + spherical_caps=geometry.spherical_caps, + ) + if isinstance(geometry, SphericalTankGeometry): + return SphericalTank(radius=geometry.radius) + if isinstance(geometry, CustomTankGeometry): + return TankGeometry(geometry_dict=dict(geometry.geometry)) + raise ValueError( + f"Unsupported tank geometry kind: {type(geometry).__name__}" + ) + + +def _build_rocketpy_fluid(fluids: TankFluids) -> Fluid: + """Convert an API TankFluids into a rocketpy Fluid. + + Scalar density is passed through (Fluid stores it as a constant). + Sampled density is converted to a 1D Temperature → Density Function + and wrapped in a ``(T, P)`` callable because rocketpy's Fluid expects + density to be a function of both temperature and pressure. Pressure + is ignored here intentionally; only temperature-dependent density + is supported in this iteration. + """ + density = fluids.density + if isinstance(density, list): + temperature_to_density = Function( + source=density, + interpolation='linear', + extrapolation='natural', + inputs=['Temperature (K)'], + outputs='Density (kg/m^3)', + ) + + def density_callable(temperature, pressure): # noqa: ARG001 + # pylint: disable=unused-argument + # Rocketpy's Fluid wraps this into a 2-input Function of + # (T, P); pressure is accepted for signature compatibility + # but intentionally ignored in this iteration. + return temperature_to_density.get_value(temperature) + + return Fluid(name=fluids.name, density=density_callable) + return Fluid(name=fluids.name, density=density) + + class MotorService: _motor: RocketPyMotor @@ -45,7 +107,6 @@ def from_motor_model(cls, motor: MotorModel) -> Self: motor_core = { "thrust_source": motor.thrust_source, - "burn_time": motor.burn_time, "nozzle_radius": motor.nozzle_radius, "dry_mass": motor.dry_mass, "dry_inertia": motor.dry_inertia, @@ -54,6 +115,13 @@ def from_motor_model(cls, motor: MotorModel) -> Self: "interpolation_method": motor.interpolation_method, "reshape_thrust_curve": reshape_thrust_curve, } + # Only forward optional rocketpy args when the client supplied them. + # Leaving them out lets rocketpy pick its own default (burn_time + # auto-detected from thrust_source span; nozzle_position = 0). + if motor.burn_time is not None: + motor_core["burn_time"] = motor.burn_time + if motor.nozzle_position is not None: + motor_core["nozzle_position"] = motor.nozzle_position match MotorKinds(motor.motor_kind): case MotorKinds.LIQUID: @@ -103,25 +171,34 @@ def from_motor_model(cls, motor: MotorModel) -> Self: **optional_params, ) case _: + # GenericMotor requires burn_time even though it's optional + # for the other motor kinds — surface the constraint at the + # API boundary instead of letting rocketpy raise a + # confusing stack trace deeper in construction. + if motor.burn_time is None: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail="burn_time is required for generic motors.", + ) + # nozzle_position is already forwarded via motor_core when + # the client supplied it; GenericMotor's own default (0) + # applies otherwise. rocketpy_motor = GenericMotor( **motor_core, chamber_radius=motor.chamber_radius, chamber_height=motor.chamber_height, chamber_position=motor.chamber_position, propellant_initial_mass=motor.propellant_initial_mass, - nozzle_position=motor.nozzle_position, ) if motor.motor_kind not in (MotorKinds.SOLID, MotorKinds.GENERIC): for tank in motor.tanks or []: tank_core = { "name": tank.name, - "geometry": TankGeometry( - geometry_dict=dict(tank.geometry) - ), + "geometry": _build_rocketpy_tank_geometry(tank.geometry), "flux_time": tank.flux_time, - "gas": tank.gas, - "liquid": tank.liquid, + "gas": _build_rocketpy_fluid(tank.gas), + "liquid": _build_rocketpy_fluid(tank.liquid), "discretize": tank.discretize, } diff --git a/src/services/rocket.py b/src/services/rocket.py index a22a2d4..70eade4 100644 --- a/src/services/rocket.py +++ b/src/services/rocket.py @@ -1,7 +1,15 @@ from typing import Self, List import dill - +import numpy as np + +from rocketpy.motors import ( + EmptyMotor, + GenericMotor, + HybridMotor, + LiquidMotor, + SolidMotor, +) from rocketpy.rocket.rocket import Rocket as RocketPyRocket from rocketpy.rocket.parachute import Parachute as RocketPyParachute from rocketpy.rocket.aero_surface import ( @@ -11,6 +19,7 @@ Fins as RocketPyFins, Tail as RocketPyTail, ) +from rocketpy.rocket.aero_surface.generic_surface import GenericSurface from fastapi import HTTPException, status @@ -18,7 +27,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 +147,426 @@ 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]: + rocket = self._rocket + motor = rocket.motor + total_csys = csys * motor._csys + nozzle_position = ( + rocket.motor_position + motor.nozzle_position * total_csys + ) + + if isinstance(motor, EmptyMotor): + return ( + MotorDrawingGeometry( + type="empty", + position=float(rocket.motor_position), + nozzle_position=float(nozzle_position), + patches=[], + ), + float(nozzle_position), + ) + + patches: list[MotorPatch] = [] + grains_cm_position: float | None = None + + if isinstance(motor, SolidMotor): + motor_type = "solid" + grains_cm_position = ( + rocket.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 = ( + rocket.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=(rocket.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=(rocket.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 — + # _MotorPlots._generate_combustion_chamber depends on grain fields that + # GenericMotor lacks. We build an equivalent rectangular chamber here + # using the GenericMotor-specific fields (chamber_radius, chamber_height, + # chamber_position) so users can see their chamber geometry in the playground. + motor_type = "generic" + chamber_center_x = ( + rocket.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)) + ) + else: + motor_type = "generic" + + # Nozzle (added after so the outline encompasses it, matching rocketpy) + nozzle_patch = motor.plots._generate_nozzle( + translate=(nozzle_position, 0), csys=csys + ) + patches.append(MotorPatch(role="nozzle", **_polygon_xy(nozzle_patch))) + + # Motor region outline. _generate_motor_region reads patch.xy arrays + # so we need matplotlib Polygon objects; rebuild them once from our + # coordinate copies. + 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(rocket.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 _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 +706,73 @@ def check_parachute_trigger(trigger) -> bool: if isinstance(trigger, (int, float)): return True return False + + +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 resulting patch can flow through _generate_motor_region for outline + computation identically to a SolidMotor chamber. + + Parameters + ---------- + center_x : float + World-frame x-coordinate of the chamber centroid (already includes + rocket.motor_position + motor.chamber_position * csys). + chamber_height : float + Axial length of the chamber (m). + chamber_radius : float + Internal radius of the chamber (m). + """ + from matplotlib.patches import ( + Polygon, + ) # local import keeps service cold-start lean + + half_len = chamber_height / 2.0 + # Top edge then mirror to the bottom, matching _generate_combustion_chamber's + # vertex ordering so motor-region assembly sees a consistent shape. + 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])) + + +def _polygon_xy(patch) -> dict: + """Extract (x, y) coordinate lists from a matplotlib Polygon patch. + + The generator helpers on rocketpy's _MotorPlots return matplotlib + Polygon objects; we only ever use them as a data carrier (patch.xy + is an Nx2 numpy array), 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)])) diff --git a/src/views/rocket.py b/src/views/rocket.py index 25d08fe..17ed05a 100644 --- a/src/views/rocket.py +++ b/src/views/rocket.py @@ -1,5 +1,5 @@ -from typing import Optional, Any -from pydantic import ConfigDict +from typing import Optional, Any, Literal +from pydantic import BaseModel, ConfigDict from src.models.rocket import RocketModel from src.views.interface import ApiBaseView from src.views.motor import MotorView, MotorSimulation @@ -64,6 +64,106 @@ class RocketView(RocketModel): motor: MotorView +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 + + +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..ae09b42 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', ) @@ -32,51 +36,117 @@ def stub_motor_dump(): @pytest.fixture def stub_tank_dump(): + # Base fixture defaults to MASS_FLOW (matches MotorTank's own default) + # so the tank validates standalone. Sub-variant fixtures below switch + # tank_kind and populate only the fields that variant requires. tank = MotorTank( - geometry=[[(0, 0), 0]], + geometry={ + 'geometry_kind': 'custom', + 'geometry': [[(0, 0), 0]], + }, gas=TankFluids(name='gas', density=0), liquid=TankFluids(name='liquid', density=0), flux_time=(0, 0), position=0, discretize=0, name='tank', + gas_mass_flow_rate_in=0, + gas_mass_flow_rate_out=0, + liquid_mass_flow_rate_in=0, + liquid_mass_flow_rate_out=0, + initial_liquid_mass=0, + initial_gas_mass=0, ) tank_json = tank.model_dump_json() return json.loads(tank_json) @pytest.fixture -def stub_level_tank_dump(stub_tank_dump): - stub_tank_dump.update({'tank_kind': TankKinds.LEVEL, 'liquid_height': 0}) +def stub_cylindrical_tank_dump(stub_tank_dump): + stub_tank_dump['geometry'] = { + 'geometry_kind': 'cylindrical', + 'radius': 0.1, + 'height': 0.5, + 'spherical_caps': False, + } return stub_tank_dump @pytest.fixture -def stub_mass_flow_tank_dump(stub_tank_dump): +def stub_spherical_tank_dump(stub_tank_dump): + stub_tank_dump['geometry'] = { + 'geometry_kind': 'spherical', + 'radius': 0.2, + } + return stub_tank_dump + + +@pytest.fixture +def stub_tank_with_sampled_density_dump(stub_tank_dump): + stub_tank_dump['liquid'] = { + 'name': 'LOX', + 'density': [[90.0, 1141.0], [120.0, 1091.0], [150.0, 1021.0]], + } + return stub_tank_dump + + +@pytest.fixture +def stub_level_tank_dump(stub_tank_dump): + # Switch out of the MASS_FLOW defaults into LEVEL, clearing the + # unused MASS_FLOW fields so the kind-specific validator passes. stub_tank_dump.update( { - 'tank_kind': TankKinds.MASS_FLOW, - 'gas_mass_flow_rate_in': 0, - 'gas_mass_flow_rate_out': 0, - 'liquid_mass_flow_rate_in': 0, - 'liquid_mass_flow_rate_out': 0, - 'initial_liquid_mass': 0, - 'initial_gas_mass': 0, + 'tank_kind': TankKinds.LEVEL, + 'liquid_height': 0, + 'gas_mass_flow_rate_in': None, + 'gas_mass_flow_rate_out': None, + 'liquid_mass_flow_rate_in': None, + 'liquid_mass_flow_rate_out': None, + 'initial_liquid_mass': None, + 'initial_gas_mass': None, } ) return stub_tank_dump +@pytest.fixture +def stub_mass_flow_tank_dump(stub_tank_dump): + # stub_tank_dump already includes all MASS_FLOW fields. + stub_tank_dump['tank_kind'] = TankKinds.MASS_FLOW + return stub_tank_dump + + @pytest.fixture def stub_ullage_tank_dump(stub_tank_dump): - stub_tank_dump.update({'tank_kind': TankKinds.ULLAGE, 'ullage': 0}) + stub_tank_dump.update( + { + 'tank_kind': TankKinds.ULLAGE, + 'ullage': 0, + 'gas_mass_flow_rate_in': None, + 'gas_mass_flow_rate_out': None, + 'liquid_mass_flow_rate_in': None, + 'liquid_mass_flow_rate_out': None, + 'initial_liquid_mass': None, + 'initial_gas_mass': None, + } + ) return stub_tank_dump @pytest.fixture def stub_mass_tank_dump(stub_tank_dump): stub_tank_dump.update( - {'tank_kind': TankKinds.MASS, 'liquid_mass': 0, 'gas_mass': 0} + { + 'tank_kind': TankKinds.MASS, + 'liquid_mass': 0, + 'gas_mass': 0, + 'gas_mass_flow_rate_in': None, + 'gas_mass_flow_rate_out': None, + 'liquid_mass_flow_rate_in': None, + 'liquid_mass_flow_rate_out': None, + 'initial_liquid_mass': None, + 'initial_gas_mass': None, + } ) return stub_tank_dump diff --git a/tests/unit/test_routes/test_motors_route.py b/tests/unit/test_routes/test_motors_route.py index c4a01d2..52fe794 100644 --- a/tests/unit/test_routes/test_motors_route.py +++ b/tests/unit/test_routes/test_motors_route.py @@ -182,6 +182,152 @@ def test_create_liquid_motor_mass_tank( ) +def test_create_liquid_motor_cylindrical_geometry( + stub_motor_dump, stub_cylindrical_tank_dump, mock_controller_instance +): + stub_cylindrical_tank_dump.update( + {'tank_kind': 'LEVEL', 'liquid_height': 0.2} + ) + stub_cylindrical_tank_dump.update( + { + 'gas_mass_flow_rate_in': None, + 'gas_mass_flow_rate_out': None, + 'liquid_mass_flow_rate_in': None, + 'liquid_mass_flow_rate_out': None, + 'initial_liquid_mass': None, + 'initial_gas_mass': None, + } + ) + stub_motor_dump.update( + {'tanks': [stub_cylindrical_tank_dump], 'motor_kind': 'LIQUID'} + ) + mock_response = AsyncMock(return_value=MotorCreated(motor_id='123')) + mock_controller_instance.post_motor = mock_response + response = client.post('/motors/', json=stub_motor_dump) + assert response.status_code == 201 + + +def test_create_liquid_motor_spherical_geometry( + stub_motor_dump, stub_spherical_tank_dump, mock_controller_instance +): + stub_spherical_tank_dump.update( + { + 'tank_kind': 'LEVEL', + 'liquid_height': 0.1, + 'gas_mass_flow_rate_in': None, + 'gas_mass_flow_rate_out': None, + 'liquid_mass_flow_rate_in': None, + 'liquid_mass_flow_rate_out': None, + 'initial_liquid_mass': None, + 'initial_gas_mass': None, + } + ) + stub_motor_dump.update( + {'tanks': [stub_spherical_tank_dump], 'motor_kind': 'LIQUID'} + ) + mock_response = AsyncMock(return_value=MotorCreated(motor_id='123')) + mock_controller_instance.post_motor = mock_response + response = client.post('/motors/', json=stub_motor_dump) + assert response.status_code == 201 + + +def test_create_liquid_motor_sampled_density( + stub_motor_dump, + stub_tank_with_sampled_density_dump, + mock_controller_instance, +): + stub_motor_dump.update( + { + 'tanks': [stub_tank_with_sampled_density_dump], + 'motor_kind': 'LIQUID', + } + ) + mock_response = AsyncMock(return_value=MotorCreated(motor_id='123')) + mock_controller_instance.post_motor = mock_response + response = client.post('/motors/', json=stub_motor_dump) + assert response.status_code == 201 + + +def test_create_motor_invalid_geometry_kind( + stub_motor_dump, stub_tank_dump, mock_controller_instance +): + stub_tank_dump['geometry'] = { + 'geometry_kind': 'pyramid', + 'radius': 0.1, + } + stub_motor_dump.update( + {'tanks': [stub_tank_dump], 'motor_kind': 'LIQUID'} + ) + response = client.post('/motors/', json=stub_motor_dump) + assert response.status_code == 422 + + +def test_create_motor_mass_kind_missing_fields( + stub_motor_dump, stub_tank_dump, mock_controller_instance +): + # stub_tank_dump defaults to MASS_FLOW with all required fields + # populated; switching to MASS without adding liquid_mass/gas_mass + # must trigger the tank_kind guard at schema validation. + stub_tank_dump['tank_kind'] = 'MASS' + stub_motor_dump.update( + {'tanks': [stub_tank_dump], 'motor_kind': 'LIQUID'} + ) + response = client.post('/motors/', json=stub_motor_dump) + assert response.status_code == 422 + body = response.json() + assert 'liquid_mass' in json.dumps(body) + assert 'gas_mass' in json.dumps(body) + + +def test_create_motor_level_kind_missing_liquid_height( + stub_motor_dump, stub_tank_dump, mock_controller_instance +): + stub_tank_dump['tank_kind'] = 'LEVEL' + stub_motor_dump.update( + {'tanks': [stub_tank_dump], 'motor_kind': 'LIQUID'} + ) + response = client.post('/motors/', json=stub_motor_dump) + assert response.status_code == 422 + assert 'liquid_height' in json.dumps(response.json()) + + +def test_create_motor_ullage_kind_missing_ullage( + stub_motor_dump, stub_tank_dump, mock_controller_instance +): + stub_tank_dump['tank_kind'] = 'ULLAGE' + stub_motor_dump.update( + {'tanks': [stub_tank_dump], 'motor_kind': 'LIQUID'} + ) + response = client.post('/motors/', json=stub_motor_dump) + assert response.status_code == 422 + assert 'ullage' in json.dumps(response.json()) + + +def test_create_motor_mass_flow_kind_missing_flow_rates( + stub_motor_dump, mock_controller_instance +): + # Build a tank payload with MASS_FLOW kind but no flow-rate fields + # so the guard rejects it. + tank_payload = { + 'geometry': { + 'geometry_kind': 'custom', + 'geometry': [[[0, 0], 0]], + }, + 'gas': {'name': 'gas', 'density': 0}, + 'liquid': {'name': 'liquid', 'density': 0}, + 'flux_time': [0, 0], + 'position': 0, + 'discretize': 0, + 'tank_kind': 'MASS_FLOW', + } + stub_motor_dump.update( + {'tanks': [tank_payload], 'motor_kind': 'LIQUID'} + ) + response = client.post('/motors/', json=stub_motor_dump) + assert response.status_code == 422 + assert 'initial_liquid_mass' in json.dumps(response.json()) + + def test_create_hybrid_motor( stub_motor_dump, stub_level_tank_dump, mock_controller_instance ): diff --git a/tests/unit/test_services/__init__.py b/tests/unit/test_services/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/test_services/test_motor_service.py b/tests/unit/test_services/test_motor_service.py new file mode 100644 index 0000000..1bb1d6c --- /dev/null +++ b/tests/unit/test_services/test_motor_service.py @@ -0,0 +1,229 @@ +"""Tests for src.services.motor conversion helpers. + +These exercise the translation from API Pydantic models to concrete +rocketpy objects. Unlike the route tests (which mock the controller), +these use the real rocketpy package so any mismatch between our +adapter and the rocketpy API surfaces immediately. +""" + +import pytest +from rocketpy import ( + CylindricalTank, + Fluid, + Function, + LevelBasedTank, + LiquidMotor, + MassBasedTank, + MassFlowRateBasedTank, + SphericalTank, + TankGeometry, + UllageBasedTank, +) + +from src.models.motor import MotorModel +from src.models.sub.tanks import MotorTank, TankFluids, TankKinds +from src.services.motor import ( + MotorService, + _build_rocketpy_fluid, + _build_rocketpy_tank_geometry, +) + + +_LIQUID_CORE = { + 'thrust_source': [[0, 1000], [1, 500]], + 'burn_time': 1.0, + 'nozzle_radius': 0.05, + 'dry_mass': 5.0, + 'dry_inertia': [0.1, 0.1, 0.01], + 'center_of_dry_mass_position': 0.0, + 'motor_kind': 'LIQUID', +} + + +def _level_tank(geometry, liquid=None): + return MotorTank( + geometry=geometry, + gas=TankFluids(name='N2', density=1.2), + liquid=liquid or TankFluids(name='LOX', density=1141.0), + flux_time=(0, 1.0), + position=0.5, + discretize=100, + tank_kind=TankKinds.LEVEL, + liquid_height=0.2, + ) + + +class TestGeometryAdapter: + def test_custom_geometry_produces_generic_tank_geometry(self): + tank = _level_tank( + { + 'geometry_kind': 'custom', + 'geometry': [[(0.0, 1.0), 0.1]], + } + ) + built = _build_rocketpy_tank_geometry(tank.geometry) + assert isinstance(built, TankGeometry) + assert not isinstance(built, (CylindricalTank, SphericalTank)) + + def test_cylindrical_geometry_produces_cylindrical_tank(self): + tank = _level_tank( + { + 'geometry_kind': 'cylindrical', + 'radius': 0.1, + 'height': 0.5, + 'spherical_caps': False, + } + ) + built = _build_rocketpy_tank_geometry(tank.geometry) + assert isinstance(built, CylindricalTank) + # height is the total cylindrical span + assert built.height == pytest.approx(0.5) + + def test_spherical_geometry_produces_spherical_tank(self): + tank = _level_tank( + {'geometry_kind': 'spherical', 'radius': 0.2} + ) + built = _build_rocketpy_tank_geometry(tank.geometry) + assert isinstance(built, SphericalTank) + + +class TestFluidAdapter: + def test_scalar_density_passes_through(self): + fluid = _build_rocketpy_fluid( + TankFluids(name='water', density=1000.0) + ) + assert isinstance(fluid, Fluid) + assert fluid.density == 1000.0 + + def test_sampled_density_becomes_callable_function(self): + fluid = _build_rocketpy_fluid( + TankFluids( + name='LOX', + density=[[90.0, 1141.0], [120.0, 1091.0], [150.0, 1021.0]], + ) + ) + assert isinstance(fluid, Fluid) + # The 2-input density function interpolates correctly. Rocketpy + # wraps our 1D callable so pressure is accepted but ignored. + assert fluid.density_function.get_value( + 105.0, 1e5 + ) == pytest.approx(1116.0) + assert fluid.density_function.get_value( + 135.0, 1e5 + ) == pytest.approx(1056.0) + + +class TestFromMotorModelLiquid: + def test_mass_flow_tank_with_custom_geometry(self): + # Sized so the propellant fits: radius 0.5 m, height 2 m → + # tank volume ≈ 1.57 m³, which comfortably holds 10 kg of LOX + # (~0.009 m³) plus 0.01 kg of N2 (~0.008 m³). + motor_model = MotorModel( + **_LIQUID_CORE, + tanks=[ + MotorTank( + geometry={ + 'geometry_kind': 'custom', + 'geometry': [[(-1.0, 1.0), 0.5]], + }, + gas=TankFluids(name='N2', density=1.2), + liquid=TankFluids(name='LOX', density=1141.0), + flux_time=(0, 1.0), + position=0.5, + discretize=100, + tank_kind=TankKinds.MASS_FLOW, + initial_liquid_mass=10.0, + initial_gas_mass=0.01, + liquid_mass_flow_rate_in=0.0, + liquid_mass_flow_rate_out=5.0, + gas_mass_flow_rate_in=0.0, + gas_mass_flow_rate_out=0.0, + ), + ], + ) + service = MotorService.from_motor_model(motor_model) + assert isinstance(service.motor, LiquidMotor) + tank = service.motor.positioned_tanks[0]['tank'] + assert isinstance(tank, MassFlowRateBasedTank) + assert isinstance(tank.geometry, TankGeometry) + + def test_level_tank_with_cylindrical_geometry(self): + motor_model = MotorModel( + **_LIQUID_CORE, + tanks=[ + _level_tank( + { + 'geometry_kind': 'cylindrical', + 'radius': 0.1, + 'height': 0.5, + } + ) + ], + ) + service = MotorService.from_motor_model(motor_model) + tank = service.motor.positioned_tanks[0]['tank'] + assert isinstance(tank, LevelBasedTank) + assert isinstance(tank.geometry, CylindricalTank) + + def test_ullage_tank_with_spherical_geometry(self): + motor_model = MotorModel( + **_LIQUID_CORE, + tanks=[ + MotorTank( + geometry={ + 'geometry_kind': 'spherical', + 'radius': 0.2, + }, + gas=TankFluids(name='N2', density=1.2), + liquid=TankFluids(name='LOX', density=1141.0), + flux_time=(0, 1.0), + position=0.5, + discretize=100, + tank_kind=TankKinds.ULLAGE, + ullage=0.01, + ) + ], + ) + service = MotorService.from_motor_model(motor_model) + tank = service.motor.positioned_tanks[0]['tank'] + assert isinstance(tank, UllageBasedTank) + assert isinstance(tank.geometry, SphericalTank) + + def test_mass_tank_keeps_sampled_density_as_function(self): + motor_model = MotorModel( + **_LIQUID_CORE, + tanks=[ + MotorTank( + geometry={ + 'geometry_kind': 'cylindrical', + 'radius': 0.1, + 'height': 0.5, + }, + gas=TankFluids(name='N2', density=1.2), + liquid=TankFluids( + name='LOX', + density=[ + [90.0, 1141.0], + [120.0, 1091.0], + [150.0, 1021.0], + ], + ), + flux_time=(0, 1.0), + position=0.5, + discretize=100, + tank_kind=TankKinds.MASS, + liquid_mass=10.0, + gas_mass=0.001, + ), + ], + ) + service = MotorService.from_motor_model(motor_model) + tank = service.motor.positioned_tanks[0]['tank'] + assert isinstance(tank, MassBasedTank) + # Liquid density survived as a Function wrapping our 1D sampler. + assert isinstance(tank.liquid.density, Function) + assert tank.liquid.density_function.get_value( + 105.0, 1e5 + ) == pytest.approx(1116.0) + # Gas density stayed scalar. + assert tank.gas.density == 1.2