ENH: class encoding refactor for large sims (#58)

Co-authored-by: Gabriel Barberini <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: Pedro Bressan <[email protected]>

Author: 4 people

.gitignore

@@ -158,3 +158,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# VSCode config
+.vscode/

src/models/rocket.py

@@ -19,7 +19,7 @@ class RocketModel(ApiBaseModel):
     radius: float
     mass: float
     motor_position: float
-    center_of_mass_without_motor: int
+    center_of_mass_without_motor: float
     inertia: Union[
         Tuple[float, float, float],
         Tuple[float, float, float, float, float, float],

src/models/sub/aerosurfaces.py

@@ -34,14 +34,15 @@ class Fins(BaseModel):
     root_chord: float
     span: float
     position: float
-
     # Optional parameters
     tip_chord: Optional[float] = None
     cant_angle: Optional[float] = None
     rocket_radius: Optional[float] = None
     airfoil: Optional[
         Tuple[List[Tuple[float, float]], Literal['radians', 'degrees']]
     ] = None
+    sweep_length: Optional[float] = None
+    sweep_angle: Optional[float] = None
 
     def get_additional_parameters(self):
         return {

src/services/environment.py

@@ -5,7 +5,7 @@
 from rocketpy.environment.environment import Environment as RocketPyEnvironment
 from src.models.environment import EnvironmentModel
 from src.views.environment import EnvironmentSimulation
-from src.utils import rocketpy_encoder, DiscretizeConfig
+from src.utils import collect_attributes
 
 
 class EnvironmentService:
@@ -54,10 +54,11 @@ def get_environment_simulation(self) -> EnvironmentSimulation:
             EnvironmentSimulation
         """
 
-        attributes = rocketpy_encoder(
-            self.environment, DiscretizeConfig.for_environment()
+        encoded_attributes = collect_attributes(
+            self.environment,
+            [EnvironmentSimulation],
         )
-        env_simulation = EnvironmentSimulation(**attributes)
+        env_simulation = EnvironmentSimulation(**encoded_attributes)
         return env_simulation
 
     def get_environment_binary(self) -> bytes:

src/services/flight.py

@@ -8,7 +8,10 @@
 from src.services.rocket import RocketService
 from src.models.flight import FlightModel
 from src.views.flight import FlightSimulation
-from src.utils import rocketpy_encoder, DiscretizeConfig
+from src.views.rocket import RocketSimulation
+from src.views.motor import MotorSimulation
+from src.views.environment import EnvironmentSimulation
+from src.utils import collect_attributes
 
 
 class FlightService:
@@ -55,10 +58,16 @@ def get_flight_simulation(self) -> FlightSimulation:
         Returns:
             FlightSimulation
         """
-        attributes = rocketpy_encoder(
-            self.flight, DiscretizeConfig.for_flight()
+        encoded_attributes = collect_attributes(
+            self.flight,
+            [
+                FlightSimulation,
+                RocketSimulation,
+                MotorSimulation,
+                EnvironmentSimulation,
+            ],
         )
-        flight_simulation = FlightSimulation(**attributes)
+        flight_simulation = FlightSimulation(**encoded_attributes)
         return flight_simulation
 
     def get_flight_binary(self) -> bytes:

src/services/motor.py

@@ -17,7 +17,7 @@
 from src.models.sub.tanks import TankKinds
 from src.models.motor import MotorKinds, MotorModel
 from src.views.motor import MotorSimulation
-from src.utils import rocketpy_encoder, DiscretizeConfig
+from src.utils import collect_attributes
 
 
 class MotorService:
@@ -140,8 +140,11 @@ def get_motor_simulation(self) -> MotorSimulation:
         Returns:
             MotorSimulation
         """
-        attributes = rocketpy_encoder(self.motor, DiscretizeConfig.for_motor())
-        motor_simulation = MotorSimulation(**attributes)
+        encoded_attributes = collect_attributes(
+            self.motor,
+            [MotorSimulation],
+        )
+        motor_simulation = MotorSimulation(**encoded_attributes)
         return motor_simulation
 
     def get_motor_binary(self) -> bytes:

src/services/rocket.py

@@ -17,7 +17,8 @@
 from src.models.sub.aerosurfaces import NoseCone, Tail, Fins
 from src.services.motor import MotorService
 from src.views.rocket import RocketSimulation
-from src.utils import rocketpy_encoder, DiscretizeConfig
+from src.views.motor import MotorSimulation
+from src.utils import collect_attributes
 
 
 class RocketService:
@@ -107,10 +108,10 @@ def get_rocket_simulation(self) -> RocketSimulation:
         Returns:
             RocketSimulation
         """
-        attributes = rocketpy_encoder(
-            self.rocket, DiscretizeConfig.for_rocket()
+        encoded_attributes = collect_attributes(
+            self.rocket, [RocketSimulation, MotorSimulation]
         )
-        rocket_simulation = RocketSimulation(**attributes)
+        rocket_simulation = RocketSimulation(**encoded_attributes)
         return rocket_simulation
 
     def get_rocket_binary(self) -> bytes:

src/utils.py

@@ -2,16 +2,23 @@
 import io
 import logging
 import json
-import copy
 from datetime import datetime
-
 from typing import NoReturn, Tuple
 
-from rocketpy import Function
+import numpy as np
+from scipy.interpolate import interp1d
+
+from rocketpy import Function, Flight
 from rocketpy._encoders import RocketPyEncoder
+
 from starlette.datastructures import Headers, MutableHeaders
 from starlette.types import ASGIApp, Message, Receive, Scope, Send
 
+from src.views.environment import EnvironmentSimulation
+from src.views.flight import FlightSimulation
+from src.views.motor import MotorSimulation
+from src.views.rocket import RocketSimulation
+
 logger = logging.getLogger(__name__)
 
 
@@ -46,78 +53,154 @@ def for_flight(cls) -> 'DiscretizeConfig':
         return cls(bounds=(0, 30), samples=200)
 
 
-def rocketpy_encoder(obj, config: DiscretizeConfig = DiscretizeConfig()):
-    """
-    Encode a RocketPy object using official RocketPy encoders.
+class InfinityEncoder(RocketPyEncoder):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
 
-    This function creates a copy of the object, discretizes callable Function
-    attributes on the copy, and then uses RocketPy's official RocketPyEncoder for
-    complete object serialization. The original object remains unchanged.
+    def default(self, obj):
+        if (
+            isinstance(obj, Function)
+            and not callable(obj.source)
+            and obj.__dom_dim__ == 1
+        ):
+            size = len(obj._domain)
+            reduction_factor = 1
+            if size > 25:
+                reduction_factor = size // 25
+            if reduction_factor > 1:
+                obj = obj.set_discrete(
+                    obj.x_array[0],
+                    obj.x_array[-1],
+                    size // reduction_factor,
+                    mutate_self=False,
+                )
+        if isinstance(obj, Flight):
+            obj._Flight__evaluate_post_process
+            solution = np.array(obj.solution)
+            size = len(solution)
+            if size > 25:
+                reduction_factor = size // 25
+                reduced_solution = np.zeros(
+                    (size // reduction_factor, solution.shape[1])
+                )
+                reduced_scale = np.linspace(
+                    solution[0, 0], solution[-1, 0], size // reduction_factor
+                )
+                for i, col in enumerate(solution.T):
+                    reduced_solution[:, i] = interp1d(
+                        solution[:, 0], col, assume_sorted=True
+                    )(reduced_scale)
+                obj.solution = reduced_solution.tolist()
 
-    Args:
-        obj: RocketPy object (Environment, Motor, Rocket, Flight)
-        config: DiscretizeConfig object with discretization parameters (optional)
+            obj.flight_phases = None
+            obj.function_evaluations = None
 
-    Returns:
-        Dictionary of encoded attributes
-    """
+        return super().default(obj)
 
-    if config is None:
-        config = DiscretizeConfig()
-    try:
-        # Create a copy to avoid mutating the original object
-        obj_copy = copy.deepcopy(obj)
-    except Exception:
-        # Fall back to a shallow copy if deep copy is not supported
-        obj_copy = copy.copy(obj)
-
-    for attr_name in dir(obj_copy):
-        if attr_name.startswith('_'):
-            continue
 
-        try:
-            attr_value = getattr(obj_copy, attr_name)
-        except Exception:
-            continue
+def rocketpy_encoder(obj):
+    """
+    Encode a RocketPy object using official RocketPy encoders.
 
-        if callable(attr_value) and isinstance(attr_value, Function):
+    Uses InfinityEncoder for serialization and reduction.
+    """
+    json_str = json.dumps(
+        obj,
+        cls=InfinityEncoder,
+        include_outputs=True,
+        include_function_data=True,
+        discretize=True,
+        allow_pickle=False,
+    )
+    encoded_result = json.loads(json_str)
+    return _fix_datetime_fields(encoded_result)
+
+
+def collect_attributes(obj, attribute_classes=None):
+    """
+    Collect attributes from various simulation classes and populate them from the flight object.
+    """
+    if attribute_classes is None:
+        attribute_classes = []
+
+    attributes = rocketpy_encoder(obj)
+
+    for attribute_class in attribute_classes:
+        if issubclass(attribute_class, FlightSimulation):
+            flight_attributes_list = [
+                attr
+                for attr in attribute_class.__annotations__.keys()
+                if attr not in ["message", "rocket", "env"]
+            ]
             try:
-                discretized_func = Function(attr_value.source)
-                discretized_func.set_discrete(
-                    lower=config.bounds[0],
-                    upper=config.bounds[1],
-                    samples=config.samples,
-                    mutate_self=True,
-                )
-
-                setattr(obj_copy, attr_name, discretized_func)
-
-            except Exception as e:
-                logger.warning(f"Failed to discretize {attr_name}: {e}")
+                for key in flight_attributes_list:
+                    if key not in attributes:
+                        try:
+                            value = getattr(obj, key)
+                            attributes[key] = value
+                        except Exception:
+                            pass
+            except Exception:
+                pass
+
+        elif issubclass(attribute_class, RocketSimulation):
+            rocket_attributes_list = [
+                attr
+                for attr in attribute_class.__annotations__.keys()
+                if attr not in ["message", "motor"]
+            ]
+            try:
+                for key in rocket_attributes_list:
+                    if key not in attributes.get("rocket", {}):
+                        try:
+                            value = getattr(obj.rocket, key)
+                            attributes.setdefault("rocket", {})[key] = value
+                        except Exception:
+                            pass
+            except Exception:
+                pass
+
+        elif issubclass(attribute_class, MotorSimulation):
+            motor_attributes_list = [
+                attr
+                for attr in attribute_class.__annotations__.keys()
+                if attr not in ["message"]
+            ]
+            try:
+                for key in motor_attributes_list:
+                    if key not in attributes.get("rocket", {}).get(
+                        "motor", {}
+                    ):
+                        try:
+                            value = getattr(obj.rocket.motor, key)
+                            attributes.setdefault("rocket", {}).setdefault(
+                                "motor", {}
+                            )[key] = value
+                        except Exception:
+                            pass
+            except Exception:
+                pass
+
+        elif issubclass(attribute_class, EnvironmentSimulation):
+            environment_attributes_list = [
+                attr
+                for attr in attribute_class.__annotations__.keys()
+                if attr not in ["message"]
+            ]
+            try:
+                for key in environment_attributes_list:
+                    if key not in attributes.get("env", {}):
+                        try:
+                            value = getattr(obj.env, key)
+                            attributes.setdefault("env", {})[key] = value
+                        except Exception:
+                            pass
+            except Exception:
+                pass
+        else:
+            continue
 
-    try:
-        json_str = json.dumps(
-            obj_copy,
-            cls=RocketPyEncoder,
-            include_outputs=True,
-            include_function_data=True,
-        )
-        encoded_result = json.loads(json_str)
-
-        # Post-process to fix datetime fields that got converted to lists
-        return _fix_datetime_fields(encoded_result)
-    except Exception as e:
-        logger.warning(f"Failed to encode with RocketPyEncoder: {e}")
-        attributes = {}
-        for attr_name in dir(obj_copy):
-            if not attr_name.startswith('_'):
-                try:
-                    attr_value = getattr(obj_copy, attr_name)
-                    if not callable(attr_value):
-                        attributes[attr_name] = str(attr_value)
-                except Exception:
-                    continue
-        return attributes
+    return rocketpy_encoder(attributes)
 
 
 def _fix_datetime_fields(data):