Fixed math and added game(tm)

src/sim/pendulum.py

@@ -1,10 +1,13 @@
 from pygame.math import Vector2
 import math
 import numpy as np
+import random
 
 # 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:
     def __init__(self, theta, length, dx, mass, color):
@@ -12,10 +15,10 @@ class Pendulum:
         Initialize a Pendulum object.
 
         Parameters:
-        theta (float): Angle in radians.
+        theta (float): Angle [rad].
-        length (float): Length of the pendulum.
+        length (float): Length of the pendulum [m].
-        dx (float): Horizontal displacement of the "cart" from the center.
+        dx (float): Horizontal displacement of the "cart" from the center [m].
-        mass (float): Mass of the pendulum for physics calculations.
+        mass (float): Mass of the pendulum for physics calculations [kg].
         color (str): Display color.
 
         Returns:
@@ -30,40 +33,63 @@ class Pendulum:
         self.dx = dx  # Horizontal displacement of "cart" from center
         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 + (1.5 * math.pi)))
         )
 
+        if abs(self.theta) == C_FALL_ANG:
+            self.fallen = True
+
     def doMath(self, dt):
         # Angle
-        ang_term1 = -(C_GRAVITY * math.sin(self.theta))
+        ang_term1 = self.a_cart * math.cos(self.theta)
-        ang_term2 = self.a_cart * math.sin(self.theta)
+        ang_term2 = self.v_cart * math.sin(self.theta)
-        ang_term3 = self.v_cart * self.v_ang * math.cos(self.theta)
+        ang_term3 = self.v_cart * self.v_ang * math.sin(self.theta)
+        ang_term4 = C_GRAVITY * math.sin(self.theta)
 
-        # 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 - ang_term4) / -(self.length)) - ((self.r_factor / 2) * 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
 
+        # Limit fall of pendulum
+        self.theta = self.clamp(self.theta, -C_FALL_ANG, C_FALL_ANG)
+
         # Cart pos
-        cart_term1 = self.length * self.a_ang * math.sin(self.theta)
+        cart_term1 = self.mass * self.length * self.a_ang * math.cos(self.theta)
-        cart_term2 = self.length * pow(self.v_ang, 2) * math.cos(self.theta)
+        cart_term2 = self.mass * self.length * self.v_ang * math.sin(self.theta)
         
-        self.a_cart = (cart_term1 - cart_term2) / 2  # Cart acceleration
+        self.a_cart = (-cart_term1 + cart_term2) / (2 * self.mass) - (self.r_factor * self.v_cart) # Cart acceleration
-        self.v_cart = (self.a_cart * (dt / 1000) + self.v_cart)  # Integrate acceleration to get velocity
+        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
+        self.s_cart = (self.v_cart * (dt / 1000) + self.s_cart)  # Cart displacement
+        self.dx = self.s_cart * C_MTPRATIO  # Convert to pixels
 
+    def clamp(self, n, minn, maxn):
+        return max(min(maxn, n), minn)
+
+    def reset(self):
+        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 update(self, dt):
     #     """

src/sim/sim.py

@@ -9,57 +9,91 @@ pygame.init()
 screen = pygame.display.set_mode((1280, 720))
 clock = pygame.time.Clock()
 running = True
+update = True
 
 # Text setup
+font_g = pygame.font.SysFont(None, 128)
 font_h = pygame.font.SysFont(None, 28)
 font_m = pygame.font.SysFont(None, 16)
 
 # Metadata plotter
-plot_y = 40
+plot_y = 50
 def plotMeta(val, desc):
     global plot_y
-    screen.blit(font_m.render(f"{desc} = {val}", True, "black"), (10, plot_y))
+    screen.blit(font_m.render(f"{desc} = {val}", True, "black"), (15, plot_y))
     plot_y += 15
 
-
 # Pendulum setup
 # Start angle in radians, length, mass, color
-p_t_start = 99 / 100 * math.pi
+pendulum = Pendulum(0, 2, 0, 0.25, "red")
-pendulum = Pendulum(p_t_start, 1, 0, 100, "red")
+pendulum.reset()
 dt = 1  # delta time
 
+# Gametime
+rt = 10 # run time
+highscore = 0 
+
 while running:
     # poll for events
     # pygame.QUIT event means the user clicked X to close your window
     for event in pygame.event.get():
         if event.type == pygame.QUIT:
             running = False
-        if event.type == pygame.KEYDOWN:
+        elif event.type == pygame.KEYDOWN: 
-            keys = pygame.key.get_pressed()
+            if event.key == pygame.K_ESCAPE:
-            if keys[pygame.K_SPACE]:
+                running = False
-                pendulum.theta = p_t_start
+            elif event.key == pygame.K_SPACE:
-            if keys[pygame.K_LEFT]:
+                pendulum.reset()
-                pendulum.dx -= 1
+                rt = 0
-            if keys[pygame.K_RIGHT]:
+            elif event.key == pygame.K_p: 
-                pendulum.dx += 1
+                if update:
+                    update = False
+                else:
+                    update = 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
 
+    keys = pygame.key.get_pressed()
+    if keys[pygame.K_LEFT]:
+        pendulum.a_cart -= 4
+    if keys[pygame.K_RIGHT]:
+        pendulum.a_cart += 4
+
     # Update pendulum
+    if not pendulum.fallen:
+        screen.fill("gray")
+        if update:
+            rt += dt
             pendulum.update(dt)
+    else:
+        screen.fill("darkgray")
+        screen.blit(font_g.render("WASTED", True, "red"), (pole.x - 200, pole.y - 120))
+        screen.blit(font_m.render("Press space to restart", True, "black"), (pole.x - 70, pole.y - 40))
+        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, "Theta")
+    plotMeta(pendulum.a_ang, "Angular acceleration")
+    
     plotMeta(pendulum.dx, "dx")
+    plotMeta(pendulum.a_cart, "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
     pygame.draw.line(screen, pendulum.color, pole + dx, pole + pendulum.vector + dx, 3)
@@ -69,6 +103,6 @@ while running:
     pygame.draw.line(screen, "black", (0, pole.y + 15), (1280, pole.y + 15), 1)
 
     pygame.display.flip()
-    dt = clock.tick(120)  # limits FPS to 120
+    dt = clock.tick(60)  # limits FPS to 120
 
 pygame.quit()