pendulum.py

# Double Pendulum Simulation
# Run with academy.cs.cmu sandbox

from math import *
import random

g = 0.78
dt = 0.008

app.stepsPerSecond = 60

def cos2(x):
    return cos(x)**2

def a0Step(a0, a1, p0, p1):
    top = 2*p0 - 3*cos(a0 - a1)*p1
    bot = 16 - 9*cos2(a0 - a1)
    return 6*top/bot
    
def a1Step(a0, a1, p0, p1):
    top = 8*p1 - 3*cos(a0 - a1)*p0
    bot = 16 - 9*cos2(a0 - a1)
    return 6*top/bot
    
def p0Step(a0, a1, a0d, a1d):
    return -0.5*(a0d*a1d*sin(a0 - a1) + 
        3*g*sin(a0))
        
def p1Step(a0, a1, a0d, a1d):
    return -0.5*(-a0d*a1d*sin(a0 - a1) + 
        g*sin(a1))

def eulerStep(state):
    a0, a1, p0, p1 = state
    a0d = a0Step(a0, a1, p0, p1)
    a1d = a1Step(a0, a1, p0, p1)
    p0d = p0Step(a0, a1, a0d, a1d)
    p1d = p1Step(a0, a1, a0d, a1d)
    return (a0 + dt*a0d, a1 + dt*a1d, p0 + dt*p0d, p1 + dt*p1d)
    
turb = random.randrange(100) / 100
app.stateA = [0.5*pi + turb, 0.5*pi, 0, 0]
app.stateB = app.stateA.copy()
app.stateB[0] += 0.02
app.stateC = app.stateB.copy()
app.stateC[0] += 0.02

line0 = Line(200, 200, 200, 275, fill='red')
line1 = Line(200, 275, 200, 350, fill='blue')
line2 = Line(200, 200, 200, 275, fill='green')
line3 = Line(200, 275, 200, 350, fill='orange')
line4 = Line(200, 200, 200, 275, fill='brown')
line5 = Line(200, 275, 200, 350, fill='skyBlue')

def stateToLines(state, lineA, lineB):
    lineA.y2 = 200 + 75*sin(0.5*pi + state[0])
    lineA.x2 = 200 + 75*cos(0.5*pi + state[0])
    lineB.x1 = lineA.x2
    lineB.y1 = lineA.y2
    lineB.y2 = lineB.y1 + 75*sin(0.5*pi + state[1])
    lineB.x2 = lineB.x1 + 75*cos(0.5*pi + state[1])
    
def onStep():
    stateToLines(app.stateA, line0, line1)
    stateToLines(app.stateB, line2, line3)
    stateToLines(app.stateC, line4, line5)
    for _x in range(0, 10):
        app.stateA = eulerStep(app.stateA)
        app.stateB = eulerStep(app.stateB)
        app.stateC = eulerStep(app.stateC)

def onMousePress(x, y):
    print(app.stateA)
    app.paused = not app.paused