double_pendulum.py

# https://matplotlib.org/examples/animation/double_pendulum_animated.html
# Double pendulum formula translated from the C code at
# http://www.physics.usyd.edu.au/~wheat/dpend_html/solve_dpend.c

import numpy as np
import scipy.integrate as integrate
from numpy import cos, sin

G = 9.8  # acceleration due to gravity, in m/s^2
L1 = 1  # length of pendulum 1 in m
L2 = 1  # length of pendulum 2 in m
M1 = 1.0  # mass of pendulum 1 in kg
M2 = 1.0  # mass of pendulum 2 in kg


def derivs(state, t):

    dydx = np.zeros_like(state)
    dydx[0] = state[1]

    del_ = state[2] - state[0]
    den1 = (M1 + M2)*L1 - M2*L1*cos(del_)*cos(del_)
    dydx[1] = (M2*L1*state[1]*state[1]*sin(del_)*cos(del_) +
               M2*G*sin(state[2])*cos(del_) +
               M2*L2*state[3]*state[3]*sin(del_) -
               (M1 + M2)*G*sin(state[0]))/den1

    dydx[2] = state[3]

    den2 = (L2/L1)*den1
    dydx[3] = (-M2*L2*state[3]*state[3]*sin(del_)*cos(del_) +
               (M1 + M2)*G*sin(state[0])*cos(del_) -
               (M1 + M2)*L1*state[1]*state[1]*sin(del_) -
               (M1 + M2)*G*sin(state[2]))/den2

    return dydx

# th1 and th2 are the initial angles (degrees)
# w10 and w20 are the initial angular velocities (degrees per second)
# initial state
#state = np.radians([th1, w1, th2, w2])
#state = np.radians([120, 0, -10, 0])

def simulate(state):
    # create a time array from 0..100 sampled at 0.05 second steps
    dt = 0.001
    t = np.arange(0.0, 0.1, dt)
    
    # integrate your ODE using scipy.integrate.
    y = integrate.odeint(derivs, state, t)
    
    x1 = L1*sin(y[:, 0])
    y1 = -L1*cos(y[:, 0])
    
    x2 = L2*sin(y[:, 2]) + x1
    y2 = -L2*cos(y[:, 2]) + y1

    return (x1[-1],y1[-1],x2[-1],y2[-1],y[-1,:])