double_pendulum.py

import math
from math import cos, sin

#### DEFINE THE mass, length and gravitational constant #####
#playing around with the gravity constant and mass ratios might cause a resonant system, which is pretty cool
m1,m2 = 2,1
L1,L2 = 1,2
g = 9.81
########


y = [ #this defined the y-functions. The entries here are the initial angles and initial momentum
math.pi/2,
math.pi/2,
0,
0
]
#########

def f(x,y): #warning: error prone stuff. This sets up the ODE's. Must be customarily set up for each problem...
    C1 = y[2]*y[3]*sin(y[0]-y[1])/(L1*L2*(m1+m2*sin(y[0]-y[1])**2))
    C2_numerator = L2*L2*m2*y[2]*y[2] + L1*L1*(m1+m2)*y[3]*y[3] - L1*L2*m2*y[2]*y[3]*cos(y[0]-y[1])
    C2_denominator = 2*L1*L1*L2*L2*(m1+m2*sin(y[0]-y[1])*sin(y[0]-y[1]))**2
    C2 = C2_numerator / C2_denominator * sin(2*(y[0]-y[1]))
        
    """This sets the f_entries"""
    num = L2*y[2] - L1*y[3]*cos(y[0]-y[1])
    den = L1*L1*L2*(m1+m2*sin(y[0]-y[1])*sin(y[0]-y[1]))
    yield num/den
    num = L1*(m1+m2)*y[3] - L2*m2*y[2]*cos(y[0]-y[1])
    den = L1*L2*L2*m2*(m1+m2*sin(y[0]-y[1])*sin(y[0]-y[1]))
    yield num/den
    yield -(m1+m2)*g*L1*sin(y[0]) - C1 + C2
    yield -m2*g*L2*sin(y[1]) + C1 - C2
    
def rk(h=0.2,N=5,x=0): #runge-kutta method for solving the ODE's.
    global y
    n = len(y) #number of funcitons
    for _ in range(N):
        a = [h*e for e in f(x,y)]
        b = [h*e for e in f(x+h/2, [y[j]+a[j]/2 for j in range(len(y))])]
        c = [h*e for e in f(x+h/2, [y[j]+b[j]/2 for j in range(len(y))])]
        d = [h*e for e in f(x+h, [y[j]+c[j] for j in range(len(y))])]        

        y = [y[i] + (a[i] + 2*b[i] + 2*c[i] + d[i])/6 for i in range(n)]
        x += h
        yield x,y

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import numpy as np
gen = rk(h=0.1,N=2000,x=0)
dL1, dL0,dL2,dL3 = [], [],[],[]
X = []
for x,y in gen:
    X.append(x)
    dL0.append(y[0])
    dL1.append(y[1])
    dL2.append(y[2])
    dL3.append(y[3])

plt.plot(X,dL0,"r2")
plt.plot(X,dL1,"b2")
plt.plot(X,dL2,"g1")
plt.plot(X,dL3,"y1")

fig = plt.figure()
ax = plt.axes(xlim=(-3,3), ylim=(-3,3))
pLine = Line2D([],[])
pend1, = ax.plot([],[],"r",lw=2)
pend2, = ax.plot([],[],"b",lw=2)

def init():
    pend1.set_data([],[])
    pend2.set_data([],[])
    return pend2,pend1,
def animate(i):
    global X, dL0, dL1, dL2, dL3
    pend1.set_data([0,L1*sin(dL0[i])],[0,-L1*cos(dL0[i])])
    pend2.set_data([pend1.get_data()[0][1], pend1.get_data()[0][1] +L2*sin(dL1[i])],[pend1.get_data()[1][1],pend1.get_data()[1][1]-L2*cos(dL1[i])])
    return pend1,pend2
from matplotlib import animation
ani = animation.FuncAnimation(fig, animate, init_func = init, frames = 2000, interval = 45, blit = True)
plt.show()