Updated sensors

testing.py

@@ -1,10 +1,148 @@
-from spatz.sensors.antenna.tx_gain import GainPattern
-import math
+import pygame
+import serial
+import re
+import numpy as np
+import time
 
-pattern = GainPattern("data/gain_pattern/farfield_all.txt")
+from numpy import array
+from pyquaternion import Quaternion
+from math import cos, sin, pi
 
-# pattern.get_gain(41,66)
-# pattern.get_gain(40,100)
-# pattern.get_gain(10,180)
-# pattern.get_gain(0,95)
-# pattern.get_gain(21,100)
+from spatz.sensors.imu.wsen_isds import WSEN_ISDS_GYRO
+from spatz.simulation import Simulation
+from spatz.simulations.rocketpy import RocketPyCSV
+
+# Blatantly stolen from: https://stackoverflow.com/questions/21019471/how-can-i-draw-a-3d-shape-using-pygame-no-other-modules
+
+X, Y, Z = 0, 1, 2
+
+
+def rotation_matrix(a, b, by):
+    """
+    rotation matrix of a, b, by radians around x, y, z axes (respectively)
+    """
+    sa, ca = sin(a), cos(a)
+    sb, cb = sin(b), cos(b)
+    sby, cby = sin(by), cos(by)
+    return (
+        (cb*cby, -cb*sby, sb),
+        (ca*sby + sa*sb*cby, ca*cby - sby*sa*sb, -cb*sa),
+        (sby*sa - ca*sb*cby, ca*sby*sb + sa*cby, ca*cb)
+    )
+
+
+class Physical:
+    def __init__(self, vertices, edges, colors):
+        """
+        a 3D object that can rotate around the three axes
+        :param vertices: a tuple of points (each has 3 coordinates)
+        :param edges: a tuple of pairs (each pair is a set containing 2 vertices' indexes)
+        """
+        self.__vertices = array(vertices)
+        self.__edges = tuple(edges)
+        self.__colors = tuple(colors)
+        self.__rotation = Quaternion(axis=(1, 0, 0), angle=0) # radians around each axis
+
+    def rotate(self, quaternion):
+        self.__rotation = quaternion
+
+    @property
+    def lines(self):
+        location = array([self.__rotation.rotate(vertex) for vertex in self.__vertices])  # an index->location mapping
+        return (((location[v1], location[v2]), color) for (v1, v2), color in zip(self.__edges, self.__colors))
+    
+
+BLACK, RED, GREEN, BLUE = (0, 0, 0), (255, 128, 128), (128, 255, 128), (128, 128, 255)
+
+LIGHTRED, LIGHTGREEN, LIGHTBLUE = (128, 64, 64), (64, 128, 64), (64, 64, 128)
+
+
+class Paint:
+    def __init__(self, shape, shape2):
+        self.__shape = shape
+        self.__shape2 = shape2
+        self.__size = 900, 450
+        self.__clock = pygame.time.Clock()
+        self.__screen = pygame.display.set_mode(self.__size)
+
+    def __fit(self, vec):
+        """
+        ignore the z-element (creating a very cheap projection), and scale x, y to the coordinates of the screen
+        """
+        # notice that len(self.__size) is 2, hence zip(vec, self.__size) ignores the vector's last coordinate
+        return [round(70 * coordinate + frame / 2) for coordinate, frame in zip(vec, self.__size)]
+
+    def __draw_shape(self, thickness=4):
+        for (start, end), color in self.__shape.lines:
+            pygame.draw.line(self.__screen, color, self.__fit((start[0]-2, start[1], start[2])), self.__fit((end[0]-2, end[1], end[2])), thickness)
+
+        for (start, end), color in self.__shape2.lines:
+            pygame.draw.line(self.__screen, color, self.__fit((start[0]+2, start[1], start[2])), self.__fit((end[0]+2, end[1], end[2])), thickness)
+
+    def draw(self):
+        self.__screen.fill(BLACK)
+        self.__draw_shape()
+        pygame.display.flip()
+        self.__clock.tick(40)          
+
+
+def main():
+    from pygame import K_q, K_w, K_a, K_s, K_z, K_x
+
+    rotation = Quaternion(axis=[1, 0, 0], angle=pi/2)
+
+    axes = Physical(
+        vertices=((0, 0, 0), (0, 0, 2), (0, 2, 0), (2, 0, 0)),
+        edges=({0, 1}, {0, 2}, {0, 3}),
+        colors=(BLUE, GREEN, RED)
+    ) 
+
+    truth = Physical(
+        vertices=((0, 0, 0), (0, 0, 2), (0, 2, 0), (2, 0, 0)),
+        edges=({0, 1}, {0, 2}, {0, 3}),
+        colors=(LIGHTBLUE, LIGHTGREEN, LIGHTRED)
+    ) 
+        
+    pygame.init()
+    pygame.display.set_caption("Simulation")
+    renderer = Paint(axes, truth)
+
+    simulation = Simulation().load(RocketPyCSV('nominal_wind.csv'))
+
+    gyro = simulation.add_sensor(WSEN_ISDS_GYRO)
+    offset = gyro.calibrate(100)
+    att = simulation.add_observer([' e0', ' e1', ' e2', ' e3'])
+    dt = 0.01
+
+    quat = np.array([-0.706864,0.018510,0.018510,-0.706864])
+
+    while True:
+        time.sleep(0.01)
+        simulation.advance(dt)
+        omegas = gyro() - offset
+        true = att()
+
+        omegas = (omegas / 1000) * np.pi / 180
+
+        matrix = np.array([
+            [1, -dt/2*omegas[0], -dt/2*omegas[1], -dt/2*omegas[2]],
+            [dt/2*omegas[0], 1, dt/2*omegas[2], -dt/2*omegas[1]],
+            [dt/2*omegas[1], -dt/2*omegas[2], 1, dt/2*omegas[0]],
+            [dt/2*omegas[2], dt/2*omegas[1], -dt/2*omegas[0], 1]
+        ])
+
+        quat = matrix @ quat
+        quat /= np.linalg.norm(quat)
+
+
+        for event in pygame.event.get():
+            if event.type == pygame.QUIT:
+                exit()
+            
+        axes.rotate(Quaternion(x=quat[0], y=quat[1], z=quat[2], w=quat[3]))
+        truth.rotate(Quaternion(x=true[0], y=true[1], z=true[2], w=true[3]))
+
+        renderer.draw()
+
+if __name__ == '__main__':
+    main()