SPATZ migration + proper directory structure

spatz/simulation.py

@@ -0,0 +1,99 @@
+from typing import List
+from numpy.random import normal
+from tqdm import tqdm
+
+from spatz.dataset import Dataset
+from spatz.logger import Logger
+from spatz.sensors import Sensor
+from dataset import Dataset
+from logger import Logger
+from sensor import Sensor
+
+
+class UniformTimeSteps:
+    def __init__(self, dt: float, mu: float = 0, sigma: float = 0, delay_only=True) -> None:
+        """_summary_
+
+        Args:
+            dt (float): _description_
+            mu (float, optional): _description_. Defaults to 0.
+            sigma (float, optional): _description_. Defaults to 0.
+            delay_only (bool, optional): _description_. Defaults to True.
+        """
+        self.__dt = dt
+        self.__mu = mu
+        self.__sigma = sigma
+        self.__delay_only = delay_only
+
+    def __call__(self, t):
+        noise = normal(self.__mu, self.__sigma)
+
+        if self.__delay_only:
+            noise = abs(noise)
+
+        return self.__dt + noise
+
+
+
+class Simulation:
+    def __init__(self, time_steps=UniformTimeSteps(0.01)):
+        self.__dataset = None
+        self.__logger = None
+        self.__sensors: List[Sensor] = []
+        self.__time_steps = time_steps
+
+    def run(self, verbose=False):
+        idx = 0
+
+        # Clear all logs and reset the dataset to the first time step.
+        self.__dataset.reset()
+        self.__logger.reset()
+
+        if verbose:
+            pbar = tqdm(total=self.__dataset.get_length())
+
+        while True:
+            t = self.__dataset.get_time()
+            dt = self.__time_steps(t)
+            t_ = t + dt
+            idx += 1
+
+            if t_ > self.__dataset.get_length():
+                break
+
+            self.__dataset.step(dt)
+            self.__logger.step(dt)
+
+            if verbose:
+                pbar.update(dt)
+
+            yield idx, t_, t_ - t
+
+        if verbose:
+            pbar.close()
+
+    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)
+        self.__logger = Logger()
+        
+        for sensor in self.__sensors:
+            sensor.set_dataset(self.__dataset)
+            sensor.set_logger(self.__logger)
+
+        return self
+
+    def add_sensor(self, sensor, *args, **kwargs) -> Sensor:
+        assert issubclass(sensor, Sensor), "Expected a subclass of Sensor."
+
+        self.__sensors.append(sensor(self.__dataset, self.__logger, *args, **kwargs))
+
+        return self.__sensors[-1]
+        
+
+