diff --git a/plot.py b/plot.py new file mode 100644 index 0000000..aea5ffa --- /dev/null +++ b/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() \ No newline at end of file diff --git a/rocket_twin/systems/rocket/rocket.py b/rocket_twin/systems/rocket/rocket.py index 25a964d..cec12c2 100644 --- a/rocket_twin/systems/rocket/rocket.py +++ b/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"]) diff --git a/rocket_twin/systems/station/station.py b/rocket_twin/systems/station/station.py index f40b86e..a3776f2 100644 --- a/rocket_twin/systems/station/station.py +++ b/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. + diff --git a/rocket_twin/tests/test_stage.py b/rocket_twin/tests/test_stage.py index 386e08b..f195d01 100644 --- a/rocket_twin/tests/test_stage.py +++ b/rocket_twin/tests/test_stage.py @@ -1,6 +1,9 @@ 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: @@ -8,3 +11,58 @@ 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)