mirror of
https://git.intern.spaceteamaachen.de/ALPAKA/SPATZ.git
synced 2025-06-10 01:55:59 +00:00
148 lines
4.7 KiB
Python
148 lines
4.7 KiB
Python
import pygame
|
|
import serial
|
|
import re
|
|
import numpy as np
|
|
import time
|
|
|
|
from numpy import array
|
|
from pyquaternion import Quaternion
|
|
from math import cos, sin, pi
|
|
|
|
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() |