Complex

wheels.py

@@ -1,76 +1,57 @@
 # Import the pygame library and initialise the game engine
 import pygame
 import numpy as np
+from verlet import Ball, Spring
 pygame.init()
 
 # Initialize Pygame
-size = (800, 800)
-CENTER = np.array(size) / 2
+SIZE = (800, 800)
 WHITE = (255, 255, 255)
-DAMPING = 0.99
-screen = pygame.display.set_mode(size)
+screen = pygame.display.set_mode(SIZE)
 pygame.display.set_caption("Patrick")
 running = True
 clock = pygame.time.Clock()
 
 
-class Ball(object):
-    def __init__(self, x, y, radius=10, mass=1):
-        self.pos = np.array([x, y])
-        self.last_pos = np.array([x, y])
-        self.velocity = np.array([0, 0])
-        self.radius = radius
-        self.mass = mass
+def physics():
+    for ball in balls:
+        ball.update()
 
-    def render(self, screen):
-        x, y = self.pos + CENTER
-        pygame.draw.circle(screen, WHITE, [int(x), int(y)], self.radius, 2)
+    for spring in springs:
+        spring.update()
 
-    def update(self):
-        self.pos[1] += .05
 
-        if self.pos[1] > 400:
-            self.pos[1] = 400
-            self.last_pos[1] = 400
-            self.velocity[1] = -.5 * self.velocity[1]
+def draw():
+    for ball in balls:
+        ball.render(screen)
 
-        self.pos += self.velocity
+    for spring in springs:
+        spring.render(screen)
 
-        self.velocity = self.pos - self.last_pos
-        self.velocity *= DAMPING
-        self.last_pos = np.array(self.pos)
 
-
-class Spring(object):
-    def __init__(self, a, b, length=None, k=0.3):
-        self.a = a
-        self.b = b
-        self.k = k
-        self.length = self.compute_length() if length is None else length
-        print(self.length)
-
-    def render(self, screen):
-        x1, y1 = self.a.pos + CENTER
-        x2, y2 = self.b.pos + CENTER
-        pygame.draw.line(screen, WHITE, [x1, y1], [x2, y2])
-
-    def compute_length(self):
-        x1, y1 = self.a.pos + CENTER
-        x2, y2 = self.b.pos + CENTER
-        return np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
-
-    def update(self):
-        current_length = self.compute_length()
-        delta = self.b.pos - self.a.pos
-        delta /= current_length
-        stretch = (current_length - self.length) / self.length
-        self.a.velocity += stretch * delta * self.k
-        self.b.velocity -= stretch * delta * self.k
+def boost():
+    vector = springs[0].b.pos - springs[0].a.pos
+    length = springs[0].compute_length()
+    vector /= length
+    norm = np.array([vector[1], -vector[0]])
+    balls[1].velocity += norm * .5
 
 
 if __name__ == '__main__':
-    balls = [Ball(0., 0.), Ball(100., 100.)]
-    springs = [Spring(balls[0], balls[1], length=100)]
+    balls = [Ball(0., 0., fixed=True),
+             Ball(0., 100.),
+             Ball(20., 120.),
+             Ball(-20., 120.),
+             Ball(0., 80.)]
+    springs = [
+        Spring(balls[0], balls[1], length=100, k=3),
+        Spring(balls[1], balls[2], k=5),
+        Spring(balls[1], balls[3], k=5),
+        Spring(balls[2], balls[3], k=5),
+        Spring(balls[1], balls[4], k=5),
+        Spring(balls[2], balls[4], k=5),
+        Spring(balls[3], balls[4], k=5)
+    ]
 
     while running:
         for event in pygame.event.get():
@@ -78,20 +59,9 @@ if __name__ == '__main__':
                 running = False
 
         screen.fill((0, 0, 0))
-
-        # Update physics
-        for ball in balls:
-            ball.update()
-
-        for spring in springs:
-            spring.update()
-
-        # Draw
-        for ball in balls:
-            ball.render(screen)
-
-        for spring in springs:
-            spring.render(screen)
+        physics()
+        boost()
+        draw()
 
         pygame.display.flip()
         clock.tick(60)