pendulum_visualizer.py

import pygame
from pygame.locals import *
from math import sin, cos
from time import sleep
import sys

class PendulumVisualizer:

    def __init__(self, t, z, y, l):
        pygame.init()
        self.window = pygame.display.set_mode((640, 480))
        pygame.display.set_caption("Pendulum")
        self.screen = pygame.display.get_surface()
        t_ = 0
        self.frame = 0
        print len(y)
        while True:
            self.input(pygame.event.get())
            t_ += 1
            if t_ >= len(y):
                sys.exit(0)
            self.draw_pendulum(z[t_], l, y[t_])
            sleep(0.033)

    def draw_pendulum(self, z, l, theta):
        plat_width = 200
        plat_height = 4
        scale = 50.0
        self.screen.fill((0,0,0))
        zpix = 240 - z * scale
        pygame.draw.rect(self.screen, (255, 255, 0), 
            pygame.Rect(320 - plat_width/2, zpix - plat_height/2, 
            plat_width, plat_height))
        pend_x = 320 + (l * scale) * sin(theta)
        pend_y = zpix + (l * scale) * cos(theta)
        pygame.draw.line(self.screen, (255, 255, 255), (320, zpix),
            (pend_x, pend_y), 2)
        pygame.draw.circle(self.screen, (255, 0, 0), (pend_x, pend_y), 10)
        #pygame.image.save(self.screen, "movie/%04d.png" % self.frame)
        self.frame += 1
        pygame.display.flip() 

    def input(self, events):
        for event in events:
            if event.type == QUIT:
                sys.exit(0)

if __name__ == "__main__":
    PendulumVisualizer()