Added display grid and array structure
This commit is contained in:
parent
7ab9feb3fb
commit
e4b327ddcf
@ -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.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.s_cart = 0 # Displacement of cart [m]
|
||||
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.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
|
||||
ang_term1 = self.a_cart * math.cos(self.theta)
|
||||
ang_term2 = self.v_cart * math.sin(self.theta)
|
||||
ang_term3 = self.v_cart * self.v_ang * math.sin(self.theta)
|
||||
ang_term4 = C_GRAVITY * math.sin(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 - ang_term4) / -(self.length)) - ((self.r_factor / 2) * self.v_ang) # 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)
|
||||
)
|
||||
|
||||
# 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_term1 = self.mass * self.length * self.a_ang * 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 * 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.s_cart = (self.v_cart * (dt / 1000) + self.s_cart) # Cart displacement
|
||||
self.dx = self.s_cart * C_MTPRATIO # Convert to pixels
|
||||
### 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.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.v_ang = 0
|
||||
self.dx = 0
|
||||
self.theta = random.choice([1, -1]) * C_P_ANG_START
|
||||
self.index = 0
|
||||
|
||||
self.a_cart = 0
|
||||
self.v_cart = 0
|
||||
self.s_cart = 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)
|
||||
|
||||
|
||||
@ -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
|
||||
highscore = 0
|
||||
rt = 10 # run time
|
||||
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]:
|
||||
pendulum.a_cart -= 4
|
||||
if keys[pygame.K_RIGHT]:
|
||||
pendulum.a_cart += 4
|
||||
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
|
||||
plotGrid(50, 0, 15)
|
||||
|
||||
# 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))
|
||||
# ol = pygame.Surface(screen.get_size())
|
||||
# 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)
|
||||
|
||||
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.a_ang, "Angular acceleration")
|
||||
|
||||
plotMeta(pendulum.theta[pendulum.index], "Theta")
|
||||
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
|
||||
|
||||
7
src/sim/uiHelpers.py
Normal file
7
src/sim/uiHelpers.py
Normal file
@ -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)
|
||||
Loading…
Reference in New Issue
Block a user