Commit c1838909 authored by Olaf's avatar Olaf

improved performance barnes hut

parent 5d248c37
......@@ -3,7 +3,7 @@ from predef import *
def create_node(root):
"""
Creates list containing 4 lists of 4 smaller nodes, the center of mass of the node, total mass of the node and the origin.
Does the same as the function create_node, but in 3D.
Parameters:
-----------
......@@ -13,69 +13,37 @@ def create_node(root):
Returns:
--------
node: list
contains 4 lists containing all information concerning the 4 new nodes, the center of mass, the total mass and the origin of the "root"
contains lists containing all information concerning the new nodes, the center of mass, the total mass and the origin of the "root"
"""
# extract data
pos = root[0]
mass = root[1]
origin = root[2]
nbody = root[3] # numbers of the bodies. use to keep track of which body has what position
r = root[4]/2 # 2r = width of square
pos_ref = pos - origin # set reference origin at 0
node = [[], [], [], [], (np.sum(pos.T*mass, axis = 1)/np.sum(mass)).tolist(), np.sum(mass), 2*r] # initialize node with origin, center of mass, total mass, width
east = (np.dot(pos_ref, [1,0])>0)
west = (np.dot(pos_ref, [1,0])<=0)
north = (np.dot(pos_ref, [0,1])>0)
south = (np.dot(pos_ref, [0,1])<=0)
northwest = north*west
northeast = north*east
southwest = south*west
southeast = south*east
nbody_NW = nbody[northwest]
nbody_NE = nbody[northeast]
nbody_ZW = nbody[southwest]
nbody_ZE = nbody[southeast]
# for each new internal node, create list of positions, masses, new origin, width of node and numbers of the bodies
# for each external node: create list of position, mass and number of the body
if nbody_NW.shape[0] != 0: # check if there are any bodies in node. if not: left as empty list
if nbody_NW.shape[0] > 1: # check if node is internal
node[0] = [pos[northwest], mass[northwest], origin + np.array([-r,r])/2, nbody_NW, r] # internal node
elif nbody_NW.shape[0] == 1: # check of node is external
node[0] = [(pos[northwest]).tolist()[0], float(mass[northwest]), int(nbody_NW)] # external node
if nbody_NE.shape[0] != 0:
if nbody_NE.shape[0] > 1:
node[1] = [pos[northeast], mass[northeast], origin + np.array([r,r])/2, nbody_NE, r]
elif nbody_NE.shape[0] == 1:
node[1] = [(pos[northeast]).tolist()[0], float(mass[northeast]), int(nbody_NE)]
if nbody_ZE.shape[0] != 0:
if nbody_ZE.shape[0] > 1:
node[2] = [pos[southeast], mass[southeast], origin + np.array([r,-r])/2, nbody_ZE, r]
elif nbody_ZE.shape[0] == 1:
node[2] = [(pos[southeast]).tolist()[0], float(mass[southeast]), int(nbody_ZE)]
if nbody_ZW.shape[0] != 0:
if nbody_ZW.shape[0] > 1:
node[3] = [pos[southwest], mass[southwest], origin + np.array([-r,-r])/2, nbody_ZW, r]
elif nbody_ZW.shape[0] == 1:
node[3] = [(pos[southwest]).tolist()[0], float(mass[southwest]), int(nbody_ZW)]
dim = origin.shape[0]
node = (2**dim + 3)*[[]]
node[-1] = 2*r
node[-2] = np.sum(mass)
node[-3] = (np.sum(pos.T*mass, axis = 1)/np.sum(mass)).tolist()
direction = (np.dot(pos - origin, np.eye(dim))>0)[:,np.newaxis]
new_origin = np.array([range(2**dim)]).T & 2**np.arange(dim) > 0
locations = (new_origin^direction).all(axis = 2).T
for i in range(2**dim):
if nbody[locations[i]].shape[0] > 1:
node[i] = [pos[locations[i]], mass[locations[i]], origin-r*(new_origin[i]-0.5), nbody[locations[i]], r]
elif nbody[locations[i]].shape[0] == 1:
node[i] = [pos[locations[i]].tolist(), float(mass[locations[i]]), int(nbody[locations[i]])]
# return
return node
def create_root(pos, mass, origin):
"""
Creates the root of the Barnes-Hut tree.
Creates the root of the Barnes-Hut tree in 3D.
Parameters:
-----------
......@@ -89,18 +57,15 @@ def create_root(pos, mass, origin):
Returns:
--------
root: list
contains 4 lists containing all information concerning the 4 new nodes, the center of mass, the total mass and the origin of the root
contains lists containing all information concerning the new nodes, the center of mass, the total mass and the origin of the root
"""
r = np.max(np.abs(pos - origin))
r = np.max(abs(pos - origin))
nbody = np.arange(mass.shape[0])
root_data = [pos, mass, origin, nbody, 2*r]
root = create_node(root_data)
root = create_node([pos, mass, origin, nbody, 2*r])
# return
return root
def build_tree(root):
"""
Builds the Barnes-Hut tree recursively.
......@@ -108,12 +73,12 @@ def build_tree(root):
Parameters:
-----------
root: list
contains 4 lists containing all information concerning the 4 new nodes, the center of mass, the total mass and the origin of the "root"
contains lists containing all information concerning the new nodes, the center of mass, the total mass and the origin of the "root"
Returns:
--------
node: list
contains 4 lists containing all information concerning the 4 new nodes, the center of mass, the total mass and the origin of the "root"
contains lists containing all information concerning the new nodes, the center of mass, the total mass and the origin of the "root"
"""
for i in range(len(root) - 3):
if len(root[i]) == 5:
......@@ -147,39 +112,13 @@ def create_tree(pos, mass):
return tree
def force_tree(theta, tree, pos):
"""
Calculates forces on the bodies by the nodes.
Parameters:
-----------
theta: float
ratio threshold of width node / distance body - center of mass node
tree: list
full Barnes-Hut tree
pos: np.array([N,D])
all positions of all bodies
Returns:
--------
force: np.array([N,D])
all net forces on all bodies
U: np.array([N,])
Potential energy
"""
force, U = force_node(pos, tree, theta, np.arange(pos.shape[0]), pos.shape)
# return
return force, U
def force_node(pos, node, theta, nbody, Mshape):
"""
Calculate all forces on the bodies at positions pos, by all nodes recursively
Parameters:
-----------
pos: list
pos: np.array([N,D])
position of the body
node: list
all information of the node that (possibly) exerts force on the body
......@@ -216,11 +155,11 @@ def force_node(pos, node, theta, nbody, Mshape):
force[nbody_for], U[nbody_for] = force_cal(pos_for, node[0], node[1])
elif len(node) == (2**Mshape[1] + 3):
d = np.linalg.norm(pos - np.asarray(node[4]), axis = 1)
d = np.linalg.norm(pos - np.asarray(node[-3]), axis = 1)
nbody_calc = nbody[node[-1]/d < theta]
if nbody_calc.shape[0] != 0:
force[nbody_calc], U[nbody_calc] = force_cal(pos[node[-1]/d < theta], node[4], node[5])
force[nbody_calc], U[nbody_calc] = force_cal(pos[node[-1]/d < theta], node[-3], node[-2])
pos_cont = pos[node[-1]/d >= theta]
nbody_cont = nbody[node[-1]/d >= theta]
......@@ -229,7 +168,7 @@ def force_node(pos, node, theta, nbody, Mshape):
force1, U1 = force_node(pos_cont, node[i], theta, nbody_cont, Mshape)
force += force1
U += U1
# return
return force, U
......@@ -249,19 +188,20 @@ def force_cal(pos1 , pos2, m):
Returns:
--------
F: np.array([1,D])
force on the body at pos1
F: np.array([N,D])
force on the body / bodies at pos
U: np.array([N,])
Potential energy
"""
dx = np.asarray(pos2) - pos1
dr = np.linalg.norm(dx, axis = 1)
F = ((G*m/(dr**3))*(dx.T)).T
U = G*m/dr
U = -G*m/dr
# return
return F, U
def force_barneshut(pos, mass, theta):
"""
Combines the creation of the Barnes-Hut tree and the force calulation using that tree.
......@@ -283,7 +223,7 @@ def force_barneshut(pos, mass, theta):
Potential energy
"""
tree = create_tree(pos, mass)
force, U = force_tree(theta, tree, pos)
force, U = force_node(pos, tree, theta, np.arange(pos.shape[0]), pos.shape)
# return
return force, -U
\ No newline at end of file
return force, U
\ No newline at end of file
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment