Added multistage fueling and dropping

plot.py

@@ -0,0 +1,83 @@
+import numpy as np
+from rocket_twin.systems import Station
+from cosapp.drivers import RungeKutta
+from cosapp.recorders import DataFrameRecorder
+import matplotlib.pyplot as plt
+
+sys = Station('sys', n_stages=3)
+T = 20.
+
+init = {'g_tank.fuel.weight_p' : 20.,
+        'g_tank.fuel.w_out_max' : 1.,
+        'controller.w_temp' : 1.}
+
+values = {'rocket.stage_1.tank.geom.height' : 1.4,
+          'rocket.stage_2.tank.geom.height' : 1.2}
+
+includes = ["g_tank.weight_prop", "rocket.weight_prop_1", "rocket.weight_prop_2", "rocket.weight_prop_3", "rocket.geom.weight"]
+
+driver = sys.add_driver(RungeKutta('rk', order=4, dt = 1))
+driver.time_interval = (0, T)
+driver.set_scenario(init=init, values = values)
+driver.add_recorder(DataFrameRecorder(includes=includes), period=1.)
+
+sys.run_drivers()
+
+data = driver.recorder.export_data()
+data = data.drop(["Section", "Status", "Error code"], axis=1)
+
+time = np.asarray(data['time'])
+fuel1 = np.asarray(data["rocket.weight_prop_1"])
+fuel2 = np.asarray(data["rocket.weight_prop_2"])
+fuel3 = np.asarray(data["rocket.weight_prop_3"])
+fuelg = np.asarray(data["g_tank.weight_prop"])
+mass = np.asarray(data["rocket.geom.weight"])
+
+plt.plot(time, fuel1, label="Stage 1 fuel")
+plt.plot(time, fuel2, label="Stage 2 fuel")
+plt.plot(time, fuel3, label="Stage 3 fuel")
+plt.plot(time, fuelg, label="Ground fuel")
+plt.plot(time, mass, label="Rocket mass")
+plt.xlabel("Time (s)")
+plt.ylabel("Fuel mass (kg)")
+plt.title("Fuel mass over time")
+plt.legend()
+plt.show()
+
+sys.drivers.clear()
+
+init = {'rocket.stage_1.controller.w_temp' : 1.}
+
+values = {'rocket.stage_1.tank.fuel.w_out_max' : 1.,
+          'rocket.stage_2.tank.fuel.w_out_max' : 1.,
+          'rocket.stage_3.tank.fuel.w_out_max' : 1.}
+
+includes = ["g_tank.weight_prop", "rocket.weight_prop_1", "rocket.weight_prop_2", "rocket.weight_prop_3", "rocket.geom.weight"]
+
+driver = sys.add_driver(RungeKutta('rk', order=4, dt = 1))
+driver.time_interval = (T, 2*T)
+driver.set_scenario(init=init, values = values)
+driver.add_recorder(DataFrameRecorder(includes=includes), period=1.)
+
+sys.run_drivers()
+
+data = driver.recorder.export_data()
+data = data.drop(["Section", "Status", "Error code"], axis=1)
+
+time = np.asarray(data['time'])
+fuel1 = np.asarray(data["rocket.weight_prop_1"])
+fuel2 = np.asarray(data["rocket.weight_prop_2"])
+fuel3 = np.asarray(data["rocket.weight_prop_3"])
+fuelg = np.asarray(data["g_tank.weight_prop"])
+mass = np.asarray(data["rocket.geom.weight"])
+
+plt.plot(time, fuel1, label="Stage 1 fuel")
+plt.plot(time, fuel2, label="Stage 2 fuel")
+plt.plot(time, fuel3, label="Stage 3 fuel")
+plt.plot(time, fuelg, label="Ground fuel")
+plt.plot(time, mass, label="Rocket mass")
+plt.xlabel("Time (s)")
+plt.ylabel("Fuel mass (kg)")
+plt.title("Fuel mass over time")
+plt.legend()
+plt.show()

rocket_twin/systems/rocket/rocket.py

@@ -27,13 +27,15 @@ def setup(self, n_stages=1):
                 wings = True
             if i == n_stages:
                 nose = True
+
             self.add_child(
                 Stage(f"stage_{i}", nose=nose, wings=wings),
-                pulling=["w_in", "weight_max", "weight_prop"],
+                pulling={"w_in" : f'w_in_{i}', "weight_max" : f'weight_max_{i}', "weight_prop" : f"weight_prop_{i}"},
             )
             shapes[i - 1] = f"stage_{i}_s"
             properties[i - 1] = f"stage_{i}"
             forces[i - 1] = f"thrust_{i}"
+
 
         self.add_child(OCCGeometry("geom", shapes=shapes, properties=properties))
         self.add_child(Dynamics("dyn", forces=forces, weights=["weight_rocket"]), pulling=["a"])

rocket_twin/systems/station/station.py

@@ -1,4 +1,5 @@
 from cosapp.base import System
+from OCC.Core.GProp import GProp_GProps
 
 from rocket_twin.systems import ControllerCoSApp, Pipe, Rocket, Tank
 
@@ -14,14 +15,56 @@ class Station(System):
     """
 
     def setup(self, n_stages=1):
+
+        self.add_inward('n_stages', n_stages, desc="Number of stages")
+        self.add_outward_modevar('stage', 1, desc="Current stage")
+
         self.add_child(ControllerCoSApp("controller"))
         self.add_child(Tank("g_tank"))
         self.add_child(Pipe("pipe"))
         self.add_child(Rocket("rocket", n_stages=n_stages))
 
         self.connect(self.g_tank.outwards, self.pipe.inwards, {"w_out": "w_in"})
-        self.connect(self.pipe.outwards, self.rocket.inwards, {"w_out": "w_in"})
+        self.connect(self.pipe.outwards, self.rocket.inwards, {"w_out": "w_in_1"})
 
         self.connect(self.controller.outwards, self.g_tank.inwards, {"w": "w_command"})
 
         self.g_tank.geom.height = 2.0
+
+        self.add_event("stage_full", trigger="rocket.weight_prop_1 == rocket.weight_max_1")
+        self.add_event("stage_empty", trigger="rocket.weight_prop_1 == 0.")
+
+    def transition(self):
+
+        if self.stage_full.present:
+            if self.stage < self.n_stages:
+                self.stage += 1
+
+                self.pop_child('pipe')
+                self.add_child(Pipe('pipe'), execution_index=2)
+                #self.exec_order = ["controller", "g_tank", "pipe", "rocket"]
+
+                self.connect(self.g_tank.outwards, self.pipe.inwards, {"w_out": "w_in"})
+                self.connect(self.pipe.outwards, self.rocket.inwards, {"w_out": f"w_in_{self.stage}"})
+
+                self.rocket[f'w_in_{self.stage - 1}'] = 0.
+                self.stage_full.trigger = f"rocket.weight_prop_{self.stage} == rocket.weight_max_{self.stage}"
+            else:
+                self.controller.w_temp = 0.
+                self.rocket[f'w_in_{self.stage}'] = 0.
+                self.stage = 1
+
+        if self.stage_empty.present:
+            if self.stage < self.n_stages:
+                stage = self.rocket.pop_child(f"stage_{self.stage}")
+                self.rocket.add_child(stage, execution_index=self.stage - 1)
+                self.rocket[f'stage_{self.stage}'].controller.w_temp = 0.
+                self.rocket.geom[f'stage_{self.stage}'] = GProp_GProps()
+
+                self.stage += 1
+                self.stage_empty.trigger = f"rocket.weight_prop_{self.stage} == 0."
+                self.rocket.Takeoff.trigger = "dyn.a == -10000000"
+                self.rocket[f'stage_{self.stage}'].controller.w_temp = 1.
+            else:
+                self.rocket[f'stage_{self.stage}'].controller.w_temp = 0.
+

rocket_twin/tests/test_stage.py

@@ -1,10 +1,68 @@
 import numpy as np
 
-from rocket_twin.systems import Rocket
+from rocket_twin.systems import Rocket, Station
+
+from cosapp.drivers import RungeKutta, NonLinearSolver
+from cosapp.recorders import DataFrameRecorder
 
 
 class TestStage:
     def test_run_once(self):
         sys = Rocket("sys")
 
         sys.run_once()
+        print(sys.w_in_1)
+
+    def test_fuel(self):
+        sys = Station('sys', n_stages=3)
+
+        init = {'g_tank.fuel.weight_p' : 20.,
+                'g_tank.fuel.w_out_max' : 1.,
+                'controller.w_temp' : 1.}
+
+        includes = ["g_tank.weight_prop", "rocket.weight_prop_1", "rocket.weight_prop_2", "rocket.weight_prop_3"]
+
+        driver = sys.add_driver(RungeKutta('rk', order=4, dt = 1))
+        driver.time_interval = (0, 20)
+        driver.set_scenario(init=init)
+        driver.add_recorder(DataFrameRecorder(includes=includes), period=1.)
+
+        sys.run_drivers()
+
+        np.testing.assert_allclose(sys.rocket.weight_prop_1, 5., rtol=10**(-1))
+        np.testing.assert_allclose(sys.rocket.weight_prop_2, 5., rtol=10**(-1))
+        np.testing.assert_allclose(sys.rocket.weight_prop_3, 5., rtol=10**(-1))
+        np.testing.assert_allclose(sys.g_tank.weight_prop, 5., atol=10**(-2))
+
+    def test_flight(self):
+
+        sys = Station('sys', n_stages=3)
+
+        init = {'controller.w_temp' : 0.,
+                'rocket.stage_1.controller.w_temp' : 1.,
+                "rocket.stage_1.tank.fuel.w_out_max" : 1.,
+                "rocket.stage_1.tank.fuel.weight_p" : 5.,
+                "rocket.stage_2.tank.fuel.w_out_max" : 1.,
+                "rocket.stage_2.tank.fuel.weight_p" : 5.,
+                "rocket.stage_3.tank.fuel.w_out_max" : 1.,
+                "rocket.stage_3.tank.fuel.weight_p" : 5.}
+
+        includes = ["g_tank.weight_prop", "rocket.weight_prop_1", "rocket.weight_prop_2", "rocket.weight_prop_3"]
+
+        driver = sys.add_driver(RungeKutta('rk', order=4, dt = 1))
+        #solver = driver.add_child(NonLinearSolver('solver'))
+        driver.time_interval = (0, 20)
+        driver.set_scenario(init=init)
+        driver.add_recorder(DataFrameRecorder(includes=includes), period=1.)
+
+        sys.run_drivers()
+
+        data = driver.recorder.export_data()
+        data1 = data.drop(["Section", "Status", "Error code", "rocket.weight_prop_2", "rocket.weight_prop_3"], axis=1)
+        data2 = data.drop(["Section", "Status", "Error code", "g_tank.weight_prop", "rocket.weight_prop_1"], axis=1)
+        print(data1)
+        print('\n')
+        print(data2)
+        print(sys.rocket.geom.weight)
+
+        np.testing.assert_allclose(1, 2, atol=0.1)