Two stages

.pre-commit-config.yaml

@@ -47,6 +47,6 @@ repos:
           - flake8-bugbear==22.1.11
           - flake8-isort==4.1.1
   - repo: https://github.com/kynan/nbstripout
-    rev: 0.5.0
+    rev: 0.6.1
     hooks:
       - id: nbstripout

rocket_twin/systems/__init__.py

@@ -4,13 +4,14 @@
 from rocket_twin.systems.physics import Dynamics
 from rocket_twin.systems.tank import Pipe, Tank
 
-from rocket_twin.systems.rocket import Rocket  # isort: skip
+from rocket_twin.systems.rocket import Stage, Rocket  # isort: skip
 from rocket_twin.systems.station import Station  # isort: skip
 
 __all__ = [
     "Clock",
     "Engine",
     "Tank",
+    "Stage",
     "Rocket",
     "Pipe",
     "Dynamics",

rocket_twin/systems/rocket/__init__.py

@@ -1,3 +1,4 @@
 from rocket_twin.systems.rocket.rocket import Rocket
+from rocket_twin.systems.rocket.stage import Stage
 
-__all__ = ["Rocket"]
+__all__ = ["Stage", "Rocket"]

rocket_twin/systems/rocket/stage.py

@@ -0,0 +1,39 @@
+from cosapp.base import System
+
+from rocket_twin.systems import ControllerCoSApp, Engine, Tank
+
+
+class Stage(System):
+    """Model of a rocket stage.
+
+    Inputs
+    ------
+
+    Outputs
+    ------
+    force [N]: float,
+        thrust force
+    weight [kg]: float,
+        weight
+    cg [m]: float,
+        center of gravity
+    """
+
+    def setup(self):
+
+        self.add_child(ControllerCoSApp("controller"))
+        self.add_child(Tank("tank"), pulling=["w_in", "weight_max", "weight_p"])
+        self.add_child(Engine("engine"))
+
+        self.connect(self.controller.outwards, self.tank.inwards, {"w": "w_command"})
+        self.connect(self.tank.outwards, self.engine.inwards, {"w_out": "w_out"})
+
+        self.add_outward("weight", 1.0, desc="Weight", unit="kg")
+        self.add_outward("cg", 1.0, desc="Center of gravity", unit="m")
+
+    def compute(self):
+
+        self.weight = self.tank.weight + self.engine.weight
+        self.cg = (self.tank.cg * self.tank.weight + self.engine.cg * self.engine.weight) / (
+            self.weight
+        )

rocket_twin/tests/test_stages.py

@@ -0,0 +1,21 @@
+import numpy as np
+from cosapp.drivers import RungeKutta
+
+from rocket_twin.systems import Stage
+
+
+class TestStage:
+    def test_single_stage(self):
+
+        sys = Stage("sys")
+        init = {"tank.weight_p": "tank.weight_max"}
+        values = {"controller.w_temp": 1.0, "tank.w_out_max": 1.0}
+        stop = "tank.weight_p == 0."
+
+        driver = sys.add_driver(RungeKutta(order=4, dt=0.1))
+        driver.time_interval = (0, 10)
+        driver.set_scenario(init=init, values=values, stop=stop)
+        sys.run_drivers()
+
+        np.testing.assert_allclose(sys.weight, 2.0, atol=10 ** (-2))
+        np.testing.assert_allclose(sys.cg, 1.0, atol=10 ** (-2))