Added nerd graphs for debugging

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.

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
+# Own objects must be imported after pygame init
-font_g = pygame.font.SysFont(None, 128)
+from pendulum import Pendulum
-font_h = pygame.font.SysFont(None, 28)
+from uiHelpers import *
-font_m = pygame.font.SysFont(None, 16)
 
-# Metadata plotter
+# UI helpers
-plot_y = 50
+ui = SimUI(screen, pole)
-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)
 
 # 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
+    ### User controls ###
-    # pygame.QUIT event means the user clicked X to close your window
     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
+    # Draw grid
-    plotGrid(50, 0, 15)
+    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())
+        ui.wasted()
-        # 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)
 
+        # 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))
+    ui.update()
-    screen.blit(font_m.render("Arne van Iterson, 2023", True, "black"), (1150, 700))
+    meta()
-
-    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]")
 
     # 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
 

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
+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