These are simple pygame based pendulum physics simulations.

README.md

auth: Craig Wm. Versek ([email protected]) dates: gist created 2016/06/09 based directly from code written circa 2008 while I was a PhD candidate at UMass Amherst

These are simple pygame based pendulum physics simulations that started with a simple idea: perform the numerical integration (the "physics") within an infinite generator loop, supplying the parameters to the graphics rendering code on demand. The thought was to separate the number crunching from the application at large and simplify its organization and portability. The code can be configured to use one of a few numerical integration methods by swapping generators with the same interface. pygame is the only dependency that needs to be installed.

phys_gen.py illustrates the basic concept with no bells & whistles (nor graphics).

pendulum.py is a pygame driven simple pendulum interactive simulation -- try clicking and dragging the bob. Can use Heun, Steormer_Verlet, or RK4 methods.

doubled_pedulum.py has two double pendula running in tandem with very close but not equal initial conditions. Can use RK4 or Steormer_Verlet methods. Note, the physical system is chaotic and should conserve energy, but numerical integration errors build up and the behavior tends towards a mode of the small angle approximation. Some day this needs to be put right!

double_pendulum.py

###############################################################################
"""
doublependulum.py

desc: A pygame animated simulation of two tandem double pendula.

auth: Craig Wm. Versek ([email protected]) circa 2008?
"""
###############################################################################
import sys,os
from math import sin,cos,pi,atan2,sqrt
import pygame
from pygame.locals import *
import pygame.surfarray

COLOR = {'black' : (0,0,0),
         'red'   : (255,0,0),
         'green' : (0,255,0),
         'blue'  : (0,0,255)}

SCREEN_WIDTH  = 1200
SCREEN_HEIGHT = 800
SCREEN_DIM    = (SCREEN_WIDTH,SCREEN_HEIGHT)
SCREEN_CENTER = (SCREEN_WIDTH//2,SCREEN_HEIGHT//2) 

def gen_doublependulum_physics_RK4(dt,theta1_0,theta2_0,theta1_dot0,theta2_dot0,m1,m2,l1,l2,g):
    #aggregate some constants
    M     = float(m1 + m2)
    L1L2  = float(l1*l2)
    L1sqr = float(l1*l1)
    L2sqr = float(l2*l2)
    L1M   = float(l1*M)
    L2M2  = float(l2*m2)
    GL1M  = float(g*L1M)
    GM2L2 = float(g*m2*l2)
    L2sqrM2 = float(L2sqr*m2)
    L1sqrM  = float(L1sqr*M)
    L1L2M2  = float(L1L2*m2)
    #initialize generalized coordinates
    q1        = theta1_0
    q2        = theta2_0
    p1        = L1sqrM*theta1_dot0 + L1L2M2*theta2_dot0*cos(theta1_0 - theta2_0)
    p2        = L2sqrM2*theta2_dot0 + L1L2M2*theta1_dot0*cos(theta1_0 - theta2_0)
    #the Equations of Motion
    def EOM(q1,q2,p1,p2):
        #compute common parts
        delta_q = q1 - q2
        sin_delq = sin(delta_q)
        cos_delq = cos(delta_q)
        A = L1L2*(m1 + m2*sin_delq*sin_delq)
        B = p1*p2*sin_delq/A
        C = sin(2*delta_q)*(L2sqrM2*p1*p1 + L1sqrM*p2*p2 - L1L2M2*p1*p2*cos_delq)/(2*A*A)
        #equations of motion
        q1_dot = (l2*p1 - l1*p2*cos_delq)/(l1*A)
        q2_dot = (L1M*p2 - L2M2*p1*cos_delq)/(l2*A)
        p1_dot = -GL1M*sin(q1)  - B + C
        p2_dot = -GM2L2*sin(q2) + B - C    
        return (dt*q1_dot,dt*q2_dot,dt*p1_dot,dt*p2_dot)
    #loop forever
    while True:
        #integrate using Runge-Kutta 4th order
        k1_1,k2_1,n1_1,n2_1 = EOM(q1         ,q2         ,p1         ,p2         )
        k1_2,k2_2,n1_2,n2_2 = EOM(q1+k1_1/2.0,q2+k2_1/2.0,p1+n1_1/2.0,p2+n2_1/2.0)
        k1_3,k2_3,n1_3,n2_3 = EOM(q1+k1_2/2.0,q2+k2_2/2.0,p1+n1_2/2.0,p2+n2_2/2.0)
        k1_4,k2_4,n1_4,n2_4 = EOM(q1+k1_3    ,q2+k2_3    ,p1+n1_3    ,p2+n2_3    )
        q1 += ((k1_1+k1_4)/2.0 + k1_2+k1_3)/3.0
        q2 += ((k2_1+k2_4)/2.0 + k2_2+k2_3)/3.0
        p1 += ((n1_1+n1_4)/2.0 + n1_2+n1_3)/3.0
        p2 += ((n2_1+n2_4)/2.0 + n2_2+n2_3)/3.0
        yield (q1,q2)

def gen_doublependulum_physics_Steomer_Verlet(dt,
                                              theta1_0,
                                              theta2_0,
                                              theta1_dot0,
                                              theta2_dot0,
                                              m1,m2,l1,l2,g):
    #aggregate some constants
    M     = float(m1 + m2)
    L1L2  = float(l1*l2)
    L1sqr = float(l1*l1)
    L2sqr = float(l2*l2)
    L1M   = float(l1*M)
    L2M2  = float(l2*m2)
    GL1M  = float(g*L1M)
    GM2L2 = float(g*m2*l2)
    L2sqrM2 = float(L2sqr*m2)
    L1sqrM  = float(L1sqr*M)
    L1L2M2  = float(L1L2*m2)
    #initialize generalized coordinates
    q1        = theta1_0
    q2        = theta2_0
    p1        = L1sqrM*theta1_dot0 + L1L2M2*theta2_dot0*cos(theta1_0 - theta2_0)
    p2        = L2sqrM2*theta2_dot0 + L1L2M2*theta1_dot0*cos(theta1_0 - theta2_0)
    #the Equations of Motion
    def EOM(q1,q2,p1,p2):
        #compute common parts
        delta_q = q1 - q2
        sin_delq = sin(delta_q)
        cos_delq = cos(delta_q)
        A = L1L2*(m1 + m2*sin_delq*sin_delq)
        B = p1*p2*sin_delq/A
        C = sin(2*delta_q)*(L2sqrM2*p1*p1 + L1sqrM*p2*p2 - L1L2M2*p1*p2*cos_delq)/(2*A*A)
        #equations of motion
        q1_dot = (l2*p1 - l1*p2*cos_delq)/(l1*A)
        q2_dot = (L1M*p2 - L2M2*p1*cos_delq)/(l2*A)
        p1_dot = -GL1M*sin(q1)  - B + C
        p2_dot = -GM2L2*sin(q2) + B - C
        return (dt*q1_dot,dt*q2_dot,dt*p1_dot,dt*p2_dot)
    #loop forever
    while True:
        #integrate using Runge-Kutta 4th order
        k1_1,k2_1,n1_1,n2_1 = EOM(q1         ,q2         ,p1         ,p2         )
        k1_2,k2_2,n1_2,n2_2 = EOM(q1+k1_1/2.0,q2+k2_1/2.0,p1+n1_1/2.0,p2+n2_1/2.0)
        k1_3,k2_3,n1_3,n2_3 = EOM(q1+k1_2/2.0,q2+k2_2/2.0,p1+n1_2/2.0,p2+n2_2/2.0)
        k1_4,k2_4,n1_4,n2_4 = EOM(q1+k1_3    ,q2+k2_3    ,p1+n1_3    ,p2+n2_3    )
        q1 += k1_1
        q2 += ((k2_1+k2_4)/2.0 + k2_2+k2_3)/3.0
        p1 += ((n1_1+n1_4)/2.0 + n1_2+n1_3)/3.0
        p2 += ((n2_1+n2_4)/2.0 + n2_2+n2_3)/3.0
        yield (q1,q2)


class DoublePendulum(pygame.sprite.Sprite):
    """renders a fixed pivot pendulum and updates motion according to differential equation"""
    def __init__(self,pivot_vect=SCREEN_CENTER,
                      length1=200,length2=100,
                      bob_radius1=20,bob_radius2=10,
                      bob_mass1=1,  bob_mass2=0.5,
                      init_angle1=pi/4,init_angularspeed1=0,
                      init_angle2=pi/4,init_angularspeed2=0,
                      gravity=9.8,dt=0.01):
        pygame.sprite.Sprite.__init__(self) #call Sprite initializer
        self.phys_gen = gen_doublependulum_physics_Steomer_Verlet(dt,
                                                                  init_angle1,
                                                                  init_angle2,
                                                                  init_angularspeed1,
                                                                  init_angularspeed2,
                                                                  bob_mass1,bob_mass2,
                                                                  length1/1000.0,length2/1000.0,
                                                                  gravity)
        self.length     = (length1    ,length2    )
        self.bob_radius = (bob_radius1,bob_radius2)
        self.bob_mass   = (bob_mass1  ,bob_mass2  )
        self.angle      = (init_angle1,init_angle2)
        #positioning attributes
        self.pivot_vect = pivot_vect                     #vector from topleft to pivot of pendulum
        swinglen = length1 + length2 + bob_radius2       #whole swing arc
        #these next two attributes are used by the RenderPlain container class
        self.image  = pygame.Surface((swinglen*2,swinglen*2)).convert()          #create surface just big enough to fit swing
        self.image.set_colorkey(COLOR['black'])
        self.rect   = self.image.get_rect()
        self.rect.topleft = (pivot_vect[0] - swinglen, pivot_vect[1] - swinglen) #place so that pivot is at center
        self.image_center = (self.rect.width//2,self.rect.height//2)
        #calculate the initial relative bob position in the image
        self.bob1_X    = int(length1*sin(init_angle1)+self.rect.width//2)
        self.bob1_Y    = int(length1*cos(init_angle1)+self.rect.height//2)
        self.bob2_X    = int(length2*sin(init_angle2)+self.bob1_X)
        self.bob2_Y    = int(length2*cos(init_angle2)+self.bob1_Y)
        #render the pendulum from the parameters
        self._render()
        
    def _render(self):
        #clear the pendulum surface
        self.image.fill(COLOR['black'])
        bob1_pos = (self.bob1_X,self.bob1_Y)
        bob2_pos = (self.bob2_X,self.bob2_Y)
        bob_radius1, bob_radius2 = self.bob_radius
        #draw the tethers
        pygame.draw.aaline(self.image,COLOR['red'],self.image_center,bob1_pos,True)
        pygame.draw.aaline(self.image,COLOR['red'],bob1_pos         ,bob2_pos,True)
        #draw the bob2
        pygame.draw.circle(self.image,COLOR['blue'] ,bob1_pos,bob_radius1,0)
        pygame.draw.circle(self.image,COLOR['green'],bob2_pos,bob_radius2,0)
        
    def update(self):
        #coords relative to pivot
        self.angle = self.phys_gen.next()
        angle1 ,angle2  = self.angle
        length1,length2 = self.length
        self.bob1_X = int(length1*sin(angle1)) + self.rect.width//2
        self.bob1_Y = int(length1*cos(angle1)) + self.rect.height//2
        self.bob2_X = int(length2*sin(angle2)  + self.bob1_X)
        self.bob2_Y = int(length2*cos(angle2)  + self.bob1_Y)
        self._render()


def main():
    """this function is called when the program starts.
    it initializes everything it needs, then runs in
    a loop until a stop event (escape or window closing) is recognized.
    """
    #Initialize Everything
    pygame.init()
    screen = pygame.display.set_mode(SCREEN_DIM)
    pygame.display.set_caption('Double Pendulum Simulation')
    #pygame.mouse.set_visible(0)

    #Create The Backgound
    background = pygame.Surface(screen.get_size())
    background = background.convert()
    background.fill(COLOR['black'])
    #Prepare Objects
    clock = pygame.time.Clock()
    p1 = DoublePendulum(pivot_vect=(  SCREEN_WIDTH//4,SCREEN_HEIGHT//2),
                        init_angle1=pi,init_angle2=pi+0.001,dt=0.01,
                        #bob_mass1=0.5,bob_radius1=10,
                        #bob_mass2=0.5,bob_radius2=1
                       )
    p2 = DoublePendulum(pivot_vect=(3*SCREEN_WIDTH//4,SCREEN_HEIGHT//2),
                        init_angle1=pi,init_angle2=pi+0.002,dt=0.01)
    free_group = pygame.sprite.RenderPlain((p1,p2))
    #Display The Background
    screen.blit(background, (0, 0))
    pygame.display.flip()

    #Main Loop
    while True:
        clock.tick(30)
        #Handle Input Events
        for event in pygame.event.get():
            if event.type == QUIT:
                return
            elif event.type == KEYDOWN and event.key == K_ESCAPE:
                return
        free_group.update()
        #send the mouse position to the held bobs so we can move them
        screen.blit(background,(0,0))
        free_group.draw(screen)
        pygame.display.flip()

if __name__ == "__main__":
    main()

pendulum.py

###############################################################################
"""
pendulum.py

desc: A mildly interactive pygame animated simulation of circular physical
      pendulum.

auth: Craig Wm. Versek ([email protected]) circa 2008?
"""
###############################################################################
import sys,os
from math import sin,cos,pi,atan2,sqrt
import pygame
from pygame.locals import *
import pygame.surfarray

COLOR = {'black' : (0,0,0),
         'red'   : (255,0,0),
         'green' : (0,255,0),
         'blue'  : (0,0,255)}

SCREEN_WIDTH  = 800
SCREEN_HEIGHT = 800
SCREEN_DIM    = (SCREEN_WIDTH,SCREEN_HEIGHT)
SCREEN_CENTER = (SCREEN_WIDTH//2,SCREEN_HEIGHT//2) 
DEFAULT_INTEGRATION_METHOD = 'Steormer-Verlet'

def gen_pendulum_physics_Heun(dt,theta0,theta_dot0,m,g,l):
    a = float(-m*g*l)
    b = float(m*l*l)
    #initialize generalized coordinates
    q = theta0        
    p = b*theta_dot0
    while True:
        #Update using Heun's Method
        q_dot1 = p/b
        p_dot1 = a*sin(q)
        q_dot2 = (p+dt*p_dot1)/b
        p_dot2 = a*sin(q+dt*q_dot1)
        q += dt/2.0*(q_dot1 + q_dot2)
        p += dt/2.0*(p_dot1 + p_dot2)
        yield q

def gen_pendulum_physics_Steormer_Verlet(dt,theta0,theta_dot0,m,g,l):
    a = float(-m*g*l)
    b = float(m*l*l)
    #initialize generalized coordinates
    q = theta0        
    p = b*theta_dot0
    #this is the force pulling back, note 'a' should be negative 
    f = lambda x: a*sin(x)
    #integrate via the leap-frog method
    while True:
        q_dot = p + 0.5*dt*f(q)
        q += dt*q_dot
        p  = q_dot + 0.5*dt*f(q)  #note that q has changed in between
        yield (q, q_dot)


def gen_pendulum_physics_RK4(dt,theta0,theta_dot0,m,g,l):
    #agregate the physical constants    
    a = 1.0/(m*l*l)    
    b = float(-m*g*l)
    #initialize generalized coordinates
    q = theta0        
    p = theta_dot0/a
    #this is the physics! ... Hamiltonian style
    q_dot = lambda p_arg: a*p_arg        
    p_dot = lambda q_arg: b*sin(q_arg)  
    #Integrate using Runge-Kutta 4th Order Method
    while True:
        k1 = dt*q_dot(p)
        h1 = dt*p_dot(q)        
        k2 = dt*q_dot(p + h1/2.0)
        h2 = dt*p_dot(q + k1/2.0)
        k3 = dt*q_dot(p + h2/2.0)
        h3 = dt*p_dot(q + k2/2.0)
        k4 = dt*q_dot(p + h3)
        h4 = dt*p_dot(q + k3)
        q += ((k1+k4)/2.0 + k2+k3)/3.0
        p += ((h1+h4)/2.0 + h2+h3)/3.0
        yield (q,q_dot(p))


PHYSICS_GENERATORS = {
    'Heun':            gen_pendulum_physics_Heun,
    'Steormer-Verlet': gen_pendulum_physics_Steormer_Verlet,
    'Runge-Kutta 4':   gen_pendulum_physics_RK4,
}

class Pendulum(pygame.sprite.Sprite):
    """renders a fixed pivot pendulum and updates motion according to differential equation"""
    def __init__(self,pivot_vect,length,bob_radius,bob_mass,init_angle, integration_method = DEFAULT_INTEGRATION_METHOD):
        pygame.sprite.Sprite.__init__(self) #call Sprite initializer
        phys_gen_init = PHYSICS_GENERATORS[integration_method]
        self.phys_gen = phys_gen_init(dt = 0.01,theta0 = init_angle,theta_dot0 = 0,m = bob_mass,g = 9.8,l = length/1000.0)        
        self.length = length
        self.bob_radius = bob_radius
        self.bob_mass   = bob_mass
        #deflection from plumb
        self.angle      = init_angle
        #positioning attributes
        self.pivot_vect = pivot_vect         #vector from topleft to pivot of pendulum
        swinglen = (length + bob_radius)     #whole
        #these next two attributes are used by the RenderPlain container class
        self.image  = pygame.Surface((swinglen*2,swinglen*2)).convert()          #create surface just big enough to fit swing
        self.rect   = self.image.get_rect()
        self.rect.topleft = (pivot_vect[0] - swinglen, pivot_vect[1] - swinglen) #place so that pivot is at center 
        #calculate the initial relative bob position in the image
        self.bob_X    = int(length*sin(init_angle)+self.rect.width//2)
        self.bob_Y    = int(length*cos(init_angle)+self.rect.height//2)
        self.bob_rect = None
        #render the pendulum from the parameters
        self._render()
    def _render(self):
        #clear the pendulum surface
        self.image.fill(COLOR['black'])
        bob_pos = (self.bob_X,self.bob_Y)
        #draw the tether
        pygame.draw.aaline(self.image,COLOR['red'],(self.rect.width//2,self.rect.height//2),bob_pos,True)
        #draw the bob
        self.bob_rect = pygame.draw.circle(self.image,COLOR['blue'],bob_pos,self.bob_radius,0)
        x1,y1 = self.bob_rect.topleft
        x2,y2 = self.rect.topleft
        #make the reference absolute
        self.bob_rect.topleft = (x1+x2,y1+y2)  #make the reference absolute
    def update(self):
        #coords relative to pivot
        self.angle = self.phys_gen.next()[0]
        angle  = self.angle
        length = self.length
        X = int(length*sin(angle))
        Y = int(length*cos(angle))
        self.bob_X = X + self.rect.width//2
        self.bob_Y = Y + self.rect.height//2
        self._render()

    def update_held(self,mouse_pos):
        m_x,m_y = mouse_pos
        self.bob_X = m_x - self.lever_x - self.bob_rect.width//2
        self.bob_Y = m_y - self.lever_y - self.bob_rect.height//2
        self.length = sqrt((self.bob_X - self.rect.width//2)**2 + (self.bob_Y - self.rect.height//2)**2)
        #calculate the new parameters        
        self.angle = atan2(self.bob_X - self.rect.width//2, self.bob_Y - self.rect.height//2)
        self._render()

    def grab(self,mouse_pos):
        m_x, m_y = mouse_pos
        b_x, b_y = self.bob_rect.center
        self.held = True
        self.lever_x = m_x - b_x
        self.lever_y = m_y - b_y
        print "Grabbed the bob a vector from center (%d,%d)" % (self.lever_x,self.lever_y)

    def point_on_bob(self,point):
        x,y = point
        return self.bob_rect.collidepoint(x,y)

    def release(self):
        #create new physics by dropping pendulum at current angle
        self.phys_gen = gen_pendulum_physics_RK4(0.01,self.angle,0,self.bob_mass,9.8,self.length/1000.0)
   

def main():
    """this function is called when the program starts.
    it initializes everything it needs, then runs in
    a loop until the function returns.
    """
    #Initialize Everything
    pygame.init()
    screen = pygame.display.set_mode(SCREEN_DIM)
    pygame.display.set_caption('Pendulum Simulation')
    #pygame.mouse.set_visible(0)

    #Create The Backgound
    background = pygame.Surface(screen.get_size())
    background = background.convert()
    background.fill(COLOR['black'])
    #Prepare Objects
    clock = pygame.time.Clock()
    p1 = Pendulum(pivot_vect=SCREEN_CENTER,length=300,bob_radius=50,bob_mass=1,init_angle=pi/5)
    free_group = pygame.sprite.RenderPlain((p1,))
    held_group = pygame.sprite.RenderPlain()
    #Display The Background
    screen.blit(background, (0, 0))
    pygame.display.flip()

    #Main Loop
    while True:
        clock.tick(60)
        #Handle Input Events
        for event in pygame.event.get():
            if event.type == QUIT:
                return
            elif event.type == KEYDOWN and event.key == K_ESCAPE:
                return
            elif event.type == MOUSEBUTTONDOWN:
                print "Mouse Button Down"
                mouse_pos = pygame.mouse.get_pos()
                for p in free_group:
                    if p.point_on_bob(mouse_pos):     #if user clicked on the bob grab it
                        p.grab(mouse_pos)
                        held_group.add(p)
                        free_group.remove(p)
            elif event.type is MOUSEBUTTONUP:
                print "Mouse Button Up"
                for p in held_group:
                    p.release()
                    free_group.add(p)
                    held_group.remove(p)
        free_group.update()
        #send the mouse position to the held bobs so we can move them
        mouse_pos = pygame.mouse.get_pos()
        for p in held_group:
            p.update_held(mouse_pos)
        screen.blit(background,(0,0))
        free_group.draw(screen)
        held_group.draw(screen)
        #pygame.draw.circle(screen, COLOR['blue'], SCREEN_CENTER, 50, 0)
        pygame.display.flip()

if __name__ == "__main__":
    main()

phys_gen.py

import time
import itertools
from math import *

def gen_pendulum_physics(dt,theta0,theta_dot0,m,g,l):
    a = float(-m*g*l)
    b = float(m*l*l)
    #initialize generalized coordinates
    q = theta0        
    p = b*theta_dot0
    while True:
        #Update using Heun's Method
        q_dot1 = p/b
        p_dot1 = a*sin(q)
        q_dot2 = (p+dt*p_dot1)/b
        p_dot2 = a*sin(q+dt*q_dot1)
        q += dt/2.0*(q_dot1 + q_dot2)
        p += dt/2.0*(p_dot1 + p_dot2)
        yield q



#TODO: add a real-time delay to time generator for simulation
def time_flow(dt=1,start=0.0):
    for i in itertools.count():	
        yield start + dt*i
        time.sleep(dt)

def func(f,i):
    for x in i:
        yield f(x)

def diff(i,dt=1.0):
    x0 = i.next()
    for x1 in i:
        yield (x1 - x0)/dt
        x0 = x1

dt = .05
start = 0
steps = 100
#t = time_flow(dt)
#x = func(lambda v: sin(v)+v,t)
#dx = diff(x,dt)
pend = gen_pendulum_physics(dt,pi/4,0,1,9.8,1)
for val in itertools.islice(pend,start,steps):
    print val