Here's an attempt I've made of implementing the Barnes-Hut n-body algorithm, or its initial stage - The quadtree. That there're lots of lengthy doc strings might excuse the lack of detailed explanation here.
My primary concerns as a beginner are to
- Avoid complexity
- Avoid library imports
- Make fast and short code whenever possible.
Being said that, I think the code is bulky (and 'unpythonic'?). Advice me on how to improve it in terms of speed, design pattern and readability. Is functional programming even suitable here?
Summary of the code
- Every node and body is identified by a global ID.
- A body is added to the quadtree by registering its mass, position and velocity. It is assigned an ID (called bid).
- A node is made by specifying the South-West vertex and box length (an ID is issued here too - called nid).
updatecomis called for updating the quadtree's (or individual nodes') centre of mass.walkis called throughout the script for navigating the quadtree.- Accelerating the bodies, updating position for each time-frame and plotting are yet to be established.
The docstrings are what makes the code lengthy which otherwise would be short. Since there're no special areas to which attention needs to be drawn, I've added the complete code here.
"""Quadtree construction module.
ESSENTIAL TERMINOLOGY
a. Quadtree : A tree structure with four quadrants that are
coplanar. Quadrants are named SW, NW, NE and SE as
in compass directions, and are usually adressed in
the same order.
b. Node : Same as a quadrant; it can be internal or external.
c. Internal Node : A node that has child nodes inside. It contains no
bodies since they are distributed among the child
nodes.
d. External Node : A node that doesn't have child nodes inside. It
contains a body.
e. Empty Node : A node which has neither child nodes nor a body
inside it.
f. Root Node : The topmost node (parent) of all nodes.
g. Occupant : A body bound within a node (lowest host node).
h. nid : Node ID
i. bid : Body ID
j. Centre of mass : The weighted average of mass distribution.
X = (m1x1+m2x2+...mnxn) / (m1+m2+...mn)
This must be replaced with centre of gravity if the
bodies are non-uniform mass distributions, unlike a
point object.
"""
from __future__ import division
# Dictionaries to map body parameters
masses = {}
positions = {}
velocities = {}
# Dictionaries to map node parameters
lengths = {0: None}
vertices = {0: None}
occupants = {0: None}
childnodes = {0: []}
totalmasses = {0: 0}
centremasses = {0: None}
# Global ID counter starts with 1 (0 is dedicated for root node)
ID = 1
def newid(reset=False):
"""newid(reset=False)
Returns an unused ID of type 'int'. Consequtive integers starting from one
are returned avoiding repetition.
If optional argument 'reset' is True, ID counter restarts from one. Those
IDs which aren't active anymore can be retrieved by this manner.
"""
global ID
if reset:
ID = 1
while any(ID in maps for maps in (masses, childnodes)):
ID += 1
return ID
def initiate(vertex, length):
"""initiate(vertex, length)
Initiates the quadtree by specifying the SW vertex and side length.
"""
lengths[0] = length
vertices[0] = vertex
def isexternal(node):
"""isexternal(node)
Returns True if the node is external.
"""
if occupants[node]:
return True
def isinternal(node):
"""isinternal(node)
Returns True if the node is internal.
"""
if childnodes[node]:
return True
def isempty(node):
"""isempty(node)
Returns True if the node is empty.
"""
if (not occupants[node]) and (not childnodes[node]):
return True
def addbody(mass, position, velocity, bid=None):
"""addbody(mass, position, velocity, bid=None)
Adds a body to the quadtree.
'mass' must be a positive quantity of type 'int' or 'float'.
'position' must be a tuple or list of two values representing the two
dimensional position coordinates: (x, y).
'velocity' must be a tuple or list of two values representing the two
dimensional velocity coordinates: (vx, vy).
Optional argument bid, if specified, must be unique (auto-generated
otherwise).
"""
if not bid or bid in masses:
bid = newid()
masses[bid] = mass
positions[bid] = position
velocities[bid] = velocity
attach(bid)
def canfit(body, node):
"""canfit(body, node)
Checks wheather the body could be accomodated inside the node. Returns
suitable status mesages;
a. EXTERNAL : Can't fit because there is another body inside.
b. INTERNAL : Can't fit because this node has child nodes.
c. EMPTY : Can fit (yay!)
d. None : Can't fit at all (out of bounds).
"""
bx, by = positions[body]
nx, ny = vertices[node]
l = lengths[node]
if (0<=bx-nx<=l) and (0<=by-ny<=l):
if isexternal(node):
return 'EXTERNAL'
if isinternal(node):
return 'INTERNAL'
return 'EMPTY'
def split(node, rearrange=True):
"""split(node, rearrange=True)
Splits the given node (external or empty) and divides it in to four
quadrants (internal).
If 'rearrange' is True, the body (if any) contained within this node will
be redistributed to its appropriate childnode and thereby will get detached
from the node itself.
"""
nx, ny = vertices[node]
h = lengths[node] / 2
hx, hy = nx + h, ny + h
for vertex in ((nx,ny), (nx,hy), (hx,hy), (hx,ny)):
nid = newid()
lengths[nid] = h
vertices[nid] = vertex
occupants[nid] = None
childnodes[nid] = []
totalmasses[nid] = 0
centremasses[nid] = None
childnodes[node].append(nid)
# Detach the body, distribute it to a child node.
if rearrange:
body = occupants[node]
occupants[node] = None
for child in childnodes[node]:
if attach(body, child):
break
def attach(body, node=0):
"""attach(body, node=0)
Attaches body to its host node and returns True upon success.
Note: If recursion is initiated, and two bodies happens to be so close to
each other, it could take considerable depth until both the bodies are
individually distributed.
"""
for nid in walk(node):
status = canfit(body, nid)
if status == 'EMPTY':
occupants[nid] = body
return True
elif status == 'EXTERNAL':
split(nid)
for child in childnodes[nid]:
if attach(body, child):
return True
def walk(top=0, topdown=True, gettop=True):
"""walk(top=0, topdown=True, gettop=True)
Walks through the quadtree generating nids.
If 'top' is not specified, the generator will walk through the entire
quadtree (only top node otherwise).
If 'topdown' is True, parent nodes will be generated before child nodes.
Otherwise, child nodes will be generated before all their parents are. eg:
parent : child nodes
0 : 1, 2, 3, 4
3 : 5, 6, 7, 8
4 : 9, 10, 11, 12
7 : 13, 14, 15, 16
topdown mode - 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16
bottomup mode - 1, 2, 5, 6, 13, 14, 15, 16, 7, 8, 3, 9, 10, 11, 12, 4, 0
if optional argument 'gettop' is False, the top node will not be yielded.
"""
if topdown:
if gettop:
yield top
# Holding a queue list is memory inefficient, needs workaround.
queue = list(childnodes[top])
for child in queue:
yield child
queue.extend(childnodes[child])
# This part is recursive.
else:
for child in childnodes[top]:
for node in walk(child, topdown, gettop=True):
yield node
if gettop:
yield top
def updatecom(node=0):
"""updatecom(node=0)
Updates centre of mass of the node and its childnodes if any.
If optional argument 'node' is not provided, the entire quadtree's centre
of mass will be updated.
"""
for nid in walk(node, topdown=False):
# Empty node, reset to None
if isempty(nid):
totalmasses[nid] = 0
centremasses[nid] = None
# External node, com is where the body is.
if isexternal(nid):
body = occupants[nid]
totalmasses[nid] = masses[body]
centremasses[nid] = positions[body]
# Internal node, add children's com to parent.
elif isinternal(nid):
tmass, comx, comy = 0, None, None
for child in childnodes[nid]:
if isinternal(child) or isexternal(child):
ccomx, ccomy = centremasses[child]
ctmass = totalmasses[child]
if comx is None:
tmass, comx, comy = ctmass, ccomx, ccomy
continue
comx += ccomx*ctmass
comy += ccomy*ctmass
tmass += ctmass
totalmasses[nid] = tmass
centremasses[nid] = (comx/tmass, comy/tmass)
def distance(pointA, pointB):
"""distance(pointA, pointB)
Calculates the distance between two points and returns a tuple of
(dx, dy, r); where r is the radial separation, dx and dy are the axial
separation.
"""
ax, ay = pointA
bx, by = pointB
dx, dy = ax - bx, ay - by
r = (dx*dx+dy*dy)**0.5
return dx, dy, r
def summary(brief=False, linecolour=False):
""" Prints a tabulated version of the quadtree information.
(for debugging only)
Body parameters:
BID, MASS, POSITION, VELOCITY
Node parametres:
NID, LENGTH, VERTEX, OCCUPANT, CHILDREN, MASS, COM
"""
def history():
global period
print ("{0} nodes were created to fit {1} bodies in {2} seconds."
.format(len(childnodes), len(masses), style[2]%period))
if brief:
history()
return
# Styles and formats
style = ['{:>7}', '{:>15}', '%.2f', '(%.2f, %.2f)', '{:>19}', '\x1b[0m',
'\x1b[7m', '\x1b[0m', '{:>9}', '{:>31}']
form1 = style[0]*2 + style[1]*2
form2 = style[0]*2 + style[4] + style[8] + style[9] + style[8] + style[1]
head1 = ["BID", "MASS", "POSITION", "VELOCITY"]
head2 = ["NID", "LENGTH", "VERTEX", "OCCUPANT", "CHILDREN", "MASS", "COM"]
print form1.format(*head1)
for i, bid in enumerate(masses):
if linecolour:
if i % 2 == 0: bg = style[6]
else: bg = style[7]
else: bg = ''
print bg+form1.format(bid, style[2]%masses[bid],
style[3]%positions[bid],
style[3]%velocities[bid])
print style[5]
print form2.format(*head2)
for i, nid in enumerate(childnodes):
if linecolour:
if i % 2 == 0: bg = style[6]
else: bg = style[7]
else: bg = ''
if centremasses[nid] is None: printcom = None
else: printcom = style[3]%centremasses[nid]
print bg+form2.format(nid, style[2]%lengths[nid],
style[3]%vertices[nid], occupants[nid],
childnodes[nid], style[2]%totalmasses[nid],
printcom)
print style[5]
history()
if __name__ == '__main__':
from random import uniform
import time
theta = 1
G = 6.673e-1
dt = 0.1
start = time.time()
# Initialize the quadtree
initiate((-5,-5), 10)
n = xrange(8)
m = (uniform(1,100) for i in n)
p = ((uniform(-5,5),uniform(-5,5)) for i in n)
v = ((uniform(-10,10),uniform(-10,10)) for i in n)
# Add bodies
for mass, pos, vel in zip(m, p, v):
addbody(mass, pos, vel)
# Refresh centre of mass
updatecom()
period = time.time() - start
summary(linecolour=True)
