diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e9a64e4 Binary files /dev/null and b/requirements.txt differ diff --git a/src/sim/pendulum.py b/src/sim/pendulum.py index ad15574..deab903 100644 --- a/src/sim/pendulum.py +++ b/src/sim/pendulum.py @@ -2,6 +2,7 @@ from pygame.math import Vector2 import math import numpy as np import random +import matplotlib.pyplot as plt # Constants C_GRAVITY = 9.81 # m/s^2 @@ -73,14 +74,14 @@ class Pendulum: # Integrate acceleration to get velocity self.v_ang.append( - self.a_ang[self.index + 1] * (dt / 1000) - + self.v_ang[self.index] # Previous velocity + self.v_ang[self.index] # Previous velocity + + (self.a_ang[self.index + 1] * (dt / 1000)) ) # Angular displacement self.theta.append( - self.v_ang[self.index + 1] * (dt / 1000) - + self.theta[self.index] # Previous angle + self.theta[self.index] # Previous angle + + (self.v_ang[self.index + 1] * (dt / 1000)) ) # Limit fall of pendulum @@ -107,14 +108,14 @@ class Pendulum: # Integrate acceleration to get velocity self.v_cart.append( - self.a_cart[self.index + 1] * (dt / 1000) - + self.v_cart[self.index] # Previous velocity + self.v_cart[self.index] # Previous velocity + + (self.a_cart[self.index + 1] * (dt / 1000)) ) # Cart displacement self.s_cart.append( - self.v_cart[self.index + 1] * (dt / 1000) - + self.s_cart[self.index] # Previous displacement + self.s_cart[self.index] # Previous displacement + + (self.v_cart[self.index + 1] * (dt / 1000)) ) self.dx = self.s_cart[self.index + 1] * C_MTPRATIO # Convert to pixels @@ -138,6 +139,29 @@ class Pendulum: self.fallen = False self.update(0) + def plot(self): + fig, axs = plt.subplots(2, 2) + fig.suptitle("Pendulum") + + axs[0,0].plot(self.theta) + axs[0,0].set_title('Angle [rad]') + axs[0,1].plot(self.v_ang) + axs[0,1].set_title('Angular velocity [rad/s]') + axs[1,0].plot(self.a_ang) + axs[1,0].set_title('Angular acceleration [rad/s^2]') + + fig, axs = plt.subplots(2, 2) + fig.suptitle("Cart") + + axs[0,0].plot(self.s_cart) + axs[0,0].set_title('Position [m]') + axs[0,1].plot(self.v_ang) + axs[0,1].set_title('Speed [m/s]') + axs[1,0].plot(self.a_ang) + axs[1,0].set_title('Acceleration [m/s^2]') + + plt.show() + # def update(self, dt): # """ # Update the pendulum's state based on the elapsed time. diff --git a/src/sim/sim.py b/src/sim/sim.py index 8c1b41c..7edf2b8 100644 --- a/src/sim/sim.py +++ b/src/sim/sim.py @@ -1,10 +1,11 @@ +# Pendulum simulator 4000 +# Arne van Iterson, 2023 + +# Imports import pygame from pygame.math import Vector2 import math -from pendulum import Pendulum -from uiHelpers import gridDark, gridLight - # pygame setup pygame.init() screen = pygame.display.set_mode((1280, 720)) @@ -13,71 +14,78 @@ running = True update = True pole = Vector2(screen.get_rect().center) # center of screen -# Text setup -font_g = pygame.font.SysFont(None, 128) -font_h = pygame.font.SysFont(None, 28) -font_m = pygame.font.SysFont(None, 16) +# Own objects must be imported after pygame init +from pendulum import Pendulum +from uiHelpers import * -# Metadata plotter -plot_y = 50 -def plotMeta(val, desc): - global plot_y - screen.blit(font_m.render(f"{desc} = {val}", True, "black"), (15, plot_y)) - plot_y += 15 - -def plotGrid(dist, Xoff=0, Yoff=0): - screen.fill("white") - cXoff = pole.x % dist - cYoff = pole.y % dist - - for i in range(0, 1280, dist): - pygame.draw.line(screen, gridLight, (i + Xoff + cXoff, 0), (i + Xoff + cXoff, 720), 1) - pygame.draw.line(screen, gridLight, (0, i + Yoff + cYoff), (1280, i + Yoff + cYoff), 1) - - pygame.draw.line(screen, gridDark, (pole.x + Xoff, 0), (pole.x + Xoff, 720), 1) - pygame.draw.line(screen, gridDark, (0, pole.y + Yoff), (1280, pole.y + Yoff), 1) - -def plotCenteredText(font, text = "", colour = "black", y = 0): - textObj = font.render(text, True, colour) - text_rect = textObj.get_rect(center=(1280/2, 720/2 - y)) - screen.blit(textObj, text_rect) +# UI helpers +ui = SimUI(screen, pole) # Pendulum setup # Start angle in radians, length, mass, color pendulum = Pendulum(0, 2, 0, 0.25, "red") pendulum.reset() +dx = 0 # x offset dt = 1 # delta time # Gametime rt = 10 # run time highscore = 0 + +# Metadata values +def meta(): + ui.meta(pendulum.theta[pendulum.index], "Theta") + ui.meta(pendulum.a_ang[pendulum.index], "Angular acceleration") + ui.meta(pendulum.dx, "dx") + ui.meta(pendulum.a_cart[pendulum.index], "Cart acceleration") + ui.meta(pendulum.pid, "Control") + ui.meta(not update, "Paused") + ui.meta(rt / 1000, "Run time [s]") + ui.meta(highscore / 1000, "Highscore [s]") + + while running: - # poll for events - # pygame.QUIT event means the user clicked X to close your window + ### User controls ### for event in pygame.event.get(): + # Quit if event.type == pygame.QUIT: running = False elif event.type == pygame.KEYDOWN: + # Quit if event.key == pygame.K_ESCAPE: running = False + # Reset simulation elif event.key == pygame.K_SPACE: + pendulum.reset() rt = 0 + # Pause simulation elif event.key == pygame.K_p: if update: update = False else: update = True + # Display plot if simulation is not running + elif event.key == pygame.K_g: + if pendulum.fallen: + pendulum.plot() + # Toggle PID controller + elif event.key == pygame.K_c: + if pendulum.pid: + pendulum.pid = False + else: + pendulum.pid = True + # Move cart keys = pygame.key.get_pressed() if keys[pygame.K_LEFT] or keys[pygame.K_a]: pendulum.a_cart[pendulum.index] -= 4 if keys[pygame.K_RIGHT] or keys[pygame.K_d]: pendulum.a_cart[pendulum.index] += 4 - # Draw x axis - plotGrid(50, 0, 15) + # Draw grid + ui.grid(50, 0, 15) # Update pendulum if not pendulum.fallen: @@ -85,41 +93,22 @@ while running: rt += dt pendulum.update(dt) else: - # ol = pygame.Surface(screen.get_size()) - # screen.fill((0, 0, 0, 255)) - - plotCenteredText(font_g, "WASTED", "red", 100) - plotCenteredText(font_m, "Press space to restart", "black", 60) - plotCenteredText(font_m, "Press G to view nerd graphs", "black", 45) + ui.wasted() + # Update highscore if rt > highscore: highscore = rt - dx = (pendulum.dx, 0) - # Draw metadata - screen.blit(font_h.render("Pendulum simulator 4000", True, "black"), (10, 10)) - screen.blit(font_m.render("Arne van Iterson, 2023", True, "black"), (1150, 700)) - - plot_y = 40 - plotMeta(pendulum.theta[pendulum.index], "Theta") - plotMeta(pendulum.a_ang[pendulum.index], "Angular acceleration") - - plotMeta(pendulum.dx, "dx") - plotMeta(pendulum.a_cart[pendulum.index], "Cart acceleration") - - plotMeta(dt, "Frame time") - plotMeta(1000 / dt, "FPS") - plotMeta(pendulum.pid, "Control") - plotMeta(not update, "Paused") - - plotMeta(rt / 1000, "Run time [s]") - plotMeta(highscore / 1000, "Highscore [s]") + ui.update() + meta() # Draw pendulum + dx = (pendulum.dx, 0) pygame.draw.line(screen, pendulum.color, pole + dx, pole + pendulum.vector + dx, 3) pygame.draw.circle(screen, "black", pole + dx, 15, 3) + # Draw frame pygame.display.flip() dt = clock.tick(60) # limits FPS to 120 diff --git a/src/sim/uiHelpers.py b/src/sim/uiHelpers.py index 44bc2bb..cc05d1b 100644 --- a/src/sim/uiHelpers.py +++ b/src/sim/uiHelpers.py @@ -1,7 +1,78 @@ import pygame +# Constants C_GRID_L_VALUE = 200 -gridLight = pygame.Color(C_GRID_L_VALUE,C_GRID_L_VALUE,C_GRID_L_VALUE) - C_GRID_D_VALUE = 100 -gridDark = pygame.Color(C_GRID_D_VALUE,C_GRID_D_VALUE,C_GRID_D_VALUE) \ No newline at end of file +C_MPLOT_START = 50 + +gridLight = pygame.Color(C_GRID_L_VALUE, C_GRID_L_VALUE, C_GRID_L_VALUE) +gridDark = pygame.Color(C_GRID_D_VALUE, C_GRID_D_VALUE, C_GRID_D_VALUE) +font_h = pygame.font.SysFont(None, 28) +font_m = pygame.font.SysFont(None, 16) + + +# UI Class +class SimUI: + def __init__(self, screen, pole): + self.screen = screen + self.pole = pole + + self.metaPlotY = 50 + + def meta(self, val, desc): + self.screen.blit( + font_m.render(f"{desc} = {val}", True, "black"), (15, self.metaPlotY) + ) + self.metaPlotY += 15 + + def grid(self, dist, Xoff=0, Yoff=0): + self.screen.fill("white") + cXoff = self.pole.x % dist + cYoff = self.pole.y % dist + + for i in range(0, 1280, dist): + pygame.draw.line( + self.screen, + gridLight, + (i + Xoff + cXoff, 0), + (i + Xoff + cXoff, 720), + 1, + ) + pygame.draw.line( + self.screen, + gridLight, + (0, i + Yoff + cYoff), + (1280, i + Yoff + cYoff), + 1, + ) + + pygame.draw.line( + self.screen, gridDark, (self.pole.x + Xoff, 0), (self.pole.x + Xoff, 720), 1 + ) + pygame.draw.line( + self.screen, + gridDark, + (0, self.pole.y + Yoff), + (1280, self.pole.y + Yoff), + 1, + ) + + def centeredText(self, font, text="", colour="black", y=0): + textObj = font.render(text, True, colour) + text_rect = textObj.get_rect(center=(1280 / 2, 720 / 2 - y)) + self.screen.blit(textObj, text_rect) + + def wasted(self): + font_g = pygame.font.SysFont(None, 128) + self.centeredText(font_g, "WASTED", "red", 100) + self.centeredText(font_m, "Press space to restart", "black", 60) + self.centeredText(font_m, "Press G to view nerd graphs", "black", 45) + + def update(self): + self.screen.blit( + font_h.render("Pendulum simulator 4000", True, "black"), (10, 10) + ) + self.screen.blit( + font_m.render("Arne van Iterson, 2023", True, "black"), (1150, 700) + ) + self.metaPlotY = C_MPLOT_START