Verlet integrator, to deal with an argument from the pub.

verlet.py 1.9KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. """
  2. Implements Verlet integration
  3. """
  4. import pygame
  5. import numpy as np
  6. SIZE = (800, 800)
  7. CENTER = np.array(SIZE) / 2
  8. WHITE = (255, 255, 255)
  9. DAMPING = 0.99
  10. class Ball(object):
  11. def __init__(self, x, y, radius=5, mass=1, fixed=False):
  12. self.pos = np.array([x, y])
  13. self.last_pos = np.array([x, y])
  14. self.velocity = np.array([0., 0.])
  15. self.radius = radius
  16. self.mass = mass
  17. self.fixed = fixed
  18. def render(self, screen):
  19. x, y = self.pos + CENTER
  20. pygame.draw.circle(screen, WHITE, [int(x), int(y)], self.radius, 4)
  21. def update(self):
  22. if self.fixed:
  23. return
  24. # Gravity
  25. #self.pos[1] += .05
  26. # Bounce
  27. if self.pos[1] > 400:
  28. self.pos[1] = 400
  29. self.last_pos[1] = 400
  30. self.velocity[1] = -.5 * self.velocity[1]
  31. self.pos += self.velocity
  32. self.velocity = self.pos - self.last_pos
  33. self.velocity *= DAMPING
  34. self.last_pos = np.array(self.pos)
  35. class Spring(object):
  36. def __init__(self, a, b, length=None, k=0.3, color=WHITE):
  37. self.a = a
  38. self.b = b
  39. self.k = k
  40. self.length = self.compute_length() if length is None else length
  41. self.color = color
  42. print(self.length)
  43. def render(self, screen):
  44. x1, y1 = self.a.pos + CENTER
  45. x2, y2 = self.b.pos + CENTER
  46. pygame.draw.line(screen, self.color, [x1, y1], [x2, y2])
  47. def compute_length(self):
  48. x1, y1 = self.a.pos + CENTER
  49. x2, y2 = self.b.pos + CENTER
  50. return np.sqrt((x2 - x1)**2 + (y2 - y1)**2)
  51. def update(self):
  52. current_length = self.compute_length()
  53. delta = self.b.pos - self.a.pos
  54. delta /= current_length
  55. stretch = (current_length - self.length) / self.length
  56. self.a.velocity += stretch * delta * self.k
  57. self.b.velocity -= stretch * delta * self.k