pendulums.py

import numpy
import pygame
import random

pygame.init()

# Params
N = 1000            # Number of pendulums
L = 10              # Length of pendulum arms
M = 1               # Mass of pendulum arms
G = 9.8             # Gravity
FPS = 60

# Initial Conditions
THETA1 = random.random() * 6.283
THETA2 = random.random() * 6.283
SPREAD = (6.283 / 360) / 4

# Doing sub-steps in between the rendered frames helps
# prevent the Euler Approx from exploding so quickly
STEPS_PER_FRAME = 10

ML2 = M * L * L


class State:

    def __init__(self, n):
        self.theta1 = numpy.linspace(THETA1, THETA1 + SPREAD, n)
        self.theta2 = numpy.linspace(THETA2, THETA2 + SPREAD, n)
        self.dtheta1 = numpy.zeros(n)
        self.dtheta2 = numpy.zeros(n)

        self.p1 = numpy.zeros(n)  # these are zero (assuming that dthetas are zero)
        self.p2 = numpy.zeros(n)
        self.dp1 = numpy.zeros(n)
        self.dp2 = numpy.zeros(n)

    def fake_update(self, dt):
        self.theta1 = self.theta1 + 3.1415 * dt
        self.theta2 = self.theta2 + 2 * 3.1415 * dt

    def update(self, dt):
        old_dtheta1 = self.dtheta1
        old_dtheta2 = self.dtheta2
        cos_theta1_minus_theta2 = numpy.cos(self.theta1 - self.theta2)
        cos_theta1_minus_theta2_squared = cos_theta1_minus_theta2 * cos_theta1_minus_theta2
        sin_theta1_minus_theta2 = numpy.sin(self.theta1 - self.theta2)

        self.dtheta1 = (6 / ML2) * (2 * self.p1 - 3 * cos_theta1_minus_theta2 * self.p2) / (16 - 9 * cos_theta1_minus_theta2_squared)
        self.dtheta2 = (6 / ML2) * (8 * self.p2 - 3 * cos_theta1_minus_theta2 * self.p1) / (16 - 9 * cos_theta1_minus_theta2_squared)

        self.dp1 = (-0.5 * ML2) * (old_dtheta1 * old_dtheta2 * sin_theta1_minus_theta2 + 3 * G / L * numpy.sin(self.theta1))
        self.dp2 = (-0.5 * ML2) * (-old_dtheta1 * old_dtheta2 * sin_theta1_minus_theta2 + G / L * numpy.sin(self.theta2))

        self.theta1 = self.theta1 + dt * self.dtheta1
        self.theta2 = self.theta2 + dt * self.dtheta2
        self.p1 = self.p1 + dt * self.dp1
        self.p2 = self.p2 + dt * self.dp2

    def render_all(self, screen):
        zoom = 5
        x0 = screen.get_width() // 2
        y0 = screen.get_height() // 2
        x1 = x0 + zoom * L * numpy.cos(self.theta1 + 3.1429 / 2)
        y1 = y0 + zoom * L * numpy.sin(self.theta1 + 3.1429 / 2)
        x2 = x1 + zoom * L * numpy.cos(self.theta2 + 3.1429 / 2)
        y2 = y1 + zoom * L * numpy.sin(self.theta2 + 3.1429 / 2)

        for i in range(0, N):
            # TBH I'm not sure if aa is doing anything for us here
            pygame.draw.aalines(screen, (int(255 * i / N), 255, 255), False, [(x0, y0), (x1[i], y1[i]), (x2[i], y2[i])], blend=True)


def start():
    screen = pygame.display.set_mode((400, 300))
    canvas = screen.convert_alpha()
    clock = pygame.time.Clock()

    state = State(N)

    running = True
    ticks = 0
    while running:
        for _ in range(STEPS_PER_FRAME):
            state.update(1 / (FPS * STEPS_PER_FRAME))

        canvas.fill((0, 0, 0, 0))
        state.render_all(canvas)

        screen.fill((0, 0, 0))
        screen.blit(canvas, (0, 0))
        pygame.display.flip()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        if ticks % 60 == 10:
            pygame.display.set_caption("Double Pendulum (FPS={:.1f}, N={})".format(clock.get_fps(), N))

        ticks += 1
        clock.tick(FPS)


if __name__ == "__main__":
    start()