ENH: Discretized and No-Pickle Encoding Options (#827)

* ENH: add an option to discretize callable sources encoding.
* TST: implement testing for discretized encoding.
* ENH: allow for disallowing pickle on encoding.
* MNT: Update CHANGELOG.
* MNT: change pickle encoding name to allow_pickle and test it.
* MNT: Tweak the discretization bounds.

---------

Co-authored-by: Mateus Stano Junqueira <69485049+MateusStano@users.noreply.github.com>
Co-authored-by: Gui-FernandesBR <63590233+Gui-FernandesBR@users.noreply.github.com>

Author: 3 people

CHANGELOG.md

@@ -32,7 +32,8 @@ Attention: The newest changes should be on top -->
 
 ### Added
 
--  ENH: Add the Coriolis Force to the Flight class [#799](https://github.com/RocketPy-Team/RocketPy/pull/799)
+- ENH: Discretized and No-Pickle Encoding Options [#827] (https://github.com/RocketPy-Team/RocketPy/pull/827)
+- ENH: Add the Coriolis Force to the Flight class [#799](https://github.com/RocketPy-Team/RocketPy/pull/799)
 
 ### Changed
 

rocketpy/_encoders.py

@@ -13,33 +13,46 @@
 
 class RocketPyEncoder(json.JSONEncoder):
     """Custom JSON encoder for RocketPy objects. It defines how to encode
-    different types of objects to a JSON supported format."""
+    different types of objects to a JSON supported format.
+    """
 
     def __init__(self, *args, **kwargs):
+        """Initializes the encoder with parameter options.
+
+        Parameters
+        ----------
+        *args : tuple
+            Positional arguments to pass to the parent class.
+        **kwargs : dict
+            Keyword arguments to configure the encoder. The following
+            options are available:
+            - include_outputs: bool, whether to include simulation outputs.
+              Default is False.
+            - include_function_data: bool, whether to include Function
+              data in the encoding. If False, Functions will be encoded by their
+              ``__repr__``. This is useful for reducing the size of the outputs,
+              but it prevents full restoration of the object upon decoding.
+              Default is True.
+            - discretize: bool, whether to discretize Functions whose source
+              are callables. If True, the accuracy of the decoding may be reduced.
+              Default is False.
+            - allow_pickle: bool, whether to pickle callable objects. If
+              False, callable sources (such as user-defined functions, parachute
+              triggers or simulation callable outputs) will have their name
+              stored instead of the function itself. This is useful for
+              reducing the size of the outputs, but it prevents full restoration
+              of the object upon decoding.
+              Default is True.
+        """
         self.include_outputs = kwargs.pop("include_outputs", False)
         self.include_function_data = kwargs.pop("include_function_data", True)
+        self.discretize = kwargs.pop("discretize", False)
+        self.allow_pickle = kwargs.pop("allow_pickle", True)
         super().__init__(*args, **kwargs)
 
     def default(self, o):
-        if isinstance(
-            o,
-            (
-                np.int_,
-                np.intc,
-                np.intp,
-                np.int8,
-                np.int16,
-                np.int32,
-                np.int64,
-                np.uint8,
-                np.uint16,
-                np.uint32,
-                np.uint64,
-            ),
-        ):
-            return int(o)
-        elif isinstance(o, (np.float16, np.float32, np.float64)):
-            return float(o)
+        if isinstance(o, np.generic):
+            return o.item()
         elif isinstance(o, np.ndarray):
             return o.tolist()
         elif isinstance(o, datetime):
@@ -50,11 +63,19 @@ def default(self, o):
             if not self.include_function_data:
                 return str(o)
             else:
-                encoding = o.to_dict(self.include_outputs)
+                encoding = o.to_dict(
+                    include_outputs=self.include_outputs,
+                    discretize=self.discretize,
+                    allow_pickle=self.allow_pickle,
+                )
                 encoding["signature"] = get_class_signature(o)
                 return encoding
         elif hasattr(o, "to_dict"):
-            encoding = o.to_dict(self.include_outputs)
+            encoding = o.to_dict(
+                include_outputs=self.include_outputs,
+                discretize=self.discretize,
+                allow_pickle=self.allow_pickle,
+            )
             encoding = remove_circular_references(encoding)
 
             encoding["signature"] = get_class_signature(o)

rocketpy/environment/environment.py

@@ -2630,7 +2630,21 @@ def decimal_degrees_to_arc_seconds(angle):
         arc_seconds = (remainder * 60 - arc_minutes) * 60
         return degrees, arc_minutes, arc_seconds
 
-    def to_dict(self, include_outputs=False):
+    def to_dict(self, **kwargs):
+        wind_velocity_x = self.wind_velocity_x
+        wind_velocity_y = self.wind_velocity_y
+        wind_heading = self.wind_heading
+        wind_direction = self.wind_direction
+        wind_speed = self.wind_speed
+        density = self.density
+        if kwargs.get("discretize", False):
+            wind_velocity_x = wind_velocity_x.set_discrete(0, self.max_expected_height)
+            wind_velocity_y = wind_velocity_y.set_discrete(0, self.max_expected_height)
+            wind_heading = wind_heading.set_discrete(0, self.max_expected_height)
+            wind_direction = wind_direction.set_discrete(0, self.max_expected_height)
+            wind_speed = wind_speed.set_discrete(0, self.max_expected_height)
+            density = density.set_discrete(0, self.max_expected_height)
+
         env_dict = {
             "gravity": self.gravity,
             "date": self.date,
@@ -2643,15 +2657,15 @@ def to_dict(self, include_outputs=False):
             "atmospheric_model_type": self.atmospheric_model_type,
             "pressure": self.pressure,
             "temperature": self.temperature,
-            "wind_velocity_x": self.wind_velocity_x,
-            "wind_velocity_y": self.wind_velocity_y,
-            "wind_heading": self.wind_heading,
-            "wind_direction": self.wind_direction,
-            "wind_speed": self.wind_speed,
+            "wind_velocity_x": wind_velocity_x,
+            "wind_velocity_y": wind_velocity_y,
+            "wind_heading": wind_heading,
+            "wind_direction": wind_direction,
+            "wind_speed": wind_speed,
         }
 
-        if include_outputs:
-            env_dict["density"] = self.density
+        if kwargs.get("include_outputs", False):
+            env_dict["density"] = density
             env_dict["barometric_height"] = self.barometric_height
             env_dict["speed_of_sound"] = self.speed_of_sound
             env_dict["dynamic_viscosity"] = self.dynamic_viscosity

rocketpy/mathutils/function.py

@@ -3565,7 +3565,7 @@ def __validate_extrapolation(self, extrapolation):
                 extrapolation = "natural"
         return extrapolation
 
-    def to_dict(self, include_outputs=False):  # pylint: disable=unused-argument
+    def to_dict(self, **kwargs):  # pylint: disable=unused-argument
         """Serializes the Function instance to a dictionary.
 
         Returns
@@ -3576,7 +3576,10 @@ def to_dict(self, include_outputs=False):  # pylint: disable=unused-argument
         source = self.source
 
         if callable(source):
-            source = to_hex_encode(source)
+            if kwargs.get("allow_pickle", True):
+                source = to_hex_encode(source)
+            else:
+                source = source.__name__
 
         return {
             "source": source,

rocketpy/mathutils/vector_matrix.py

@@ -403,7 +403,7 @@ def zeros():
         """Returns the zero vector."""
         return Vector([0, 0, 0])
 
-    def to_dict(self, include_outputs=False):  # pylint: disable=unused-argument
+    def to_dict(self, **kwargs):  # pylint: disable=unused-argument
         """Returns the vector as a JSON compatible element."""
         return list(self.components)
 
@@ -1007,7 +1007,7 @@ def __repr__(self):
             + f"       [{self.zx}, {self.zy}, {self.zz}])"
         )
 
-    def to_dict(self, include_outputs=False):  # pylint: disable=unused-argument
+    def to_dict(self, **kwargs):  # pylint: disable=unused-argument
         """Returns the matrix as a JSON compatible element."""
         return [list(row) for row in self.components]
 

rocketpy/motors/fluid.py

@@ -61,7 +61,7 @@ def __str__(self):
 
         return f"Fluid: {self.name}"
 
-    def to_dict(self, include_outputs=False):  # pylint: disable=unused-argument
+    def to_dict(self, **kwargs):  # pylint: disable=unused-argument
         return {"name": self.name, "density": self.density}
 
     @classmethod

rocketpy/motors/hybrid_motor.py

@@ -641,8 +641,8 @@ def draw(self, *, filename=None):
         """
         self.plots.draw(filename=filename)
 
-    def to_dict(self, include_outputs=False):
-        data = super().to_dict(include_outputs)
+    def to_dict(self, **kwargs):
+        data = super().to_dict(**kwargs)
         data.update(
             {
                 "grain_number": self.grain_number,
@@ -660,13 +660,18 @@ def to_dict(self, include_outputs=False):
             }
         )
 
-        if include_outputs:
+        if kwargs.get("include_outputs", False):
+            burn_rate = self.solid.burn_rate
+            if kwargs.get("discretize", False):
+                burn_rate = burn_rate.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
             data.update(
                 {
                     "grain_inner_radius": self.solid.grain_inner_radius,
                     "grain_height": self.solid.grain_height,
                     "burn_area": self.solid.burn_area,
-                    "burn_rate": self.solid.burn_rate,
+                    "burn_rate": burn_rate,
                 }
             )
 

rocketpy/motors/liquid_motor.py

@@ -497,8 +497,8 @@ def draw(self, *, filename=None):
         """
         self.plots.draw(filename=filename)
 
-    def to_dict(self, include_outputs=False):
-        data = super().to_dict(include_outputs)
+    def to_dict(self, **kwargs):
+        data = super().to_dict(**kwargs)
         data.update(
             {
                 "positioned_tanks": [