diff --git a/wheels.py b/wheels.py index b002cd4..bfb08b4 100644 --- a/wheels.py +++ b/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)