EV5_Modcon/src/sim/pendulum.py

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


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
        """
        self.vector = None  # Vector2 object

        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.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

        self.pid = False
        self.fallen = False

    def update(self, dt):
        self.doMath(dt)
        self.vector = Vector2.from_polar(
            ((self.length * C_MTPRATIO), math.degrees(self.theta[self.index] + (1.5 * math.pi)))
        )

        if abs(self.theta[self.index]) == C_FALL_ANG:
            self.fallen = True

    def doMath(self, dt):
        ### 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])

        # 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
        self.theta.append(
            self.theta[self.index]  # Previous angle
            + (self.v_ang[self.index + 1] * (dt / 1000))
        )

        # 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
        self.a_cart.append((-cart_term1 + cart_term2) / (2 * self.mass))

        # Integrate acceleration to get velocity
        self.v_cart.append(
            self.v_cart[self.index]  # Previous velocity
            + (self.a_cart[self.index + 1] * (dt / 1000))
        )

        # 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.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_ang)
        axs[0,1].set_title('Speed [m/s]')
        axs[1,0].plot(self.a_ang)
        axs[1,0].set_title('Acceleration [m/s^2]')
        
        plt.show()

    # 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