234 lines
7.7 KiB
Python
234 lines
7.7 KiB
Python
from pygame.math import Vector2
|
|
import math
|
|
import numpy as np
|
|
import random
|
|
import matplotlib.pyplot as plt
|
|
|
|
# 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
|
|
|
|
C_ANG_RES_LIMIT = ((2 * math.pi) / 360) # rad
|
|
C_V_CART_LIMIT = 0.403 # m/s
|
|
C_A_CART_LIMIT = C_V_CART_LIMIT / 0.1 # m/s^2
|
|
|
|
class Pendulum:
|
|
def __init__(self, theta, length, dx, mass, color):
|
|
"""
|
|
Initialize a Pendulum object.
|
|
|
|
Parameters:
|
|
theta (float): Angle [rad].
|
|
length (float): Length of the pendulum [m].
|
|
dx (float): Horizontal displacement of the "cart" from the center [m].
|
|
mass (float): Mass of the pendulum for physics calculations [kg].
|
|
color (str): Display color.
|
|
|
|
Returns:
|
|
None
|
|
"""
|
|
## Game variables
|
|
self.vector = None # Vector2 object
|
|
self.fallen = False # Stop when pendulum falls over
|
|
|
|
## Physics variables
|
|
self.index = 0 # Index helper for plotting graphs
|
|
self.theta = [theta] # Angle in radians
|
|
self.a_ang = [0] # Angular acceleration
|
|
self.v_ang = [0] # Angular velocity
|
|
|
|
## Gyro offset
|
|
self.theta_offset = 0
|
|
|
|
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.50 # Damping factor
|
|
|
|
self.length = length # Length of pendulum
|
|
self.mass = mass # Mass of pendulum for physics
|
|
self.color = color # Display color
|
|
|
|
## PID variables
|
|
self.pid = True
|
|
self.pidDelay = 0
|
|
# self.kp = 1.3
|
|
# self.ki = 0.0
|
|
# self.kd = 0.1
|
|
self.kp = 7.5
|
|
self.ki = 0.0
|
|
self.kd = 1.15
|
|
|
|
self.kp_m = 0.07
|
|
self.kd_m = 0.1
|
|
|
|
def update(self, dt):
|
|
self.doMath(dt)
|
|
self.vector = Vector2.from_polar(
|
|
((self.length * C_MTPRATIO), math.degrees(self.theta[self.index] - (0.5 * math.pi)))
|
|
)
|
|
|
|
if self.pid:
|
|
self.pidDelay += dt
|
|
|
|
if self.pidDelay > 21:
|
|
self.pidControl(self.pidDelay)
|
|
self.theta_offset += self.pidDelay * (C_ANG_RES_LIMIT / 1000) ## One degree a second
|
|
self.pidDelay = 0
|
|
|
|
if abs(self.theta[self.index]) == C_FALL_ANG:
|
|
self.fallen = True
|
|
|
|
# def update(self, dt):
|
|
# """
|
|
# Update the pendulum's state based on the elapsed time.
|
|
|
|
# Parameters:
|
|
# - dt (float): The elapsed time in milliseconds.
|
|
|
|
# Returns:
|
|
# None
|
|
# """
|
|
# a_ang = (-(C_GRAVITY * math.sin(self.theta)) / (self.length)) - (self.r_factor * self.v_ang) # Angular acceleration
|
|
|
|
# v_ang = a_ang * (dt/1000) + self.v_ang # Integrate acceleration to get velocity
|
|
# s_ang = v_ang * (dt/1000) # Angular displacement
|
|
|
|
# self.theta += s_ang # Update value
|
|
|
|
# self.vector = Vector2.from_polar(((self.length * 150), math.degrees(self.theta + math.pi/2)))
|
|
|
|
# self.a_ang = a_ang # Update value
|
|
# self.v_ang = v_ang # Update value
|
|
|
|
def doMath(self, dt):
|
|
### ANGLE ###
|
|
ang_term1 = self.a_cart[self.index] * math.cos(self.theta[self.index] + self.theta_offset)
|
|
ang_term2 = self.v_cart[self.index] * math.sin(self.theta[self.index] + self.theta_offset)
|
|
ang_term3 = (
|
|
self.v_cart[self.index] # Previous cart velocity
|
|
* self.v_ang[self.index] # previous angle velocity
|
|
* math.sin(self.theta[self.index] + self.theta_offset) # Sin previous angle
|
|
)
|
|
ang_term4 = C_GRAVITY * math.sin(self.theta[self.index] + self.theta_offset)
|
|
|
|
# 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.v_ang[self.index] # Previous velocity
|
|
+ (self.a_ang[self.index + 1] * (dt / 1000))
|
|
)
|
|
|
|
# Angular displacement
|
|
theta_res = self.theta[self.index] + (self.v_ang[self.index + 1] * (dt / 1000))
|
|
theta_round = theta_res % C_ANG_RES_LIMIT
|
|
if theta_round > 0.5 * C_ANG_RES_LIMIT:
|
|
theta_res = theta_res + theta_round
|
|
else:
|
|
theta_res = theta_res - theta_round
|
|
self.theta.append(theta_res)
|
|
|
|
# Limit fall of pendulum
|
|
self.theta[self.index + 1] = self.clamp(
|
|
self.theta[self.index + 1], -C_FALL_ANG, C_FALL_ANG
|
|
)
|
|
|
|
### 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
|
|
a_res = (-cart_term1 + cart_term2) / (2 * self.mass)
|
|
self.a_cart.append(self.clamp(a_res, -C_A_CART_LIMIT, C_A_CART_LIMIT))
|
|
# if abs(a_res) < C_A_CART_LIMIT:
|
|
# self.a_cart.append(a_res)
|
|
# else:
|
|
# self.a_cart.append(C_A_CART_LIMIT * (a_res / abs(a_res)))
|
|
|
|
# Integrate acceleration to get velocity
|
|
v_res = self.v_cart[self.index] + (self.a_cart[self.index + 1] * (dt / 1000)) # Previous velocity
|
|
|
|
if abs(v_res) < C_V_CART_LIMIT:
|
|
self.v_cart.append(v_res)
|
|
else:
|
|
self.v_cart.append(C_V_CART_LIMIT * (v_res / abs(v_res)))
|
|
|
|
# Cart displacement
|
|
self.s_cart.append(
|
|
self.s_cart[self.index] # Previous displacement
|
|
+ (self.v_cart[self.index + 1] * (dt / 1000))
|
|
)
|
|
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.index = 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.theta_offset = 0
|
|
self.fallen = False
|
|
self.update(0)
|
|
|
|
def plot(self):
|
|
fig, axs = plt.subplots(2, 2)
|
|
fig.suptitle("Pendulum")
|
|
|
|
axs[0,0].plot(self.theta)
|
|
axs[0,0].set_title('Angle [rad]')
|
|
axs[0,1].plot(self.v_ang)
|
|
axs[0,1].set_title('Angular velocity [rad/s]')
|
|
axs[1,0].plot(self.a_ang)
|
|
axs[1,0].set_title('Angular acceleration [rad/s^2]')
|
|
|
|
fig, axs = plt.subplots(2, 2)
|
|
fig.suptitle("Cart")
|
|
|
|
axs[0,0].plot(self.s_cart)
|
|
axs[0,0].set_title('Position [m]')
|
|
axs[0,1].plot(self.v_cart)
|
|
axs[0,1].set_title('Speed [m/s]')
|
|
axs[1,0].plot(self.a_cart)
|
|
axs[1,0].set_title('Acceleration [m/s^2]')
|
|
|
|
plt.show()
|
|
|
|
def pidControl(self, dt):
|
|
error = self.theta[self.index]
|
|
dt = (dt/1000)
|
|
# result = (self.kp * error) + (self.ki * sum(self.theta)) + (self.kd * self.v_ang[self.index]) #PID
|
|
result = (self.kp * error) \
|
|
+ (((error - self.theta[self.index - 1]) / (dt + 1 / 1000)) \
|
|
* self.kd) \
|
|
+ (self.kp_m * self.s_cart[self.index]) \
|
|
+ (self.kd_m * self.v_cart[self.index])
|
|
self.a_cart[self.index] = self.clamp(result * 10, -C_A_CART_LIMIT, C_A_CART_LIMIT) |