Spaces:
Running
Running
File size: 7,716 Bytes
6237635 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 |
from __future__ import division
import numpy as np
import math
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import matplotlib.patches as patches
import random
theta = 0.5
AU = (149.6e6 * 1000) # 149.6 million km, in meters.
G = 6.67408e-11 #m^3 kg^-1 s^-2
fig1 = plt.figure()
sim = fig1.add_subplot(111, aspect='equal')
fig2 = plt.figure()
quadt = fig2.add_subplot(111, aspect='equal')
class Node:
children = None
mass = None
center_of_mass = None
bbox = None
vx = vy = None
def quad_insert(root, x, y, m):
if root.mass is None: #when the root is empty, add the first particle
root.mass = m
root.center_of_mass = [x,y]
return
elif root.children is None:
root.children = [None,None,None,None]
old_quadrant = quadrant_of_particle(root.bbox, root.center_of_mass[0], root.center_of_mass[1])
if root.children[old_quadrant] is None:
root.children[old_quadrant] = Node()
root.children[old_quadrant].bbox = quadrant_bbox(root.bbox,old_quadrant)
quad_insert(root.children[old_quadrant], root.center_of_mass[0], root.center_of_mass[1], root.mass)
new_quadrant = quadrant_of_particle(root.bbox, x, y)
if root.children[new_quadrant] is None:
root.children[new_quadrant] = Node()
root.children[new_quadrant].bbox = quadrant_bbox(root.bbox,new_quadrant)
quad_insert(root.children[new_quadrant], x, y, m)
root.center_of_mass[0] = (root.center_of_mass[0]*root.mass + x*m) / (root.mass + m)
root.center_of_mass[1] = (root.center_of_mass[1]*root.mass + y*m) / (root.mass + m)
root.mass = root.mass + m
else:
new_quadrant = quadrant_of_particle(root.bbox, x, y)
if root.children[new_quadrant] is None:
root.children[new_quadrant] = Node()
root.children[new_quadrant].bbox = quadrant_bbox(root.bbox, new_quadrant)
quad_insert(root.children[new_quadrant], x, y, m)
root.center_of_mass[0] = (root.center_of_mass[0]*root.mass + x*m) / (root.mass + m)
root.center_of_mass[1] = (root.center_of_mass[1]*root.mass + y*m) / (root.mass + m)
root.mass = root.mass + m
def display(root):
if root.mass is None:
return
if root.children is not None:
x = (root.bbox[0] + root.bbox[1]) / 2
y = (root.bbox[2] + root.bbox[3]) / 2
width = x-root.bbox[0]
plt_node(root.bbox[0], root.bbox[2], width)
plt_node(root.bbox[0], y, width)
plt_node(x, root.bbox[2], width)
plt_node(x, y, width)
for i in xrange(4):
if root.children[i] is not None:
display(root.children[i])
else:
quadt.scatter(root.center_of_mass[0], root.center_of_mass[1])
def integrate(particles):
bodies = particles
n = len(bodies)
timestep = 24*3600 #one day
years = 2 * 365 #how many Earth years that simulate
for day in xrange(years):
particles_force = {}
root = Node()
root.center_of_mass = []
root.bbox = find_root_bbox(bodies)
for i in xrange(n):
quad_insert(root, bodies[i][3], bodies[i][4], bodies[i][2])
for i in xrange(n):
total_fx, total_fy = compute_force(root,bodies[i][3],bodies[i][4],bodies[i][2])
particles_force[bodies[i][0]] = (total_fx, total_fy)
for i in xrange(n):
fx, fy = particles_force[bodies[i][0]]
bodies[i][5] += fx / bodies[i][2] * timestep
bodies[i][6] += fy / bodies[i][2] * timestep
bodies[i][3] += bodies[i][5] * timestep
bodies[i][4] += bodies[i][6] * timestep
sim.scatter(bodies[i][3], bodies[i][4], c=bodies[i][1])
display(root)
quadt.scatter(root.center_of_mass[0], root.center_of_mass[1], c='red', marker='x')
def compute_force(root,x,y,m):
if root.mass is None:
return 0, 0
if root.center_of_mass[0] == x and root.center_of_mass[1] == y and root.mass == m:
return 0, 0
d = root.bbox[1]-root.bbox[0]
r = distance(x,y, root.center_of_mass[0], root.center_of_mass[1])
if d/r < theta or root.children is None:
return force(m, x, y, root.mass, root.center_of_mass[0], root.center_of_mass[1])
else:
fx = 0.0
fy = 0.0
for i in xrange(4):
if root.children[i] is not None:
fx += compute_force(root.children[i],x,y,m)[0]
fy += compute_force(root.children[i],x,y,m)[1]
return fx, fy
################################################# SUPPORTING FUNCTION ##############################################################
def force(m, x, y, mcm, xcm, ycm):
d = distance(x, y, xcm, ycm)
f = G*m*mcm/(d**2)
dx = xcm - x
dy = ycm - y
angle = math.atan2(dy, dx)
fx = math.cos(angle) * f
fy = math.sin(angle) * f
return fx, fy
def distance(x1, y1, x2, y2):
return math.sqrt((x2-x1)**2+(y2-y1)**2)
def plt_node(x, y, width):
quadt.add_patch(patches.Rectangle((x, y), width, width, fill = False))
def find_root_bbox(array):
""" Create a suitable square boundary box for the input particles
"""
if len(array) == 0 or len(array) == 1:
return None
xmin, xmax, ymin, ymax = array[0][3], array[0][3], array[0][4], array[0][4]
for i in xrange(len(array)):
if array[i][3] > xmax:
xmax = array[i][3]
if array[i][3] < xmin:
xmin = array[i][3]
if array[i][4] > ymax:
ymax = array[i][4]
if array[i][4] < ymin:
ymin = array[i][4]
if xmax - xmin == ymax - ymin:
return xmin, xmax, ymin, ymax
elif xmax - xmin > ymax - ymin:
return xmin, xmax, ymin, ymax+(xmax-xmin-ymax+ymin)
else:
return xmin, xmax+(ymax-ymin-xmax+xmin), ymin, ymax
def quadrant_of_particle(bbox, x, y):
"""Return position of quadrant of the particle (x,y)
"""
if y >= (bbox[3] + bbox[2])/2:
if x <= (bbox[1] + bbox[0])/2:
return 0
else:
return 1
else:
if x >= (bbox[1] + bbox[0])/2:
return 2
else:
return 3
def quadrant_bbox(bbox,quadrant):
"""Return the coordinate of the quadrant
"""
x = (bbox[0] + bbox[1]) / 2
y = (bbox[2] + bbox[3]) / 2
#Quadrant 0: (xmin, x, y, ymax)
if quadrant == 0:
return bbox[0], x, y, bbox[3]
#Quadrant 1: (x, xmax, y, ymax)
elif quadrant == 1:
return x, bbox[1], y, bbox[3]
#Quadrant 2: (x, xmax, ymin, y)
elif quadrant == 2:
return x, bbox[1], bbox[2], y
#Quadrant 3: (xmin, x, ymin, y)
elif quadrant == 3:
return bbox[0], x, bbox[2], y
def data_from_file(filename, array):
with open(filename) as f:
for line in f:
if line[0] == '#':
continue
else:
name,color,m,x,y,vx,vy = line.split(',')
array.append([name,color,float(m),float(x)*AU,float(y)*AU,float(vx)*1000,float(vy)*1000])
if __name__ == '__main__':
filename = ('solar-system.txt')
particles = []
data_from_file(filename, particles)
#root = Node()
#root.center_of_mass = []
#root.bbox = find_root_bbox(particles)
#for i in xrange(len(particles)):
# quad_insert(root, particles[i][3], particles[i][4], particles[i][2])
#print 'Boundary box: ',root.bbox
#print 'Total mass: ',root.mass
#print 'Coordinate of center of mass: ',root.center_of_mass
#plt.scatter(root.center_of_mass[0], root.center_of_mass[1], c='r', marker='x', s=50)
#print 'Theta: ', theta
integrate(particles)
plt.show() |