runge_kutta.py

import matplotlib.pyplot as plt
import numpy as np

C_MOLA: np.float64 = 0.15  # comprimento da mola relaxada
K_MOLA: np.float64 = 1.0  # constante da mola
L_MOLA: np.float64 = 0.05  # local onde a mola foi posicionada nos pêndulos
G: np.float64 = 9.78  # valor da gravidade


# força feita no pêndulo da esquerda
def forca(t: np.float64,
          theta: np.float64,
          dtheta: np.float64,
          comp: np.float64,
          massa: np.float64,
          x2: np.float64) -> np.float64:
    
    x = L_MOLA * np.sin(theta)
    const = -1.0 if np.abs(x - x2) < C_MOLA else 1.0

    return ((-massa * G * np.sin(theta) * comp) + const * (K_MOLA * np.abs(x - x2) * np.abs(comp - L_MOLA))) / (
            massa * (comp * comp))


# força feita no pêndulo da direita
def forca2(t: np.float64,
           theta: np.float64,
           dtheta: np.float64,
           comp: np.float64,
           massa: np.float64,
           x2: np.float64) -> np.float64:
    
    x = L_MOLA * np.sin(theta)
    const = -1.0 if np.abs(x - x2) > C_MOLA else 1.0

    return ((-massa * G * np.sin(theta) * comp) + const * (K_MOLA * np.abs(x - x2) * (comp - L_MOLA))) / (
            massa * (comp * comp))


def angulo(t: np.float64, 
           theta: np.float64,
           dtheta: np.float64) -> np.float64:
    
    return dtheta


def create_graph(nums,
                 y1,
                 y2):
    
    plt.plot(nums, y1)
    plt.plot(nums, y2)
    plt.legend(['pend-1', 'pend-2'])
    plt.xlabel('iteracao')
    plt.ylabel('angulo')
    plt.title('Angulo x Tempo')
    plt.show()


if __name__ == '__main__':
    comp1: np.float64 = 0.10
    comp2: np.float64 = 0.15

    h: np.float64 = 0.10
    dtheta1, dtheta2, theta1, theta2 = [], [], [], []
    massa1: np.float64 = 1.0
    massa2: np.float64 = 1.0
    x = [x for x in range(0, 10000)]

    theta1.append(0.2618)
    theta2.append(0.0)
    dtheta1.append(0.0)
    dtheta2.append(0.0)

    for i in x[1:]:
        k1 = h * forca(i, theta1[i - 1], dtheta1[i - 1], comp1, massa1, comp2 * np.sin(theta2[0]) + 0.15)
        k2 = h * forca(i, theta1[i - 1] + h / 2.0, dtheta1[i - 1] + k1 / 2.0, comp1, massa1,
                       comp2 * np.sin(theta2[0]) + 0.15)
        k3 = h * forca(i, theta1[i - 1] + h / 2.0, dtheta1[i - 1] + k2 / 2.0, comp1, massa1,
                       comp2 * np.sin(theta2[0]) + 0.15)
        k4 = h * forca(i, theta1[i - 1] + h, dtheta1[i - 1] + k3, comp1, massa1, comp2 * np.sin(theta2[0]) + 0.15)

        dtheta1.append(dtheta1[i - 1] + (k1 + 2 * k2 + 2 * k3 + k4) / 6.0)

        k1 = h * angulo(i , theta1[i-1], dtheta1[i - 1])
        k2 = h * angulo(i, theta1[i - 1], dtheta1[i-1])
        k3 = h * angulo(i, theta1[i - 1], dtheta1[i-1])
        k4 = h * angulo(i, theta1[i-1], dtheta1[i - 1])

        theta1.append(theta1[i - 1] + (k1 + 2 * k2 + 2 * k3 + k4) / 6.0)

    for i in x[1:]:
        k1 = h * forca(dtheta2[i - 1], theta2[i - 1], comp2, massa2, comp1 * np.sin(theta1[0]) + 0.15)
        k2 = h * forca(dtheta2[i - 1] + h / 2.0, theta2[i - 1] + k1 / 2.0, comp2, massa2,
                       comp1 * np.sin(theta1[0]) + 0.15)
        k3 = h * forca(dtheta2[i - 1] + h / 2.0, theta2[i - 1] + k2 / 2.0, comp2, massa2,
                       comp1 * np.sin(theta1[0]) + 0.15)
        k4 = h * forca(dtheta2[i - 1] + h, theta2[i - 1] + k3, comp2, massa2, comp1 * np.sin(theta1[0]) + 0.15)

        dtheta2.append(dtheta2[i - 1] + (k1 + 2 * k2 + 2 * k3 + k4) / 6.0)

        k1 = h * angulo(dtheta2[i - 1])
        k2 = h * angulo(dtheta2[i - 1])
        k3 = h * angulo(dtheta2[i - 1])
        k4 = h * angulo(dtheta2[i - 1])

        theta2.append(theta2[i - 1] + (k1 + 2 * k2 + 2 * k3 + k4) / 6.0)

    x = np.array(x, np.int32)
    theta1 = np.array(theta1, np.float64)
    theta2 = np.array(theta2, np.float64)

    create_graph(x, theta1, theta2)