 ### 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 mass = root origin = root nbody = root # numbers of the bodies. use to keep track of which body has what position r = root/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: # check if there are any bodies in node. if not: left as empty list if nbody_NW.shape > 1: # check if node is internal node = [pos[northwest], mass[northwest], origin + np.array([-r,r])/2, nbody_NW, r] # internal node elif nbody_NW.shape == 1: # check of node is external node = [(pos[northwest]).tolist(), float(mass[northwest]), int(nbody_NW)] # external node if nbody_NE.shape != 0: if nbody_NE.shape > 1: node = [pos[northeast], mass[northeast], origin + np.array([r,r])/2, nbody_NE, r] elif nbody_NE.shape == 1: node = [(pos[northeast]).tolist(), float(mass[northeast]), int(nbody_NE)] if nbody_ZE.shape != 0: if nbody_ZE.shape > 1: node = [pos[southeast], mass[southeast], origin + np.array([r,-r])/2, nbody_ZE, r] elif nbody_ZE.shape == 1: node = [(pos[southeast]).tolist(), float(mass[southeast]), int(nbody_ZE)] if nbody_ZW.shape != 0: if nbody_ZW.shape > 1: node = [pos[southwest], mass[southwest], origin + np.array([-r,-r])/2, nbody_ZW, r] elif nbody_ZW.shape == 1: node = [(pos[southwest]).tolist(), float(mass[southwest]), int(nbody_ZW)] dim = origin.shape 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 > 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 == 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) 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), 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, node) elif len(node) == (2**Mshape + 3): d = np.linalg.norm(pos - np.asarray(node), 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: force[nbody_calc], U[nbody_calc] = force_cal(pos[node[-1]/d < theta], node, node) 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), 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!