Added display grid and array structure

src/sim/pendulum.py

@@ -5,10 +5,11 @@ import random
 
 # Constants
 C_GRAVITY = 9.81  # m/s^2
-C_MTPRATIO = 100 # Pixels per meter
+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):
         """
@@ -26,16 +27,17 @@ class Pendulum:
         """
         self.vector = None  # Vector2 object
 
-        self.theta = theta  # Angle in radians
+        self.index = 0
-        self.a_ang = 0  # Angular acceleration
+        self.theta = [theta]  # Angle in radians
-        self.v_ang = 0  # Angular velocity
+        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.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.s_cart = [0]  # Displacement of cart [m]
 
-        self.r_factor = 0.50  # Damping factor
+        # self.r_factor = 0.50  # Damping factor
 
         self.length = length  # Length of pendulum
         self.mass = mass  # Mass of pendulum for physics
@@ -47,47 +49,92 @@ class Pendulum:
     def update(self, dt):
         self.doMath(dt)
         self.vector = Vector2.from_polar(
-            ((self.length * C_MTPRATIO), math.degrees(self.theta + (1.5 * math.pi)))
+            ((self.length * C_MTPRATIO), math.degrees(self.theta[self.index] + (1.5 * math.pi)))
         )
 
-        if abs(self.theta) == C_FALL_ANG:
+        if abs(self.theta[self.index]) == C_FALL_ANG:
             self.fallen = True
 
     def doMath(self, dt):
-        # Angle
+        ### ANGLE ###
-        ang_term1 = self.a_cart * math.cos(self.theta)
+        ang_term1 = self.a_cart[self.index] * math.cos(self.theta[self.index])
-        ang_term2 = self.v_cart * math.sin(self.theta)
+        ang_term2 = self.v_cart[self.index] * math.sin(self.theta[self.index])
-        ang_term3 = self.v_cart * self.v_ang * math.sin(self.theta)
+        ang_term3 = (
-        ang_term4 = C_GRAVITY * math.sin(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])
 
-        self.a_ang = ((ang_term1 - ang_term2 + ang_term3 - ang_term4) / -(self.length)) - ((self.r_factor / 2) * self.v_ang) # Angular acceleration        
+        # Angular acceleration
-        self.v_ang = self.a_ang * (dt / 1000) + self.v_ang  # Integrate acceleration to get velocity
+        self.a_ang.append(
-        self.theta = self.v_ang * (dt / 1000) + self.theta  # Angular displacement
+            (ang_term1 - ang_term2 + ang_term3 - ang_term4) / -(self.length)
+        )
+
+        # Integrate acceleration to get velocity
+        self.v_ang.append(
+            self.a_ang[self.index + 1] * (dt / 1000)
+            + self.v_ang[self.index]  # Previous velocity
+        )
+
+        # Angular displacement
+        self.theta.append(
+            self.v_ang[self.index + 1] * (dt / 1000)
+            + self.theta[self.index]  # Previous angle
+        )
 
         # Limit fall of pendulum
-        self.theta = self.clamp(self.theta, -C_FALL_ANG, C_FALL_ANG)
+        self.theta[self.index + 1] = self.clamp(
+            self.theta[self.index + 1], -C_FALL_ANG, C_FALL_ANG
+        )
 
-        # Cart pos
+        ### CART ###
-        cart_term1 = self.mass * self.length * self.a_ang * math.cos(self.theta)
+        cart_term1 = (
-        cart_term2 = self.mass * self.length * self.v_ang * math.sin(self.theta)
+            self.mass  # Mass
-        
+            * self.length  # Length
-        self.a_cart = (-cart_term1 + cart_term2) / (2 * self.mass) - (self.r_factor * self.v_cart) # Cart acceleration
+            * self.a_ang[self.index + 1]  # Current angle acceleration
-        self.v_cart = self.a_cart * (dt / 1000) + self.v_cart  # Integrate acceleration to get velocity
+            * math.cos(self.theta[self.index + 1])  # Current angle
-        self.s_cart = (self.v_cart * (dt / 1000) + self.s_cart)  # Cart displacement
+        )
-        self.dx = self.s_cart * C_MTPRATIO  # Convert to pixels
+        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.a_cart[self.index + 1] * (dt / 1000)
+            + self.v_cart[self.index]  # Previous velocity
+        )
+
+        # Cart displacement
+        self.s_cart.append(
+            self.v_cart[self.index + 1] * (dt / 1000)
+            + self.s_cart[self.index]  # Previous displacement
+        )
+        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.a_ang = 0
+        self.index = 0
-        self.v_ang = 0
-        self.dx = 0
-        self.theta = random.choice([1, -1]) * C_P_ANG_START
         
-        self.a_cart = 0
+        self.a_ang = [0]
-        self.v_cart = 0
+        self.v_ang = [0]
-        self.s_cart = 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)
 

src/sim/sim.py

@@ -3,6 +3,7 @@ from pygame.math import Vector2
 import math
 
 from pendulum import Pendulum
+from uiHelpers import gridDark, gridLight
 
 # pygame setup
 pygame.init()
@@ -10,6 +11,7 @@ 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_g = pygame.font.SysFont(None, 128)
@@ -23,6 +25,23 @@ def plotMeta(val, desc):
     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")
@@ -30,8 +49,8 @@ pendulum.reset()
 dt = 1  # delta time
 
 # Gametime
-rt = 10 # run time
+rt = 10  # run time
-highscore = 0 
+highscore = 0
 
 while running:
     # poll for events
@@ -39,68 +58,67 @@ while running:
     for event in pygame.event.get():
         if event.type == pygame.QUIT:
             running = False
-        elif event.type == pygame.KEYDOWN: 
+        elif event.type == pygame.KEYDOWN:
             if event.key == pygame.K_ESCAPE:
                 running = False
             elif event.key == pygame.K_SPACE:
                 pendulum.reset()
                 rt = 0
-            elif event.key == pygame.K_p: 
+            elif event.key == pygame.K_p:
                 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]:
+    if keys[pygame.K_LEFT] or keys[pygame.K_a]:
-        pendulum.a_cart -= 4
+        pendulum.a_cart[pendulum.index] -= 4
-    if keys[pygame.K_RIGHT]:
+    if keys[pygame.K_RIGHT] or keys[pygame.K_d]:
-        pendulum.a_cart += 4
+        pendulum.a_cart[pendulum.index] += 4
+
+    # Draw x axis
+    plotGrid(50, 0, 15)
 
     # Update pendulum
     if not pendulum.fallen:
-        screen.fill("gray")
         if update:
             rt += dt
             pendulum.update(dt)
     else:
-        screen.fill("darkgray")
+        # ol = pygame.Surface(screen.get_size())
-        screen.blit(font_g.render("WASTED", True, "red"), (pole.x - 200, pole.y - 120))
+        # screen.fill((0, 0, 0, 255))
-        screen.blit(font_m.render("Press space to restart", True, "black"), (pole.x - 70, pole.y - 40))
+        
+        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)
+
         if rt > highscore:
             highscore = rt
-        
-    dx = (pendulum.dx, 0)
 
+    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.theta[pendulum.index], "Theta")
-    plotMeta(pendulum.a_ang, "Angular acceleration")
+    plotMeta(pendulum.a_ang[pendulum.index], "Angular acceleration")
-    
+
     plotMeta(pendulum.dx, "dx")
-    plotMeta(pendulum.a_cart, "Cart acceleration")
+    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
     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)
 
     pygame.display.flip()
     dt = clock.tick(60)  # limits FPS to 120

src/sim/uiHelpers.py

@@ -0,0 +1,7 @@
+import pygame
+
+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)