Added nerd graphs for debugging

src/sim/pendulum.py

@@ -1,9 +1,14 @@
 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
+C_MTPRATIO = 100  # Pixels per meter
+C_P_ANG_START = 1 / 1000 * math.pi
+C_FALL_ANG = 52.5 / 100 * math.pi
 
 
 class Pendulum:
@@ -12,10 +17,10 @@ class Pendulum:
         Initialize a Pendulum object.
 
         Parameters:
-        theta (float): Angle in radians.
-        length (float): Length of the pendulum.
-        dx (float): Horizontal displacement of the "cart" from the center.
-        mass (float): Mass of the pendulum for physics calculations.
+        theta (float): Angle [rad].
+        length (float): Length of the pendulum [m].
+        dx (float): Horizontal displacement of the "cart" from the center [m].
+        mass (float): Mass of the pendulum for physics calculations [kg].
         color (str): Display color.
 
         Returns:
@@ -23,47 +28,139 @@ class Pendulum:
         """
         self.vector = None  # Vector2 object
 
-        self.theta = theta  # Angle in radians
-        self.a_ang = 0  # Angular acceleration
-        self.v_ang = 0  # Angular velocity
+        self.index = 0
+        self.theta = [theta]  # Angle in radians
+        self.a_ang = [0]  # Angular acceleration
+        self.v_ang = [0]  # Angular velocity
 
         self.dx = dx  # Horizontal displacement of "cart" from center
-        self.a_cart = 0  # Acceleration of cart
-        self.v_cart = 0  # Velocity of cart
+        self.a_cart = [0]  # Acceleration of cart
+        self.v_cart = [0]  # Velocity of cart
+        self.s_cart = [0]  # Displacement of cart [m]
 
-        self.r_factor = 0.99  # Damping factor
+        # self.r_factor = 0.50  # Damping factor
 
         self.length = length  # Length of pendulum
         self.mass = mass  # Mass of pendulum for physics
         self.color = color  # Display color
 
         self.pid = False
+        self.fallen = False
 
     def update(self, dt):
         self.doMath(dt)
         self.vector = Vector2.from_polar(
-            ((self.length * 150), math.degrees(self.theta + math.pi / 2))
+            ((self.length * C_MTPRATIO), math.degrees(self.theta[self.index] + (1.5 * math.pi)))
         )
 
+        if abs(self.theta[self.index]) == C_FALL_ANG:
+            self.fallen = True
+
     def doMath(self, dt):
-        # Angle
-        ang_term1 = -(C_GRAVITY * math.sin(self.theta))
-        ang_term2 = self.a_cart * math.sin(self.theta)
-        ang_term3 = self.v_cart * self.v_ang * math.cos(self.theta)
+        ### ANGLE ###
+        ang_term1 = self.a_cart[self.index] * math.cos(self.theta[self.index])
+        ang_term2 = self.v_cart[self.index] * math.sin(self.theta[self.index])
+        ang_term3 = (
+            self.v_cart[self.index]  # Previous cart velocity
+            * self.v_ang[self.index]  # previous angle velocity
+            * math.sin(self.theta[self.index])  # Sin previous angle
+        )
+        ang_term4 = C_GRAVITY * math.sin(self.theta[self.index])
 
-        # self.a_ang = ((ang_term1 - ang_term2 - ang_term3) / self.length) - (self.r_factor * self.v_ang) # Angular acceleration        
-        self.a_ang = ((ang_term1 - ang_term2 - ang_term3) / self.length) # Angular acceleration        
-        self.v_ang = self.a_ang * (dt / 1000) + self.v_ang  # Integrate acceleration to get velocity
-        self.theta = self.v_ang * (dt / 1000) + self.theta  # Angular displacement
+        # Angular acceleration
+        self.a_ang.append(
+            (ang_term1 - ang_term2 + ang_term3 - ang_term4) / -(self.length)
+        )
 
-        # Cart pos
-        cart_term1 = self.length * self.a_ang * math.sin(self.theta)
-        cart_term2 = self.length * pow(self.v_ang, 2) * math.cos(self.theta)
+        # Integrate acceleration to get velocity
+        self.v_ang.append(
+            self.v_ang[self.index]  # Previous velocity
+            + (self.a_ang[self.index + 1] * (dt / 1000))
+        )
 
-        self.a_cart = (cart_term1 - cart_term2) / 2  # Cart acceleration
-        self.v_cart = (self.a_cart * (dt / 1000) + self.v_cart)  # Integrate acceleration to get velocity
-        self.dx = self.v_cart * (dt / 1000) + self.dx  # Cart displacement
+        # Angular displacement
+        self.theta.append(
+            self.theta[self.index]  # Previous angle
+            + (self.v_ang[self.index + 1] * (dt / 1000))
+        )
 
+        # Limit fall of pendulum
+        self.theta[self.index + 1] = self.clamp(
+            self.theta[self.index + 1], -C_FALL_ANG, C_FALL_ANG
+        )
+
+        ### CART ###
+        cart_term1 = (
+            self.mass  # Mass
+            * self.length  # Length
+            * self.a_ang[self.index + 1]  # Current angle acceleration
+            * math.cos(self.theta[self.index + 1])  # Current angle
+        )
+        cart_term2 = (
+            self.mass  # Mass
+            * self.length  # Length
+            * self.v_ang[self.index + 1]  # Current angle velocity
+            * math.sin(self.theta[self.index + 1])  # Current angle
+        )
+
+        # Cart acceleration
+        self.a_cart.append((-cart_term1 + cart_term2) / (2 * self.mass))
+
+        # Integrate acceleration to get velocity
+        self.v_cart.append(
+            self.v_cart[self.index]  # Previous velocity
+            + (self.a_cart[self.index + 1] * (dt / 1000))
+        )
+
+        # Cart displacement
+        self.s_cart.append(
+            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
+
+        # Update index
+        self.index += 1
+
+    def clamp(self, n, minn, maxn):
+        return max(min(maxn, n), minn)
+
+    def reset(self):
+        self.index = 0
+        
+        self.a_ang = [0]
+        self.v_ang = [0]
+        self.dx = [0]
+        self.theta = [random.choice([1, -1]) * C_P_ANG_START]
+
+        self.a_cart = [0]
+        self.v_cart = [0]
+        self.s_cart = [0]
+        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):
     #     """

src/sim/sim.py

@@ -1,74 +1,115 @@
+# Pendulum simulator 4000
+# Arne van Iterson, 2023
+
+# Imports
 import pygame
 from pygame.math import Vector2
 import math
 
-from pendulum import Pendulum
-
 # pygame setup
 pygame.init()
 screen = pygame.display.set_mode((1280, 720))
 clock = pygame.time.Clock()
 running = True
+update = True
+pole = Vector2(screen.get_rect().center)  # center of screen
 
-# Text setup
-font_h = pygame.font.SysFont(None, 28)
-font_m = pygame.font.SysFont(None, 16)
-
-# Metadata plotter
-plot_y = 40
-def plotMeta(val, desc):
-    global plot_y
-    screen.blit(font_m.render(f"{desc} = {val}", True, "black"), (10, plot_y))
-    plot_y += 15
+# Own objects must be imported after pygame init
+from pendulum import Pendulum
+from uiHelpers import *
 
+# UI helpers
+ui = SimUI(screen, pole)
 
 # Pendulum setup
 # Start angle in radians, length, mass, color
-p_t_start = 99 / 100 * math.pi
-pendulum = Pendulum(p_t_start, 1, 0, 100, "red")
+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
-        if event.type == pygame.KEYDOWN:
-            keys = pygame.key.get_pressed()
-            if keys[pygame.K_SPACE]:
-                pendulum.theta = p_t_start
-            if keys[pygame.K_LEFT]:
-                pendulum.dx -= 1
-            if keys[pygame.K_RIGHT]:
-                pendulum.dx += 1
+        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
 
-    # fill the screen with a color to wipe away anything from last frame
-    screen.fill("gray")
-    pole = Vector2(screen.get_rect().center)  # center of screen
+    # 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 grid
+    ui.grid(50, 0, 15)
 
     # Update pendulum
-    pendulum.update(dt)
-    dx = (pendulum.dx, 0)
+    if not pendulum.fallen:
+        if update:
+            rt += dt
+            pendulum.update(dt)
+    else:
+        ui.wasted()
+
+        # Update highscore
+        if rt > highscore:
+            highscore = rt
 
     # Draw metadata
-    screen.blit(font_h.render("Pendulum simulator 4000", True, "black"), (10, 10))
-    
-    plot_y = 40
-    plotMeta(pendulum.theta, "Theta")
-    plotMeta(pendulum.dx, "dx")
-    plotMeta(dt, "Frame time")
-    plotMeta(1000 / dt, "FPS")
-    plotMeta(pendulum.pid, "Control")
+    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 x axis
-    pygame.draw.line(screen, "black", (0, pole.y + 15), (1280, pole.y + 15), 1)
 
+    # Draw frame
     pygame.display.flip()
-    dt = clock.tick(120)  # limits FPS to 120
+    dt = clock.tick(60)  # limits FPS to 120
 
 pygame.quit()

src/sim/uiHelpers.py

@@ -0,0 +1,78 @@
+import pygame
+
+# Constants
+C_GRID_L_VALUE = 200
+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