Springs work now

wheels.py

@@ -7,6 +7,7 @@ pygame.init()
 size = (800, 800)
 CENTER = np.array(size) / 2
 WHITE = (255, 255, 255)
+DAMPING = 0.99
 screen = pygame.display.set_mode(size)
 pygame.display.set_caption("Patrick")
 running = True
@@ -16,6 +17,8 @@ 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
 
@@ -23,22 +26,51 @@ class Ball(object):
         x, y = self.pos + CENTER
         pygame.draw.circle(screen, WHITE, [int(x), int(y)], self.radius, 2)
 
+    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]
+
+        self.pos += self.velocity
+
+        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, k=0.1):
+    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
+
 
 if __name__ == '__main__':
-    balls = [Ball(0, 0), Ball(100, 1)]
-    springs = [Spring(balls[0], balls[1])]
+    balls = [Ball(0., 0.), Ball(100., 100.)]
+    springs = [Spring(balls[0], balls[1], length=100)]
 
     while running:
         for event in pygame.event.get():
@@ -47,6 +79,14 @@ if __name__ == '__main__':
 
         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)