simulate and plot a double pendulum with rk4 in python

main.py
import numpy as np
import matplotlib.pyplot as plt

# define the parameters of the double pendulum
m1 = 1    # mass of the first pendulum
m2 = 1    # mass of the second pendulum
L1 = 1    # length of the first pendulum
L2 = 2    # length of the second pendulum
g = 9.81  # acceleration due to gravity

# define the equations of motion for the double pendulum
def dydt(t, y):
    theta1, omega1, theta2, omega2 = y
    dtheta1_dt = omega1
    domega1_dt = (-g*(2*m1+m2)*np.sin(theta1) - m2*g*np.sin(theta1-2*theta2) - 2*np.sin(theta1-theta2)*m2*(omega2**2*L2 + omega1**2*L1*np.cos(theta1-theta2))) / (L1*(2*m1 + m2 - m2*np.cos(2*theta1-2*theta2)))
    dtheta2_dt = omega2
    domega2_dt = (2*np.sin(theta1-theta2)*(omega1**2*L1*(m1+m2) + g*(m1+m2)*np.cos(theta1) + omega2**2*L2*m2*np.cos(theta1-theta2))) / (L2*(2*m1 + m2 - m2*np.cos(2*theta1-2*theta2)))
    return [dtheta1_dt, domega1_dt, dtheta2_dt, domega2_dt]

# define the RK4 method
def rk4(t0, y0, dt, t_end):
    t = []
    y = []
    t.append(t0)
    y.append(y0)
    while t[-1] < t_end:
        k1 = dt * np.array(dydt(t[-1], y[-1]))
        k2 = dt * np.array(dydt(t[-1] + dt/2, np.array(y[-1]) + np.array(k1)/2))
        k3 = dt * np.array(dydt(t[-1] + dt/2, np.array(y[-1]) + np.array(k2)/2))
        k4 = dt * np.array(dydt(t[-1] + dt, np.array(y[-1]) + np.array(k3)))
        y_new = np.array(y[-1]) + (1/6)*(k1 + 2*k2 + 2*k3 + k4)
        y.append(y_new)
        t.append(t[-1] + dt)
    return np.array(t), np.array(y)

# set the initial conditions
t0 = 0
y0 = [np.pi/2, 0, np.pi/3, 0]

# simulate the double pendulum using RK4
dt = 0.01    # time step
t_end = 10   # end time
t, y = rk4(t0, y0, dt, t_end)

# extract the positions of the two pendulums
x1 = L1*np.sin(y[:,0])
y1 = -L1*np.cos(y[:,0])
x2 = x1 + L2*np.sin(y[:,2])
y2 = y1 - L2*np.cos(y[:,2])

# plot the motion of the two pendulums
fig, ax = plt.subplots(figsize=(8,8))
ax.set_xlim((-L1-L2)*1.2, (L1+L2)*1.2)
ax.set_ylim((-L1-L2)*1.2, (L1+L2)*1.2)
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_title('Double Pendulum Motion')
line1, = ax.plot(x1, y1, 'r-', linewidth=2)
line2, = ax.plot(x2, y2, 'b-', linewidth=2)
plt.show()
2156 chars
61 lines