Skip to content

Adds features, does some code clean up, and moves Body class to it's own file #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open

Adds features, does some code clean up, and moves Body class to it's own file #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9f45d0a
Add features and some code clean up
quonic Jun 6, 2019
35976a8
Create Planet.py
quonic Jun 6, 2019
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 35 additions & 0 deletions Planet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
import math

G = 6.67408e-11 * 100_000_000


class Body:
def __init__(self, pos, velocity, mass, id_num):
self.pos = pos # pos is a list of x and y position of that body in pixels eg : [500,600]
self.velocity = velocity # b is a list of x and y components of velocity of that body in pixel units
self.mass = mass # m is the mass of that object
self.size = math.log2(self.mass) # size is used for only visualisation of the body
self.id = id_num # id is the unique identifier for this body
self.hit = False # tell us if this object was hit by another
self.invuln = False # a temporary flag for showing what is invulnerable

def update_size(self):
self.size = math.log2(self.mass)

def n_body(self, other_bodies, n_fx_total, n_fy_total):
for body_b in other_bodies:
if body_b.pos != self.pos:
fx, fy = self.calculate_forces(self.pos, body_b.pos, self.mass, body_b.mass)
n_fx_total += fx
n_fy_total += fy
return [n_fx_total, n_fy_total]

@staticmethod
def calculate_forces(c_pos_a: list, c_pos_b: list, c_m_a: float, c_m_b: float):
x_diff = c_pos_b[0] - c_pos_a[0]
y_diff = c_pos_b[1] - c_pos_a[1]
f = G * c_m_a * c_m_b / math.sqrt(x_diff ** 2 + y_diff ** 2)
angle = math.atan2(y_diff, x_diff)
cfx = f * math.cos(angle)
cfy = f * math.sin(angle)
return cfx, cfy
216 changes: 122 additions & 94 deletions gravity_simulation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,100 +1,128 @@
# Very simple implementation of simulation of gravity on bodies in 2D. Does not handle the case when 2 or more
# bodies collide with each other
import random
import math
import sys
import math
import pygame

class Body:
def __init__(self, pos, a, v, m):
self.pos = pos # pos is a list of x and y position of that body in pixels eg : [500,600]
self.a = a # a is a list of x and y components of accelaration of that body in pixel units
self.v = v # b is a list of x and y components of velocity of that body in pixel units
self.m = m # m is the mass of that object


def calculate_forces(pos_a, pos_b, m_a, m_b):
x_diff = pos_b[0] - pos_a[0]
y_diff = pos_b[1] - pos_a[1]
hypotenuse = math.sqrt(((x_diff) ** 2 + (y_diff) ** 2))
sin = x_diff / hypotenuse
cos = y_diff / hypotenuse
f = G * m_a * m_b / hypotenuse ** 2
fx = f * sin
fy = f * cos

return fx, fy


G = 6.67408e-11 * 100_000_000 # Otherwise the bodies would not move given the small value of gravitational constant
NUM_OF_BODIES = 10
WIDTH = 900
HEIGHT = 800
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BLUE = (109, 196, 255)

bodies = []
for i in range(NUM_OF_BODIES):
px = random.randint(10, WIDTH - 10)
py = random.randint(10, HEIGHT - 10)
m = random.randint(1, 25)
bodies.append(Body([px, py], [0, 0], [0, 0], m))

# Some predefined bodies for the purpose of testing
# bodies.append(Body([500,500],[0,0],[0,0],20))
# bodies.append(Body([510,503],[0,0],[0,0],7))
# bodies.append(Body([400,400],[0,0],[0,0],14))
# bodies.append(Body([10,600],[0,0],[0,0],9))
# bodies.append(Body([250,198],[0,0],[0,0],18))
# bodies.append(Body([340,700],[0,0],[0,0],24))

pygame.init()
size = WIDTH, HEIGHT
screen = pygame.display.set_mode(size)

font = pygame.font.SysFont('Arial', 16)
text = font.render('0', True, BLUE)
textRect = text.get_rect()
while True:
screen.fill(BLACK)
for event in pygame.event.get():
if event.type == pygame.QUIT: sys.exit()

for body_a in bodies:
pos_a = body_a.pos
m_a = body_a.m
fx_total = 0
fy_total = 0

for body_b in bodies:
if body_b.pos == pos_a:
continue
fx, fy = calculate_forces(pos_a, body_b.pos, m_a, body_b.m)
fx_total += fx
fy_total += fy

body_a_acceleration = body_a.a

body_a_acceleration[0] = fx_total / m_a
body_a_acceleration[1] = fy_total / m_a

body_a.v[0] = body_a.v[0] + body_a_acceleration[0]
body_a.v[1] = body_a.v[1] + body_a_acceleration[1]

pos_a[0] = pos_a[0] + body_a.v[0]
pos_a[1] = pos_a[1] + body_a.v[1]

mass_text = 'M={0}'.format(m_a)
# force_text = 'F=({0},{1})'.format(fx_total.__round__(3), fy_total.__round__(3))
# velocity_text = 'V=({},{})'.format(body_a.v[0].__round__(3),body_a.v[1].__round__(3))
# text_str = mass_text + ' ' + force_text + ' ' + velocity_text
text_str = mass_text

text = font.render(text_str, True, BLUE)
textRect.center = (pos_a[0] + m_a + 10, pos_a[1] + m_a + 10)

screen.blit(text, textRect)

pygame.draw.rect(screen, (255, 255, 255), pygame.Rect(pos_a[0], pos_a[1], m_a, m_a))
pygame.display.flip()
from Planet import Body

if __name__ == "__main__":
NUM_OF_BODIES = 20
WIDTH = 900
HEIGHT = 800
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BLUE = (109, 196, 255)

bodies = []
for i in range(NUM_OF_BODIES):
px = random.randint(10, WIDTH - 10)
py = random.randint(10, HEIGHT - 10)
m = random.randint(1, 20)
bodies.append(Body([px, py], [0, 0], m, i))

pygame.init()
size = WIDTH, HEIGHT
screen = pygame.display.set_mode(size)

font = pygame.font.SysFont('Arial', 16)
text = font.render('0', True, BLUE)
textRect = text.get_rect()
velocity_diff = [0, 0]
while True:
screen.fill(BLACK)
for event in pygame.event.get():
if event.type == pygame.QUIT:
sys.exit()

# Don't process bodies that where hit/removed
bodies = [body for body in bodies if not body.hit]

# Get Bodies and find the center of mass of all
x = [p.pos[0] for p in bodies]
y = [p.pos[1] for p in bodies]
centroid = (sum(x) / len(bodies), sum(y) / len(bodies))
lx = bodies[0].pos[0]
ly = bodies[0].pos[1]
for body in bodies:
body.pos[0] = body.pos[0] - centroid[0] + WIDTH / 2
body.pos[1] = body.pos[1] - centroid[1] + HEIGHT / 2

# Draw circles and line for reference point of origin
velocity_diff[0] = lx - centroid[0]
velocity_diff[1] = ly - centroid[1]
# Origin Point
textRect.center = (velocity_diff[0] + 10, velocity_diff[1] + 10)
screen.blit(font.render("{0},{1}".format(0, 0), True, BLUE), textRect)
pygame.draw.circle(screen, (255, 255, 127), [int(velocity_diff[0]), int(velocity_diff[1])], 3, 1)
# Center of all objects
textRect.center = (WIDTH/2 + 10, HEIGHT/2 + 10)
screen.blit(font.render("{0},{1}".format(
int(math.floor(velocity_diff[0])),
int(math.floor(velocity_diff[1]))), True, BLUE), textRect)
# Draw line from origin to center of all objects
pygame.draw.line(screen, (255, 255, 0), (WIDTH/2, HEIGHT/2), velocity_diff)
pygame.draw.circle(screen, (255, 255, 127), [int(WIDTH/2), int(HEIGHT/2)], 3, 1)

for body_a in bodies:
# Remove invulnerability flag
body_a.invuln = False

f_totals = body_a.n_body(bodies, 0, 0)

body_a.velocity[0] = body_a.velocity[0] + f_totals[0] / body_a.mass
body_a.velocity[1] = body_a.velocity[1] + f_totals[1] / body_a.mass

body_a.pos[0] = body_a.pos[0] + body_a.velocity[0]
body_a.pos[1] = body_a.pos[1] + body_a.velocity[1]

mass_text = 'M={0}'.format(body_a.mass)
# force_text = 'F=({0},{1})'.format(fx_total.__round__(3), fy_total.__round__(3))
# velocity_text = 'V=({},{})'.format(body_a.v[0].__round__(3),body_a.v[1].__round__(3))
# text_str = mass_text + ' ' + force_text + ' ' + velocity_text
text_str = mass_text

text = font.render(text_str, True, BLUE)
textRect.center = (
body_a.pos[0] + body_a.size + 10,
body_a.pos[1] + body_a.size + 10)

screen.blit(text, textRect)

pygame.draw.circle(
screen,
(255, 255, 255),
[int(body_a.pos[0]), int(body_a.pos[1])], int(body_a.size))

# Get a list of bodies, except for body_a
next_bodies = [body for body in bodies if not body_a.id == body.id]
for body in next_bodies:
# if body is invulnerable then continue on to the next body
if body.invuln:
continue
# Find the distance to body_a
distance = math.sqrt(
((body_a.pos[0] - body.pos[0]) * (body_a.pos[0] - body.pos[0])) +
((body_a.pos[1] - body.pos[1]) * (body_a.pos[1] - body.pos[1])))
# If bodied touch then "remove" one by setting flag and making the other invulnerable
if distance < int(body_a.size) + int(body.size):
if body_a.mass >= body.mass:
body_a.mass += body.mass
body_a.update_size()
body_a.velocity[0] = (body_a.mass * body_a.velocity[0] + body.mass * body.velocity[0]
) / (body_a.mass + body.mass)
body_a.velocity[1] = (body_a.mass * body_a.velocity[1] + body.mass * body.velocity[1]
) / (body_a.mass + body.mass)
body.hit = True
body_a.invuln = True
else:
body.mass += body_a.mass
body.update_size()
body.velocity[0] = (body_a.mass * body_a.velocity[0] + body.mass * body.velocity[0]
) / (body_a.mass + body.mass)
body.velocity[1] = (body_a.mass * body_a.velocity[1] + body.mass * body.velocity[1]
) / (body_a.mass + body.mass)
body_a.hit = True
body.invuln = True
pygame.display.flip()