src/sim/pendulum.py

@@ -1,14 +1,9 @@
 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:
@@ -17,10 +12,10 @@ class Pendulum:
         Initialize a Pendulum object.
 
         Parameters:
-        theta (float): Angle [rad].
+        theta (float): Angle in radians.
-        length (float): Length of the pendulum [m].
+        length (float): Length of the pendulum.
-        dx (float): Horizontal displacement of the "cart" from the center [m].
+        dx (float): Horizontal displacement of the "cart" from the center.
-        mass (float): Mass of the pendulum for physics calculations [kg].
+        mass (float): Mass of the pendulum for physics calculations.
         color (str): Display color.
 
         Returns:
@@ -28,139 +23,47 @@ class Pendulum:
         """
         self.vector = None  # Vector2 object
 
-        self.index = 0
+        self.theta = theta  # Angle in radians
-        self.theta = [theta]  # Angle in radians
+        self.a_ang = 0  # Angular acceleration
-        self.a_ang = [0]  # Angular acceleration
+        self.v_ang = 0  # Angular velocity
-        self.v_ang = [0]  # Angular velocity
 
         self.dx = dx  # Horizontal displacement of "cart" from center
-        self.a_cart = [0]  # Acceleration of cart
+        self.a_cart = 0  # Acceleration of cart
-        self.v_cart = [0]  # Velocity of cart
+        self.v_cart = 0  # Velocity of cart
-        self.s_cart = [0]  # Displacement of cart [m]
 
-        # self.r_factor = 0.50  # Damping factor
+        self.r_factor = 0.99  # 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 * C_MTPRATIO), math.degrees(self.theta[self.index] + (1.5 * math.pi)))
+            ((self.length * 150), math.degrees(self.theta + math.pi / 2))
         )
 
-        if abs(self.theta[self.index]) == C_FALL_ANG:
-            self.fallen = True
-
     def doMath(self, dt):
-        ### ANGLE ###
+        # Angle
-        ang_term1 = self.a_cart[self.index] * math.cos(self.theta[self.index])
+        ang_term1 = -(C_GRAVITY * math.sin(self.theta))
-        ang_term2 = self.v_cart[self.index] * math.sin(self.theta[self.index])
+        ang_term2 = self.a_cart * math.sin(self.theta)
-        ang_term3 = (
+        ang_term3 = self.v_cart * self.v_ang * math.cos(self.theta)
-            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])
 
-        # Angular acceleration
+        # self.a_ang = ((ang_term1 - ang_term2 - ang_term3) / self.length) - (self.r_factor * self.v_ang) # Angular acceleration        
-        self.a_ang.append(
+        self.a_ang = ((ang_term1 - ang_term2 - ang_term3) / self.length) # Angular acceleration        
-            (ang_term1 - ang_term2 + ang_term3 - ang_term4) / -(self.length)
+        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
 
-        # Integrate acceleration to get velocity
+        # Cart pos
-        self.v_ang.append(
+        cart_term1 = self.length * self.a_ang * math.sin(self.theta)
-            self.v_ang[self.index]  # Previous velocity
+        cart_term2 = self.length * pow(self.v_ang, 2) * math.cos(self.theta)
-            + (self.a_ang[self.index + 1] * (dt / 1000))
-        )
 
-        # Angular displacement
+        self.a_cart = (cart_term1 - cart_term2) / 2  # Cart acceleration
-        self.theta.append(
+        self.v_cart = (self.a_cart * (dt / 1000) + self.v_cart)  # Integrate acceleration to get velocity
-            self.theta[self.index]  # Previous angle
+        self.dx = self.v_cart * (dt / 1000) + self.dx  # Cart displacement
-            + (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,115 +1,74 @@
-# 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
 
-# Own objects must be imported after pygame init
+# Text setup
-from pendulum import Pendulum
+font_h = pygame.font.SysFont(None, 28)
-from uiHelpers import *
+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
 
-# UI helpers
-ui = SimUI(screen, pole)
 
 # Pendulum setup
 # Start angle in radians, length, mass, color
-pendulum = Pendulum(0, 2, 0, 0.25, "red")
+p_t_start = 99 / 100 * math.pi
-pendulum.reset()
+pendulum = Pendulum(p_t_start, 1, 0, 100, "red")
-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:
-    ### User controls ###
+    # poll for events
+    # 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:
+        if 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]:
+            if keys[pygame.K_SPACE]:
-        pendulum.a_cart[pendulum.index] -= 4
+                pendulum.theta = p_t_start
-    if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
+            if keys[pygame.K_LEFT]:
-        pendulum.a_cart[pendulum.index] += 4
+                pendulum.dx -= 1
+            if keys[pygame.K_RIGHT]:
+                pendulum.dx += 1
 
-    # Draw grid
+    # fill the screen with a color to wipe away anything from last frame
-    ui.grid(50, 0, 15)
+    screen.fill("gray")
+    pole = Vector2(screen.get_rect().center)  # center of screen
 
     # Update pendulum
-    if not pendulum.fallen:
-        if update:
-            rt += dt
     pendulum.update(dt)
-    else:
+    dx = (pendulum.dx, 0)
-        ui.wasted()
-
-        # Update highscore
-        if rt > highscore:
-            highscore = rt
 
     # Draw metadata
-    ui.update()
+    screen.blit(font_h.render("Pendulum simulator 4000", True, "black"), (10, 10))
-    meta()
+    
+    plot_y = 40
+    plotMeta(pendulum.theta, "Theta")
+    plotMeta(pendulum.dx, "dx")
+    plotMeta(dt, "Frame time")
+    plotMeta(1000 / dt, "FPS")
+    plotMeta(pendulum.pid, "Control")
 
     # 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
+    # Draw x axis
+    pygame.draw.line(screen, "black", (0, pole.y + 15), (1280, pole.y + 15), 1)
+
     pygame.display.flip()
-    dt = clock.tick(60)  # limits FPS to 120
+    dt = clock.tick(120)  # limits FPS to 120
 
 pygame.quit()

src/sim/uiHelpers.py

@@ -1,78 +0,0 @@
-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