ENH: added .rpy file functionality (see issue 668)

This commit add 'save_to_rpy' and 'load_from_rpy' functions, that allows saving and loading flights.

Author: ArthurJWH

rocketpy/_encoders.py

@@ -7,6 +7,8 @@
 import numpy as np
 
 from rocketpy.mathutils.function import Function
+from rocketpy.prints.flight_prints import _FlightPrints
+from rocketpy.plots.flight_plots import _FlightPlots
 
 
 class RocketPyEncoder(json.JSONEncoder):
@@ -75,6 +77,7 @@ class RocketPyDecoder(json.JSONDecoder):
     different types of objects from a JSON supported format."""
 
     def __init__(self, *args, **kwargs):
+        self.resimulate = kwargs.pop("resimulate", False)
         super().__init__(object_hook=self.object_hook, *args, **kwargs)
 
     def object_hook(self, obj):
@@ -84,7 +87,50 @@ def object_hook(self, obj):
             try:
                 class_ = get_class_from_signature(signature)
 
-                if hasattr(class_, "from_dict"):
+                if class_.__name__ == "Flight" and not self.resimulate:
+                    new_flight = class_.__new__(class_)
+                    new_flight.prints = _FlightPrints(new_flight)
+                    new_flight.plots = _FlightPlots(new_flight)
+                    new_flight.rocket = obj["rocket"]
+                    new_flight.env = obj["env"]
+                    new_flight.rail_length = obj["rail_length"]
+                    new_flight.inclination = obj["inclination"]
+                    new_flight.heading = obj["heading"]
+                    new_flight.terminate_on_apogee = obj["terminate_on_apogee"]
+                    new_flight.max_time = obj["max_time"]
+                    new_flight.max_time_step = obj["max_time_step"]
+                    new_flight.min_time_step = obj["min_time_step"]
+                    new_flight.rtol = obj["rtol"]
+                    new_flight.atol = obj["atol"]
+                    new_flight.time_overshoot = obj["time_overshoot"]
+                    new_flight.name = obj["name"]
+                    new_flight.solution = obj["solution"]
+                    new_flight.out_of_rail_time = obj["out_of_rail_time"]
+                    new_flight.apogee_time = obj["apogee_time"]
+                    new_flight.apogee = obj["apogee"]
+                    new_flight.parachute_events = obj["parachute_events"]
+                    new_flight.impact_state = obj["impact_state"]
+                    new_flight.impact_velocity = obj["impact_velocity"]
+                    new_flight.x_impact = obj["x_impact"]
+                    new_flight.y_impact = obj["y_impact"]
+                    new_flight.t_final = obj["t_final"]
+                    new_flight.flight_phases = obj["flight_phases"]
+                    new_flight.ax = obj["ax"]
+                    new_flight.ay = obj["ay"]
+                    new_flight.az = obj["az"]
+                    new_flight.out_of_rail_time_index = obj["out_of_rail_time_index"]
+                    new_flight.function_evaluations = obj["function_evaluations"]
+                    new_flight.alpha1 = obj["alpha1"]
+                    new_flight.alpha2 = obj["alpha2"]
+                    new_flight.alpha3 = obj["alpha3"]
+                    new_flight.R1 = obj["R1"]
+                    new_flight.R2 = obj["R2"]
+                    new_flight.R3 = obj["R3"]
+                    new_flight.M1= obj["M1"]
+                    new_flight.M2 = obj["M2"]
+                    new_flight.M3 = obj["M3"]
+                    return new_flight
+                elif hasattr(class_, "from_dict"):
                     return class_.from_dict(obj)
                 else:
                     # Filter keyword arguments
@@ -118,7 +164,7 @@ def get_class_signature(obj):
         Signature of the class.
     """
     class_ = obj.__class__
-    name = getattr(class_, '__qualname__', class_.__name__)
+    name = getattr(class_, "__qualname__", class_.__name__)
 
     return {"module": class_.__module__, "name": name}
 

rocketpy/simulation/flight.py

@@ -3433,24 +3433,42 @@ def to_dict(self, include_outputs=False):
             "time_overshoot": self.time_overshoot,
             "name": self.name,
             "equations_of_motion": self.equations_of_motion,
+            "solution": self.solution, ##
+            "out_of_rail_time": self.out_of_rail_time, #
+            "apogee_time": self.apogee_time, #
+            "apogee": self.apogee, #
+            "parachute_events": self.parachute_events, ##
+            "impact_state": self.impact_state, #
+            "impact_velocity": self.impact_velocity, #
+            "x_impact": self.x_impact, #
+            "y_impact": self.y_impact, #
+            "t_final": self.t_final, ##
+            "flight_phases": self.flight_phases, ##
+            "ax": self.ax, #
+            "ay": self.ay, #
+            "az": self.az, #
+            "out_of_rail_time_index": self.out_of_rail_time_index, ##
+            "function_evaluations": self.function_evaluations, ##
+            "alpha1": self.alpha1, #
+            "alpha2": self.alpha2, #
+            "alpha3": self.alpha3, #
+            "R1": self.R1, ##
+            "R2": self.R2, ##
+            "R3": self.R3, ##
+            "M1": self.M1, ##
+            "M2": self.M2, ##
+            "M3": self.M3, ##
         }
 
         if include_outputs:
             data.update(
                 {
                     "time": self.time,
-                    "out_of_rail_time": self.out_of_rail_time,
                     "out_of_rail_velocity": self.out_of_rail_velocity,
                     "out_of_rail_state": self.out_of_rail_state,
-                    "apogee": self.apogee,
-                    "apogee_time": self.apogee_time,
                     "apogee_x": self.apogee_x,
                     "apogee_y": self.apogee_y,
                     "apogee_state": self.apogee_state,
-                    "x_impact": self.x_impact,
-                    "y_impact": self.y_impact,
-                    "impact_velocity": self.impact_velocity,
-                    "impact_state": self.impact_state,
                     "x": self.x,
                     "y": self.y,
                     "z": self.z,
@@ -3464,12 +3482,6 @@ def to_dict(self, include_outputs=False):
                     "w1": self.w1,
                     "w2": self.w2,
                     "w3": self.w3,
-                    "ax": self.ax,
-                    "ay": self.ay,
-                    "az": self.az,
-                    "alpha1": self.alpha1,
-                    "alpha2": self.alpha2,
-                    "alpha3": self.alpha3,
                     "altitude": self.altitude,
                     "mach_number": self.mach_number,
                     "stream_velocity_x": self.stream_velocity_x,

rocketpy/utilities.py

@@ -2,7 +2,11 @@
 import inspect
 import traceback
 import warnings
+import json
+import os
 
+from importlib.metadata import version
+from datetime import date
 import matplotlib.pyplot as plt
 import numpy as np
 from scipy.integrate import solve_ivp
@@ -12,6 +16,7 @@
 from .plots.plot_helpers import show_or_save_plot
 from .rocket.aero_surface import TrapezoidalFins
 from .simulation.flight import Flight
+from ._encoders import RocketPyEncoder, RocketPyDecoder
 
 
 def compute_cd_s_from_drop_test(
@@ -685,3 +690,63 @@ def get_instance_attributes(instance):
         if not inspect.ismethod(member[1]) and not member[0].startswith("__"):
             attributes_dict[member[0]] = member[1]
     return attributes_dict
+
+
+def save_to_rpy(flight: Flight, filename: str, include_output = False):
+    """Saves a .rpy file into the given path, containing key simulation informations to reproduce the results.
+
+    Parameters
+    ----------
+    flight : rocketpy.Flight
+        Flight object containing the rocket's flight data
+    filename : str
+        Path where the file will be saved in
+    include_output : bool, optional
+        If True, the function will include extra outputs into the file, by default False
+        
+    Returns
+    -------
+    None
+    """
+    file = os.path.splitext(filename)[0] + ".rpy"
+    with open(file, "w") as f:
+        data = {"date": str(date.today()), "version": version("rocketpy")}
+        print(flight)
+        data["simulation"] = flight
+        print(data)
+        print(flight.rocket)
+        json.dump(
+            data,
+            f,
+            cls=RocketPyEncoder,
+            indent=2,
+            include_outputs=include_output,
+        )
+
+
+def load_from_rpy(filename: str, resimulate = False):
+    """Loads the saved data from the .rpy into a Flight object.
+
+    Parameters
+    ----------
+    filename : str
+        Path where the file to be loaded is
+    resimulate : bool, optional
+        If True, the function will resimulate the object Flight, by default False
+        
+    Returns
+    -------
+    rocketpy.Flight
+        Flight object containing simulation information from the .rpy file
+    """
+    ext = os.path.splitext(os.path.basename(filename))[1]
+    if ext == ".rpy":
+        with open(filename, "r") as f:
+            data = json.load(f)
+            if data["version"] < version("rocketpy"):
+                warnings.warn("The file was saved in an older version of RocketPy than the one in the current environment")
+            simulation = json.dumps(data["simulation"])
+            flight = json.loads(simulation, cls=RocketPyDecoder, resimulate=resimulate)
+        return flight
+    else:
+        raise ValueError(f"Invalid file extension: {ext}. Allowed: .rpy")