Major rework using rocketpy

spatz/simulation.py

@@ -2,6 +2,8 @@ from typing import List
 from numpy.random import normal
 from tqdm import tqdm
 
+from spatz.simulations.advanceable import Advanceable
+from spatz.simulations.data_source import DataSource
 from spatz.dataset import Dataset
 from spatz.logger import Logger
 from spatz.sensors import Sensor
@@ -35,9 +37,11 @@ class UniformTimeSteps:
         return self.__dt + noise
 
 
-class Simulation:
+class Simulation(Advanceable):
     def __init__(self, time_steps=UniformTimeSteps(0.01)):
-        self.__dataset = None
+        super().__init__()
+
+        self.__data_source = None
         self.__logger = None
         self.__sensors: List[Sensor] = []
         self.__time_steps = time_steps
@@ -46,26 +50,25 @@ class Simulation:
         idx = 0
 
         # Clear all logs and reset the dataset to the first time step.
-        self.__dataset.reset()
+        self.__data_source.reset()
         self.__logger.reset()
 
         if verbose:
-            pbar = tqdm(total=self.__dataset.get_length())
+            pbar = tqdm(total=self.__data_source.get_length())
 
         while True:
-            t = self.__dataset.get_time()
+            t = self.get_time()
             dt = self.__time_steps(t)
             t_ = t + dt
             idx += 1
 
-            if t_ > self.__dataset.get_length():
+            if t_ > self.__data_source.get_length():
                 break
 
-            if until is not None and self.__dataset.get_phase() == until:
+            if until is not None and self.__data_source.get_phase() == until:
                 break
 
-            self.__dataset.step(dt)
-            self.__logger.step(dt)
+            self.advance(dt)
 
             if verbose:
                 pbar.update(dt)
@@ -75,31 +78,26 @@ class Simulation:
         if verbose:
             pbar.close()
 
-    def advance_to(self, t: float):
-        t_old = self.__dataset.get_time()
-        self.__dataset.step(t - t_old)
-        self.__logger.step(t - t_old)
+    def _on_step(self, dt: float):
+        self.__data_source.advance(dt)
+        self.__logger.advance(dt)
 
-    def advance(self, dt: float):
-        self.__dataset.step(dt)
-        self.__logger.step(dt)
-
-    def get_dataset(self) -> Dataset:
-        return self.__dataset
-    
-    def get_logger(self) -> Logger:
-        return self.__logger
-
-    def load(self, path: str):
-        self.__dataset = Dataset(path)
+    def load(self, source: DataSource):
+        self.__data_source = source
         self.__logger = Logger()
         
         for sensor in self.__sensors:
-            sensor.set_dataset(self.__dataset)
+            sensor.set_dataset(self.__data_source)
             sensor.set_logger(self.__logger)
 
         return self
 
+    def get_data_source(self) -> DataSource:
+        return self.__data_source
+    
+    def get_logger(self) -> Logger:
+        return self.__logger
+
     def add_sensor(self, sensor, *args, **kwargs) -> Sensor:
         """Register a new sensor for this simulation. A registered sensor can be called like a function and returns
         the current measurements. The class' constructor arguments have to be given aswell.
@@ -112,7 +110,7 @@ class Simulation:
         """
         assert issubclass(sensor, Sensor), "Expected a subclass of Sensor."
 
-        self.__sensors.append(sensor(self.__dataset, self.__logger, *args, **kwargs))
+        self.__sensors.append(sensor(self.__data_source, self.__logger, *args, **kwargs))
 
         return self.__sensors[-1]
 
@@ -132,10 +130,10 @@ class Simulation:
             attributes = observer_or_attributes
             assert len(attributes) != 0, "Observed attributes list must be nonempty."
 
-            self.__sensors.append(Observer(self.__dataset, self.__logger, attributes))
+            self.__sensors.append(Observer(self.__data_source, self.__logger, attributes))
         else:
             observer = observer_or_attributes
-            self.__sensors.append(observer(self.__dataset, self.__logger))
+            self.__sensors.append(observer(self.__data_source, self.__logger))
 
         return self.__sensors[-1]