Major rework using rocketpy

spatz/simulations/advanceable.py

@@ -0,0 +1,47 @@
+from abc import abstractmethod
+
+
+class Advanceable:
+    def __init__(self) -> None:
+        self.reset()
+
+    def advance(self, dt: float):
+        """Advances the simulation data in time.
+
+        Args:
+            dt (float): The step in time to make.
+        """
+        self.__t += dt
+        self._on_step(dt)
+
+    def advance_to(self, t: float):
+        """Advances the simulation data to a new point in time.
+
+        Args:
+            t (float): The target point in time.
+        """
+        assert t > self.__t, 'Advanceable can only move forward in time.'
+
+        self.advance(t - self.__t)
+
+    def reset(self):
+        """
+        Reset the Avanceable object to its initial state.
+        """
+        self.__t = 0
+        self._on_reset()
+
+    @abstractmethod
+    def _on_step(self, dt: float):
+        pass
+
+    @abstractmethod
+    def _on_reset(self):
+        pass
+
+    def get_time(self) -> float:
+        """
+        Returns:
+            float: Returns the current time of the Advanceable.
+        """
+        return self.__t

spatz/simulations/csv_source.py

@@ -0,0 +1,91 @@
+import pandas as pd
+import numpy as np
+
+from typing import Literal, List
+from numpy.typing import NDArray
+
+from spatz.simulations.data_source import DataSource
+
+
+
+class CSVSource(DataSource):
+    def __init__(self, path: str, time_col: str, interpolation: Literal['linear']='linear') -> None:
+        """A data source that extracts all its data from a csv file.
+
+        Args:
+            time_col (str): The name of the column that contains time data.
+        """
+        super().__init__()
+
+        self._df = pd.read_csv(path)
+        self._time_col = time_col
+        self._idx = 0
+        self._interpolation = interpolation
+
+    def get_length(self) -> float:
+        return max(self._df[self._time_col])
+
+    def _on_reset(self):
+        pass
+
+    def _get_closest_idx(self, t: float) -> int:
+        """Gets an index _idx_ for the dataframe _df_ such that the values at the given time _t_ are somewhere between 
+        _idx_ and _idx+1_. 
+
+        Args:
+            t (float): The requested time.
+
+        Returns:
+            int: The computed index.
+        """
+        idx = (self._df[self._time_col] - t).abs().idxmin()
+        idx = idx if self._df[self._time_col].loc[idx] <= t else idx - 1
+
+        return idx
+
+    def _on_step(self, _: float):
+        self._idx = self._get_closest_idx(self.get_time())
+
+    def fetch_value(self, name: str, t: float | None = None, custom_interpolation=None) -> float:
+        """Get a specific value from the dataframe.
+
+        Args:
+            name (str): The name of the value to fetch.
+            t (float): Allows specification of a different time instead of the current time. None for current time.
+
+        Returns:
+            float: Returns the requested value.
+        """
+        idx = self._idx if t is None else self._get_closest_idx(t)
+
+        if self._interpolation == 'linear':
+            t_min = self._df.at[idx, self._time_col]
+            t_max = self._df.at[idx + 1, self._time_col]
+
+            # Sometimes no time passes in-between two samples.
+            if t_max == t_min:
+                return self._df.at[name, idx]
+            
+            # Compute the weight for interpolation.
+            alpha = (self.get_time() - t_min) / (t_max - t_min)
+
+            if custom_interpolation is not None:
+                a = self._df.at[idx, name]
+                b = self._df.at[idx + 1, name]
+
+                return custom_interpolation(a, b, alpha)
+
+            # Interpolate linearly between the two data points.
+            return (1 - alpha) * self._df.at[idx, name] + alpha * self._df.at[idx + 1, name]
+
+    def fetch_values(self, names: List[str], t: float | None = None, custom_interpolation=None) -> NDArray:
+        """Get specific values from the dataframe.
+
+        Args:
+            names (List[str]): Names of the values to get.
+            t (float): Allows specification of a different time instead of the current time. None for current time.
+
+        Returns:
+            np.array: Returns a numpy array containing the requested values in the same order as in the input list.
+        """
+        return np.asarray([self.fetch_value(name, t, custom_interpolation) for name in names])

spatz/simulations/data_source.py

@@ -0,0 +1,74 @@
+import numpy as np
+
+from spatz.simulations.advanceable import Advanceable
+
+from numpy.typing import NDArray
+from abc import abstractmethod
+from typing import Literal
+
+from ambiance import Atmosphere
+
+
+class DataSource(Advanceable):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def get_speed_of_sound(self) -> float:
+        return Atmosphere(self.get_altitude()).speed_of_sound
+    
+    def get_mach_number(self) -> float:
+        speed = np.linalg.norm(self.get_velocity('global'))
+
+        return speed / self.get_speed_of_sound()
+
+    @abstractmethod
+    def get_length(self) -> float:
+        pass
+
+    @abstractmethod
+    def get_position(self) -> NDArray:
+        pass
+
+    @abstractmethod
+    def get_velocity(self, frame: Literal['global', 'local']) -> NDArray:
+        pass
+
+    @abstractmethod
+    def get_acceleration(self, frame: Literal['global', 'local']) -> NDArray:
+        pass
+
+    @abstractmethod
+    def get_attitude(self) -> NDArray:
+        pass
+
+    @abstractmethod
+    def local_to_global(self) -> NDArray:
+        pass
+
+    @abstractmethod
+    def global_to_local(self) -> NDArray:
+        pass
+
+    @abstractmethod
+    def get_angular_velocity(self) -> NDArray:
+        pass
+
+    @abstractmethod
+    def get_static_pressure(self) -> float:
+        pass
+
+    @abstractmethod
+    def get_temperature(self) -> float:
+        pass
+
+    @abstractmethod
+    def get_longitude(self) -> float:
+        pass
+
+    @abstractmethod
+    def get_latitude(self) -> float:
+        pass
+
+    @abstractmethod
+    def get_altitude(self) -> float:
+        pass

spatz/simulations/rocketpy.py

@@ -0,0 +1,87 @@
+import numpy as np
+
+from typing import Literal
+from numpy.typing import NDArray
+
+from spatz.simulations.csv_source import CSVSource
+
+
+class RocketPyCSV(CSVSource):
+    def __init__(self, path: str, interpolation: Literal['linear'] = 'linear') -> None:
+        super().__init__(path, ' Time (s)', interpolation)
+
+    def get_position(self) -> NDArray:
+        return self.fetch_values([' X (m)', ' Y (m)', ' Z (m)'])
+
+    def get_velocity(self, frame: Literal['global', 'local']) -> NDArray:
+        vel_global = self.fetch_values([' Vx (m/s)', ' Vy (m/s)', ' Vz (m/s)'])
+
+        if frame == 'global':
+            return vel_global
+        
+        return self.global_to_local() @ vel_global
+
+    def get_acceleration(self, frame: Literal['global', 'local']) -> NDArray:
+        acc_global = self.fetch_values([' Ax (m/s²)', ' Ay (m/s²)', ' Az (m/s²)'])
+
+        if frame == 'global':
+            return acc_global
+        
+        return self.global_to_local() @ acc_global
+
+    def get_attitude(self) -> NDArray:
+        t_min = self._df.at[self._idx, self._time_col]
+        t_max = self._df.at[self._idx + 1, self._time_col]
+
+        def slerp(a, b, alpha):
+            theta = np.arccos(np.clip(np.dot(a, b), -1, 1))
+
+            if np.isclose(theta, 0):
+                return a
+                
+            return (a * np.sin((1-alpha) * theta) + b * np.sin(alpha * theta)) / np.sin(theta)
+        
+        qa = np.array([self._df.at[self._idx, ' e0'], self._df.at[self._idx, ' e1'], self._df.at[self._idx, ' e2'], self._df.at[self._idx, ' e3']])
+        qb = np.array([self._df.at[self._idx+1, ' e0'], self._df.at[self._idx+1, ' e1'], self._df.at[self._idx+1, ' e2'], self._df.at[self._idx+1, ' e3']])
+
+        alpha = (self.get_time() - t_min) / (t_max - t_min)
+
+        return slerp(qa, qb, alpha)
+
+    def global_to_local(self) -> NDArray:
+        quat = self.get_attitude()
+
+        e0, e1, e2, e3 = quat[0], quat[1], quat[2], quat[3]
+
+        # Taken from:
+        # https://docs.rocketpy.org/en/latest/technical/equations_of_motion.html
+        mat = np.array([
+            [e0**2 + e1**2 - e2**2 - e3**2, 2*(e1*e2+e0*e3), 2*(e1*e3 - e0*e2)], 
+            [2*(e1*e2 - e0*e3), e0**2 - e1**2 + e2**2 - e3**2, 2*(e2*e3 + e0*e1)], 
+            [2*(e1*e3 + e0*e2), 2*(e2*e3 - e0*e1), e0**2 - e1**2 - e2**2 + e3**2]
+        ])
+        
+        return mat
+
+    def local_to_global(self) -> NDArray:
+        return self.global_to_local().T
+
+    def get_angular_velocity(self) -> NDArray:
+        return self.fetch_values([' ω1 (rad/s)',' ω2 (rad/s)',' ω3 (rad/s)'])
+
+    def get_static_pressure(self) -> float:
+        return self.fetch_value(' Pressure (Pa)')
+
+    def get_temperature(self) -> float:
+        raise NotImplementedError()
+
+    def get_longitude(self) -> float:
+        return self.fetch_value(' Longitude (°)')
+
+    def get_latitude(self) -> float:
+        return self.fetch_value(' Latitude (°)')
+
+    def get_altitude(self) -> float:
+        return self.fetch_value(' Z (m)')
+
+