SPATZ/spatz/simulations/astos_source.py
2024-10-08 14:06:54 +02:00

87 lines
2.6 KiB
Python

import numpy as np
import math
from typing import Literal
from numpy.typing import NDArray
from spatz.simulations.csv_source import CSVSource
def T1(angle):
# return Rotation.from_euler('X', angle, degrees=False).as_matrix()
return np.array([
[1, 0, 0],
[0, math.cos(angle), math.sin(angle)],
[0, -math.sin(angle), math.cos(angle)],
])
def T2(angle):
# return Rotation.from_euler('Y', angle, degrees=False).as_matrix()
return np.array([
[math.cos(angle), 0, -math.sin(angle)],
[0, 1, 0],
[math.sin(angle), 0, math.cos(angle)]
])
def T3(angle):
# return Rotation.from_euler('Z', angle, degrees=False).as_matrix()
return np.array([
[math.cos(angle), math.sin(angle), 0],
[-math.sin(angle), math.cos(angle), 0],
[0, 0, 1]
])
class ASTOSSource(CSVSource):
def __init__(self, path: str, interpolation: Literal['linear'] = 'linear') -> None:
super().__init__(path, 'time', interpolation)
def get_position(self) -> NDArray:
return self.fetch_values(['px_L', 'py_L', 'pz_L'])
def get_velocity(self, frame: Literal['global', 'local']) -> NDArray:
vel = self.fetch_values(['vx_L', 'vy_L', 'vz_L'])
if frame == 'local':
return self.global_to_local() @ vel
return vel
def get_acceleration(self, frame: Literal['global', 'local']) -> NDArray:
if frame == 'global':
acc = self.fetch_values(['ax_L', 'ay_L', 'az_L']) + np.array([-self.fetch_value('gravity'), 0, 0])
else:
acc = self.fetch_values(['ax_B', 'ay_B', 'az_B']) + self.fetch_init_values(['gx_B', 'gy_B', 'gz_B'])
return acc
def get_attitude(self) -> NDArray:
raise NotImplementedError()
def local_to_global(self) -> NDArray:
return self.global_to_local().T
def global_to_local(self) -> NDArray:
# ASTOS local to body
rots = self.fetch_values(['pitch_l', 'yaw_l', 'roll_l'])
pitch_l, yaw_l, roll_l = rots[0], rots[1], rots[2]
local_to_body = T1(roll_l) @ T2(pitch_l - math.pi/2) @ T1(-yaw_l)
return local_to_body
def get_angular_velocity(self) -> NDArray:
return self.fetch_values(['OMEGA_X', 'OMEGA_Y', 'OMEGA_Z'])
def get_static_pressure(self) -> float:
return self.fetch_value('pressure')
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('altitude')