134 lines
3.5 KiB
Python
134 lines
3.5 KiB
Python
# Pendulum simulator 4000
|
|
# Arne van Iterson, 2023
|
|
|
|
# Imports
|
|
import pygame_widgets
|
|
import pygame
|
|
from pygame_widgets.slider import Slider
|
|
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, 0.2, 0, 0.7, "red")
|
|
pendulum.reset()
|
|
dx = 0 # x offset
|
|
dt = 1 # delta time
|
|
|
|
# Sliders
|
|
slider_kp = Slider(screen, 910, 590, 320, 16, initial=pendulum.kp, min=0, max=2, step=0.001)
|
|
slider_ki = Slider(screen, 910, 620, 320, 16, initial=pendulum.ki, min=0, max=0.5, step=0.001)
|
|
slider_kd = Slider(screen, 910, 650, 320, 16, initial=pendulum.kd, min=0, max=0.5, step=0.001)
|
|
|
|
# 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(pendulum.kd, "Kd")
|
|
ui.meta(pendulum.ki, "Ki")
|
|
ui.meta(pendulum.kp, "Kp")
|
|
|
|
ui.meta(dt, "dt")
|
|
|
|
ui.meta(not update, "Paused")
|
|
ui.meta(rt / 1000, "Run time [s]")
|
|
ui.meta(highscore / 1000, "Highscore [s]")
|
|
|
|
|
|
while running:
|
|
events = pygame.event.get()
|
|
### User controls ###
|
|
for event in events:
|
|
# 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:
|
|
update = not update
|
|
# Display plot if simulation is not running
|
|
elif event.key == pygame.K_g:
|
|
if pendulum.fallen:
|
|
pendulum.plot()
|
|
else:
|
|
update = False
|
|
pendulum.plot()
|
|
# Toggle PID controller
|
|
elif event.key == pygame.K_c:
|
|
pendulum.pid = not pendulum.pid
|
|
|
|
# Move pendulum
|
|
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 PID values
|
|
# pendulum.kp = slider_kp.getValue()
|
|
# pendulum.ki = slider_ki.getValue()
|
|
# pendulum.kd = slider_kd.getValue()
|
|
|
|
# Update pendulum
|
|
if not pendulum.fallen:
|
|
if update:
|
|
rt += dt
|
|
pendulum.update(dt)
|
|
else:
|
|
ui.gameover(rt)
|
|
|
|
# Update highscore
|
|
if rt > highscore:
|
|
highscore = rt
|
|
|
|
# Draw metadata
|
|
ui.update(dt)
|
|
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_widgets.update(events)
|
|
pygame.display.flip()
|
|
dt = clock.tick(60) # limits FPS to 120
|
|
|
|
pygame.quit()
|