ENH: add an option to discretize callable sources encoding.

Author: phmbressan

rocketpy/_encoders.py

@@ -17,29 +17,13 @@ class RocketPyEncoder(json.JSONEncoder):
 
     def __init__(self, *args, **kwargs):
         self.include_outputs = kwargs.pop("include_outputs", False)
+        self.discretize = kwargs.pop("discretize", False)
         self.include_function_data = kwargs.pop("include_function_data", 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 +34,15 @@ 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
+                )
                 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
+            )
             encoding = remove_circular_references(encoding)
 
             encoding["signature"] = get_class_signature(o)

rocketpy/environment/environment.py

@@ -2702,7 +2702,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,
@@ -2715,15 +2729,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

@@ -3567,7 +3567,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

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": [

rocketpy/motors/motor.py

@@ -1231,14 +1231,7 @@ def get_attr_value(obj, attr_name, multiplier=1):
             # Write last line
             file.write(f"{self.thrust.source[-1, 0]:.4f} {0:.3f}\n")
 
-    def to_dict(self, include_outputs=False):
-        thrust_source = self.thrust_source
-
-        if isinstance(thrust_source, str):
-            thrust_source = self.thrust.source
-        elif callable(thrust_source) and not isinstance(thrust_source, Function):
-            thrust_source = Function(thrust_source)
-
+    def to_dict(self, **kwargs):
         data = {
             "thrust_source": self.thrust,
             "dry_I_11": self.dry_I_11,
@@ -1258,31 +1251,94 @@ def to_dict(self, include_outputs=False):
             "reference_pressure": self.reference_pressure,
         }
 
-        if include_outputs:
+        if kwargs.get("include_outputs", False):
+            total_mass = self.total_mass
+            propellant_mass = self.propellant_mass
+            mass_flow_rate = self.total_mass_flow_rate
+            center_of_mass = self.center_of_mass
+            center_of_propellant_mass = self.center_of_propellant_mass
+            exhaust_velocity = self.exhaust_velocity
+            I_11 = self.I_11
+            I_22 = self.I_22
+            I_33 = self.I_33
+            I_12 = self.I_12
+            I_13 = self.I_13
+            I_23 = self.I_23
+            propellant_I_11 = self.propellant_I_11
+            propellant_I_22 = self.propellant_I_22
+            propellant_I_33 = self.propellant_I_33
+            propellant_I_12 = self.propellant_I_12
+            propellant_I_13 = self.propellant_I_13
+            propellant_I_23 = self.propellant_I_23
+            if kwargs.get("discretize", False):
+                total_mass = total_mass.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
+                propellant_mass = propellant_mass.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
+                mass_flow_rate = mass_flow_rate.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
+                center_of_mass = center_of_mass.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
+                center_of_propellant_mass = (
+                    center_of_propellant_mass.set_discrete_based_on_model(
+                        self.thrust, mutate_self=False
+                    )
+                )
+                exhaust_velocity = exhaust_velocity.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
+                I_11 = I_11.set_discrete_based_on_model(self.thrust, mutate_self=False)
+                I_22 = I_22.set_discrete_based_on_model(self.thrust, mutate_self=False)
+                I_33 = I_33.set_discrete_based_on_model(self.thrust, mutate_self=False)
+                I_12 = I_12.set_discrete_based_on_model(self.thrust, mutate_self=False)
+                I_13 = I_13.set_discrete_based_on_model(self.thrust, mutate_self=False)
+                I_23 = I_23.set_discrete_based_on_model(self.thrust, mutate_self=False)
+                propellant_I_11 = propellant_I_11.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
+                propellant_I_22 = propellant_I_22.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
+                propellant_I_33 = propellant_I_33.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
+                propellant_I_12 = propellant_I_12.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
+                propellant_I_13 = propellant_I_13.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
+                propellant_I_23 = propellant_I_23.set_discrete_based_on_model(
+                    self.thrust, mutate_self=False
+                )
             data.update(
                 {
                     "vacuum_thrust": self.vacuum_thrust,
-                    "total_mass": self.total_mass,
-                    "propellant_mass": self.propellant_mass,
-                    "mass_flow_rate": self.mass_flow_rate,
-                    "center_of_mass": self.center_of_mass,
-                    "center_of_propellant_mass": self.center_of_propellant_mass,
+                    "total_mass": total_mass,
+                    "propellant_mass": propellant_mass,
+                    "mass_flow_rate": mass_flow_rate,
+                    "center_of_mass": center_of_mass,
+                    "center_of_propellant_mass": center_of_propellant_mass,
                     "total_impulse": self.total_impulse,
-                    "exhaust_velocity": self.exhaust_velocity,
+                    "exhaust_velocity": exhaust_velocity,
                     "propellant_initial_mass": self.propellant_initial_mass,
                     "structural_mass_ratio": self.structural_mass_ratio,
-                    "I_11": self.I_11,
-                    "I_22": self.I_22,
-                    "I_33": self.I_33,
-                    "I_12": self.I_12,
-                    "I_13": self.I_13,
-                    "I_23": self.I_23,
-                    "propellant_I_11": self.propellant_I_11,
-                    "propellant_I_22": self.propellant_I_22,
-                    "propellant_I_33": self.propellant_I_33,
-                    "propellant_I_12": self.propellant_I_12,
-                    "propellant_I_13": self.propellant_I_13,
-                    "propellant_I_23": self.propellant_I_23,
+                    "I_11": I_11,
+                    "I_22": I_22,
+                    "I_33": I_33,
+                    "I_12": I_12,
+                    "I_13": I_13,
+                    "I_23": I_23,
+                    "propellant_I_11": propellant_I_11,
+                    "propellant_I_22": propellant_I_22,
+                    "propellant_I_33": propellant_I_33,
+                    "propellant_I_12": propellant_I_12,
+                    "propellant_I_13": propellant_I_13,
+                    "propellant_I_23": propellant_I_23,
                 }
             )
 
@@ -1510,7 +1566,7 @@ def center_of_propellant_mass(self):
         Function
             Function representing the center of mass of the motor.
         """
-        return self.chamber_position
+        return Function(self.chamber_position).set_discrete_based_on_model(self.thrust)
 
     @funcify_method("Time (s)", "Inertia I_11 (kg m²)")
     def propellant_I_11(self):
@@ -1532,11 +1588,11 @@ def propellant_I_11(self):
         ----------
         https://en.wikipedia.org/wiki/Moment_of_inertia#Inertia_tensor
         """
-        return (
+        return Function(
             self.propellant_mass
             * (3 * self.chamber_radius**2 + self.chamber_height**2)
             / 12
-        )
+        ).set_discrete_based_on_model(self.thrust)
 
     @funcify_method("Time (s)", "Inertia I_22 (kg m²)")
     def propellant_I_22(self):
@@ -1580,19 +1636,21 @@ def propellant_I_33(self):
         ----------
         https://en.wikipedia.org/wiki/Moment_of_inertia#Inertia_tensor
         """
-        return self.propellant_mass * self.chamber_radius**2 / 2
+        return Function(
+            self.propellant_mass * self.chamber_radius**2 / 2
+        ).set_discrete_based_on_model(self.thrust)
 
     @funcify_method("Time (s)", "Inertia I_12 (kg m²)")
     def propellant_I_12(self):
-        return Function(0)
+        return Function(0).set_discrete_based_on_model(self.thrust)
 
     @funcify_method("Time (s)", "Inertia I_13 (kg m²)")
     def propellant_I_13(self):
-        return Function(0)
+        return Function(0).set_discrete_based_on_model(self.thrust)
 
     @funcify_method("Time (s)", "Inertia I_23 (kg m²)")
     def propellant_I_23(self):
-        return Function(0)
+        return Function(0).set_discrete_based_on_model(self.thrust)
 
     @staticmethod
     def load_from_eng_file(
@@ -1862,8 +1920,8 @@ def all_info(self):
         self.prints.all()
         self.plots.all()
 
-    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(
             {
                 "chamber_radius": self.chamber_radius,

rocketpy/motors/solid_motor.py

@@ -765,8 +765,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(
             {
                 "nozzle_radius": self.nozzle_radius,
@@ -781,13 +781,18 @@ def to_dict(self, include_outputs=False):
             }
         )
 
-        if include_outputs:
+        if kwargs.get("include_outputs", False):
+            burn_rate = self.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.grain_inner_radius,
                     "grain_height": self.grain_height,
                     "burn_area": self.burn_area,
-                    "burn_rate": self.burn_rate,
+                    "burn_rate": burn_rate,
                     "Kn": self.Kn,
                 }
             )