-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathBHTree.py
100 lines (93 loc) · 4.02 KB
/
BHTree.py
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
#################################################################
# Name: BHTree.py #
# Authors: Chris (Yuan Qi) Ni #
# Michael Battaglia #
# Date: December 17, 2016 #
# Function: Program is object definition for a node in a #
# Barnes-Hut tree for Barnes-Hut N-body simulation. #
#################################################################
#essential modules
import numpy as np
import matplotlib.pyplot as plt
#essential imports
from Body import Body
from Quad import Quad
#class: node of Barnes-Hut tree
class BHTree:
"""definition of node of Barnes-Hut tree"""
def __init__(self, quad):
self.quad = quad
def insertBody(self, body):
#add body to Barnes-Hut node
if hasattr(self, 'body'):
#node is not empty
if self.external:
#node is external node, make into internal node
self.external = False
#create 4 child trees
self.SW = BHTree(self.quad.SW())
self.SE = BHTree(self.quad.SE())
self.NW = BHTree(self.quad.NW())
self.NE = BHTree(self.quad.NE())
#sort node body into child trees
if self.body.inQuad(self.quad.SW()):
#new body in SW quadrant node
self.SW.insertBody(self.body)
if self.body.inQuad(self.quad.SE()):
#new body in SW quadrant node
self.SE.insertBody(self.body)
if self.body.inQuad(self.quad.NW()):
#new body in SW quadrant node
self.NW.insertBody(self.body)
if self.body.inQuad(self.quad.NE()):
#new body in SW quadrant node
self.NE.insertBody(self.body)
#sort new body into child trees
if body.inQuad(self.quad.SW()):
#new body in SW quadrant node
self.SW.insertBody(body)
if body.inQuad(self.quad.SE()):
#new body in SW quadrant node
self.SE.insertBody(body)
if body.inQuad(self.quad.NW()):
#new body in SW quadrant node
self.NW.insertBody(body)
if body.inQuad(self.quad.NE()):
#new body in SW quadrant node
self.NE.insertBody(body)
#add to node body aggregate mass
R, M = self.body.r, self.body.m
r, m = body.r, body.m
R = (M*R + m*r)/(M + m)
M += m
self.body = Body(M, R[0], R[1])
else:
#if node is empty, add body and make external node
self.body = body
self.external = True
def applyForce(self, body, theta, epsilon):
#evaluate force on body from tree with resolution theta
if hasattr(self, 'body'):
#not an empty node
if (self.body.r != body.r).any():
#not self force
d = body.distanceTo(self.body) #distance of node body to body
if self.quad.L/d < theta or self.external:
#box sufficiently far away for its size, compute force
body.addForce(self.body, epsilon)
else:
#box too close, compute forces from children instead
self.SW.applyForce(body, theta, epsilon)
self.SE.applyForce(body, theta, epsilon)
self.NW.applyForce(body, theta, epsilon)
self.NE.applyForce(body, theta, epsilon)
def plot(self):
if hasattr(self, 'body'):
#not an empty node
self.quad.plot()
if not self.external:
#plot quadrants in every child node
self.SW.plot()
self.SE.plot()
self.NW.plot()
self.NE.plot()