n-bodies.py

# Simulates n point masses interacting with gravity. The masses, starting positions
# and velocites are initialized to random values. The partices can pass through 
# each other for the time being. The red and blue lines point in the direction of
# the velocity vector and force vector, respectively. 


import pygame
import random
import math
import pdb

running = True
number_of_particles = 3
background_colour = (255,255,255)
velo_line_colour = (255,0,0)
force_line_colour = (100,0,255)
(width, height) = (1000, 1000)
drag = 1
elasticity = .83
gravity = (math.pi, 0.00)
G = 10
timestep = 0.5


def addVectors(v1, v2):
    angle1 = v1[0]
    length1 = v1[1]
    angle2 = v2[0]
    length2 = v2[1]
    x  = math.cos(angle1) * length1 + math.cos(angle2) * length2
    y  = math.sin(angle1) * length1 + math.sin(angle2) * length2
    angle = math.atan2(y, x)
    length  = math.hypot(x, y)
    return [angle, length]

def distanceVector(a,b):
    x = b.x - a.x
    y = b.y - a.y
    angle = math.atan2(y,x)
    magnitude = math.hypot(x,y)
    return (angle, magnitude)

def acceleration(particle):
    resultant = [0,0]
    for other in my_particles:
        if other != particle:
            gMag = G*other.mass/(distanceVector(particle,other)[1]**2)
            gArg = distanceVector(particle,other)[0]
            resultant = addVectors(resultant,[gArg, gMag])
    return resultant

def eulerIntegrate(particle,dt):
    x = particle.x + dt * particle.speed * math.cos(particle.angle)
    y = particle.y + dt * particle.speed * math.sin(particle.angle)
    [angle,speed] = addVectors([particle.angle,particle.speed],particle.f)
    particle.x = x
    particle.y = y
    particle.angle = angle
    particle.speed = speed

def rungeKutta4(particle,dt):
    x0 = particle.x
    y0 = particle.y
    s0 = particle.speed
    a0 = particle.angle

    x1 = x0
    y1 = y0
    a1 = a0
    s1 = s0
    acc1 = acceleration(particle)

    x2 = x0 + 0.5 * s1 * math.cos(a1) * dt
    y2 = y0 + 0.5 * s1 * math.sin(a1) * dt
    (a2,s2) = addVectors([a1,s1],[acc1[0],acc1[1]*0.5*dt])
    particle.x = x2
    particle.y = y2
    acc2 = acceleration(particle)

    x3 = x0 + 0.5 * s2 * math.cos(a2) * dt
    y3 = y0 + 0.5 * s2 * math.sin(a2) * dt
    (a3,s3) = addVectors([a2,s2],[acc2[0],acc2[1]*0.5*dt])
    particle.x = x3
    particle.y = y3
    acc3 = acceleration(particle)

    x4 = x0 + s3 * math.cos(a3) * dt
    y4 = y0 + s3 * math.sin(a3) * dt
    (a4,s4) = addVectors([a3,s3],[acc3[0],acc3[1]*dt])
    particle.x = x4
    particle.y = y4
    acc4 = acceleration(particle)

    xf = x0 + (dt/6.0)*(s1*math.cos(a1)+2.0*s2*math.cos(a2)+2.0*s3*math.cos(a3)+s4*math.cos(a4))
    yf = y0 + (dt/6.0)*(s1*math.sin(a1)+2.0*s2*math.sin(a2)+2.0*s3*math.sin(a3)+s4*math.sin(a4))
    (af,sf) = addVectors([a0,s0],addVectors([acc1[0],acc1[1]*dt/6.0],addVectors([acc2[0],acc2[1]*2.0*dt/6.0],addVectors([acc3[0],acc3[1]*2.0*dt/6.0],[acc4[0],acc4[1]*dt/6.0]))))

    particle.x = xf
    particle.y = yf
    particle.angle = af
    particle.speed = sf

class Particle():
    def __init__(self, (x, y), size):
        self.x = x
        self.y = y
        self.size = size
        self.mass = size*size
        self.colour = (0, 0, 255)
        self.thickness = 2
        self.speed = 0
        self.angle = 0

    def display(self):
        pygame.draw.circle(screen, self.colour, (int(self.x), int(self.y)), self.size, self.thickness)
            
        vHyp = 30#2.0*self.size*self.speed/5
        vEnd = (self.x+vHyp*math.cos(self.angle), self.y + vHyp*math.sin(self.angle))
        pygame.draw.aaline(screen, velo_line_colour, (self.x, self.y), vEnd)

        self.f = acceleration(self)
        fHyp = 30#2.0*self.size*f[1]*300
        fEnd = (self.x + fHyp*math.cos(self.f[0]), self.y + fHyp*math.sin(self.f[0]))
        pygame.draw.aaline(screen, force_line_colour, (self.x, self.y), fEnd)
        self.f = (self.f[0],self.f[1]/self.mass)

    def move(self):
        # eulerIntegrate(self,timestep)
        rungeKutta4(self,timestep)
        self.speed *= drag
        self.bounce()

    def bounce(self):
        if self.x > width - self.size:
            self.x = 2*(width - self.size) - self.x
            self.angle = math.pi - self.angle
            self.speed *= elasticity

        elif self.x < self.size:
            self.x = 2*self.size - self.x
            self.angle = math.pi - self.angle
            self.speed *= elasticity

        if self.y > height - self.size:
            self.y = 2*(height - self.size) - self.y
            self.angle = 2*math.pi - self.angle
            self.speed *= elasticity

        elif self.y < self.size:
            self.y = 2*self.size - self.y
            self.angle = 2*math.pi - self.angle
            self.speed *= elasticity


screen = pygame.display.set_mode((width, height))
pygame.display.set_caption('N-Bodies Interacting Via Gravity')

my_particles = []

for n in range(number_of_particles):
    size = random.randint(10, 20)
    x = random.randint(size, width-size)
    y = random.randint(size, height-size)

    # x = random.randint(600,800)
    # y = random.randint(600,800)

    # x = 400 + 200*(n)
    # y = 400 + 200*n

    particle = Particle((x, y), size)
    particle.speed = random.uniform(2,3)
    particle.angle = random.uniform(0, math.pi*2)

    my_particles.append(particle)
if __name__ == "__main__":
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
        screen.fill(background_colour)
        for particle in my_particles:
            particle.display()
            particle.move()
            # particle.bounce()
        first = False
        pygame.display.flip()