orbital/orbital.py

#!/usr/bin/python
import pygame
from pygame.locals import *
import os
import sys
import math
import random

"""
Units of measure:
	1 pixel represents 10^9 m
	1 unit of speed represents 10^9 m/day
	1 unit of mass represents 10^24 kg
	graviational constant
"""
grav_const = 6.6742

pause = True
lock_index = 0
def handle_key_ev(key):
	global bounce, speed_player, trails, dir_player, color_player, lock, lock_index, pause
	if key == K_b:
		bounce = not bounce
	elif key == K_q:
		sys.exit(0)
	elif key == K_c:
		scr.fill((0,0,0))
	elif key == K_t:
		trails = not trails
	elif key == K_s:
		lock = sun
	elif key == K_e:
		lock = earth
	elif key == K_u:
		lock = None
	elif key == K_n:
		lock_index += 1
		if lock_index >= len(slaves):
			lock_index = 0
		lock = slaves[lock_index]
	elif key == K_EQUALS:
		rescale(scale * 1.5)
	elif key == K_MINUS:
		rescale(scale / 1.5)
	elif key == K_p:
		pause = not pause

def rescale(newscale):
	global scale
	scale = newscale
	scr.fill((0,0,0))
	
	
def handle_keys(keys):
	global bounce, speed_player, trails, dir_player, color_player
	for key in keys:
		if key == K_UP:
			speed_player += 0.01
		elif key == K_DOWN:
			speed_player -= 0.01
		elif key == K_LEFT:
			dir_player -= (math.pi / 180.0)
		elif key == K_RIGHT:
			dir_player += (math.pi / 180.0)


def xy2dir(coords):
	if abs(coords[0]) < 0.00000001 and abs(coords[1]) > 0.000001:
		res = math.pi
	elif coords[0] == 0:
		res = math.pi
	else:
		res = math.atan(coords[1] / coords[0]) + (math.pi/2)
		
	if coords[0] < 0:
		res += math.pi
	return res

class Vector(object):
	def __init__(self, dir, mag):
		self.dir = float(dir)
		self.mag = float(mag)
	def __add__(self, other):
		va = self.resolve()
		vb = other.resolve()
		vf = [va[0] + vb[0], va[1] + vb[1]]
		
		newmag = math.sqrt(vf[0]**2 + vf[1]**2)
		newdir = xy2dir(vf)
		return Vector(newdir, newmag)

	def __sub__(self, other):
		return self + Vector(other.dir, -other.mag)

	def __mul__(self, scalar):
		self.mag *= scalar
	def __div__(self, scalar):
		self.mag /= scalar
	
	def resolve(self):
		return [(math.sin(self.dir) * self.mag), -(math.cos(self.dir) * self.mag)]
	def __repr__(self):
		return "<Vector(dir=%f, mag=%f)>" % (self.dir, self.mag)

class GravitySlave(object):
	def __init__(self, coords, vector, color):
		self.coords = [float(x) for x in coords]
		self.vector = vector
		self.trails = False
		self.last_coords = None
		self.color = color
	def move(self):
		self.last_coords = [x for x in self.coords]
		offset = self.vector.resolve()
		self.coords[0] += offset[0]
		self.coords[1] += offset[1]
	def draw(self):
		lc = scale_for_screen(self.last_coords)
		c = scale_for_screen(self.coords)
		scr.lock()
		if not trails:
			if lc:
				pygame.draw.circle(scr, (0,0,0), lc, 2, 1)
		pygame.draw.circle(scr, self.color, c, 2, 1)
		scr.unlock()
		


scale = 1.0
def scale_for_screen(coords):
	if not coords:
		return None
	
	x, y = coords
	x, y = [x - origin_x, y - origin_y]
	x, y = [x * scale, y * scale]
	x, y = [x + origin_x, y + origin_y]
	return [int(x), int(y)]

class GravityEmitter(GravitySlave):
	def __init__(self, coords=[0,0], vector=Vector(0,0), mass=1, color=(255,255,255)):
		GravitySlave.__init__(self, coords, vector, color)
		self.mass = float(mass)
		

	def vectorize(self, coords):
		coords = [float(x) for x in coords]
		xdist = -coords[0] + self.coords[0]
		ydist = -coords[1] + self.coords[1]
		dist_sqr = (xdist**2 + ydist**2)
		mag = self.mass * (grav_const * 10**-5) / dist_sqr
		mag = mag * 8.64**2 * 10**-1
		dir = xy2dir([xdist, ydist])
		return Vector(dir, mag)



def apply_emitters(slave):
	for emit in emitters:
		if id(emit) != id(slave):
			slave.vector += emit.vectorize(slave.coords)
def apply_all_emitters():
	for slave in slaves:
		apply_emitters(slave)
	
if __name__ == '__main__':
	pygame.init()

	dir_player = math.pi

	size = width, height = (800,600)
	scr = pygame.display.set_mode(size)

	#options
	color_player = (255,255,255)
	bounce = False
	speed_player = 5.0
	trails = False

	scr.fill((0,0,0))

	keys = []
	origin_x = 400.0
	origin_y = 300.0
	earth = GravityEmitter(coords=[origin_x - 147.098074, origin_y], 
			mass=5.9742,
			vector=Vector(math.pi, 2.616797),
			color=(32,32,255)
			)
	mars = GravityEmitter(coords=[origin_x + 249.228730, origin_y], 
			mass=0.64185, 
			vector=Vector(0, 1.8983808),
			color=(220,0,0)
			)
	venus = GravityEmitter(coords=[origin_x, origin_y + 108.941849], 
			mass=4.8685,
			vector=Vector(math.pi/2, 3.0053376),
			color=(255,255,180)
			)
	merc = GravityEmitter(coords=[origin_x, origin_y - 69.817079], 
			mass=0.3302,
			vector=Vector(math.pi*3/2, 3.357504),
			color=(255,120,0)
			)
	
	sun = GravityEmitter(coords=[origin_x, origin_y], 
			mass=1989100,
			color=(255,255,32)
			)
	evilsun = GravityEmitter(coords=[800.0, 600.0], 
			mass=989100,
			color=(120,120,120),
			vector=Vector(math.pi*3/16, 1.623)
			)
	slaves = [earth, sun, mars, venus, merc]
	emitters = [x for x in slaves if type(x) == GravityEmitter]
	lock = None

	while True:
		events = pygame.event.get((KEYDOWN, KEYUP))
		for ev in events:
			if ev.type == KEYUP:
				try:
					del keys[keys.index(ev.key)]
				except IndexError:
					pass
			else:
				handle_key_ev(ev.key)
				keys += [ev.key]
			
				
		handle_keys(keys)	
		pygame.event.pump()

		if not pause:
			#ball.trails = trails
			apply_all_emitters()
			for o in slaves:
				o.move()

		if lock != None:
			lv = Vector(lock.vector.dir, lock.vector.mag)
			lc = [origin_x - lock.coords[0], origin_y - lock.coords[1]]
			for o in slaves:
				o.vector = o.vector - lv
				o.coords[0] += lc[0]
				o.coords[1] += lc[1]

		for o in slaves:
			o.draw()
			
		pygame.display.flip()