spring.py

from math import pi
import random
import pygame
import PyParticles

(width, height) = (400, 400)
screen = pygame.display.set_mode((width, height))
pygame.display.set_caption('Springs')

universe = PyParticles.Environment((width, height))
universe.colour = (255,255,255)
universe.addFunctions(['move', 'bounce', 'collide', 'drag', 'accelerate'])
universe.acceleration = (pi, 0.01)
universe.mass_of_air = 0.02

for p in range(3):
    universe.addParticles(mass=100, size=16, speed=2, elasticity=1, colour=(20,40,200))
    #universe.addParticles(mass=40, size=10, speed=1, elasticity=1, colour=(20,200,40)

universe.addSpring(0, 1, length=100, strength=0.05)
universe.addSpring(1, 2, length=100, strength=0.05)
universe.addSpring(2, 0, length=100, strength=0.05)

selected_particle = None
paused = False
running = True



def update(self):
    dx = self.p1.x - self.p2.x
    dy = self.p1.y - self.p2.y
    dist = math.hypot(dx, dy)
    theta = math.atan2(dy, dx)
    force = (self.length - dist) * self.strength
    self.p1.accelerate((theta + 0.5 * math.pi, force/self.p1.mass))
    self.p2.accelerate((theta - 0.5 * math.pi, force/self.p2.mass))



while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_SPACE:
                paused = (True, False)[paused]
        elif event.type == pygame.MOUSEBUTTONDOWN:
            selected_particle = universe.findParticle(pygame.mouse.get_pos())
        elif event.type == pygame.MOUSEBUTTONUP:
            selected_particle = None

    if selected_particle:
        selected_particle.mouseMove(pygame.mouse.get_pos())
    if not paused:
        universe.update()
        
    screen.fill(universe.colour)
    
    for p in universe.particles:
        pygame.draw.circle(screen, p.colour, (int(p.x), int(p.y)), p.size, 0)
        
    for s in universe.springs:
        pygame.draw.aaline(screen, (0,0,0), (int(s.p1.x), int(s.p1.y)), (int(s.p2.x), int(s.p2.y)))
    pygame.display.flip()