116 lines
2.9 KiB
Python
116 lines
2.9 KiB
Python
# Pendulum simulator 4000
|
|
# Arne van Iterson, 2023
|
|
|
|
# Imports
|
|
import pygame
|
|
from pygame.math import Vector2
|
|
import math
|
|
|
|
# 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
|
|
from pendulum import Pendulum
|
|
from uiHelpers import *
|
|
|
|
# UI helpers
|
|
ui = SimUI(screen, pole)
|
|
|
|
# Pendulum setup
|
|
# Start angle in radians, length, mass, color
|
|
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:
|
|
### User controls ###
|
|
for event in pygame.event.get():
|
|
# Quit
|
|
if event.type == pygame.QUIT:
|
|
running = False
|
|
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
|
|
|
|
# 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
|
|
if not pendulum.fallen:
|
|
if update:
|
|
rt += dt
|
|
pendulum.update(dt)
|
|
else:
|
|
ui.wasted()
|
|
|
|
# Update highscore
|
|
if rt > highscore:
|
|
highscore = rt
|
|
|
|
# Draw metadata
|
|
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 frame
|
|
pygame.display.flip()
|
|
dt = clock.tick(60) # limits FPS to 120
|
|
|
|
pygame.quit()
|