# N-dim Barnes-Hut algorithm for data vizualization

I've recently written an implementation of Barnes-Hut Tree algorithm for n-dimensions. All python implementations I've found were working for 2 dimensions, while I needed on working for 3... So, while I might've just written one for 3 dimensions only, I've decided to write one that works for all dimensions.
I'm going to need it for a n-dim ForceAtlas2 I'm trying to build.

Usage:

• User initializes BHTree by creating first Node:

tree = ndbh.Node(pos = [0, 0, 0], length = 10)

• Then builds up a tree:

for i in range(100000):
tree.fit([10 * np.random.random(3) - 10, np.random.random()])

• Then calculates all centers of mass:

tree.calculate_coms()

• Finally, finds all neighbors of any given body, which are returned in a list of (position, mass) tuples:

tree.neighbors([1, 2, 3])


It's the first somewhat-complete-ish, and completely amateurish, python module I've written. Any feedback will be greatly appreciated!

I have a feeling that all this input sanitation substantially slows the code. Any ideas on how to handle that better would be super helpful.

Here's the code:

"""Implementation of Barnes-Hut tree algorithm (quadtree, octatree...) in n-dimensions.

Barnes-Hut algorithm is an approximation algorithm for performing an n-body simulation.
BHTree generation recursively divides n-dim space into cells, which contain 0 or 1 bodies.
This algorithm is used to approximate forces acting on a body. Group of bodies sufficently away
from queried body can be approximated to one center of mass.

See:
https://en.wikipedia.org/wiki/Barnes%E2%80%93Hut_simulation
http://arborjs.org/docs/barnes-hut

See example implementation:

Terminology:
1.  Node            - Basic element of BHTree structure. Sector can be either:
a.  Empty       - Doesn't contain any bodies,
b.  External    - Contains precisely one body,
c.  Internal    - Contains 2**ndim children nodes.
2.  Body            - Element fit into BHTree topology. It has position in n-dimensions and mass.
"""

import numpy as np

class Node:
"""Node is a basic element of a BHTree. It is described by its:
a. position     - position of center of the cell in n-dimensional space,
b. length       - cell extends by length/2 in every direction.
c. type         - EMPTY, EXTERNAL, INTERNAL
d. body         - for EXTERNAL nodes only.
e. children     - for INTERNAL nodes only. Children are stored in an array with length of 2**ndim.
f. mass,
g. center of mass.
"""

def __init__(self, pos, length):
if isinstance(pos, list):
pos = np.array(pos)
if isinstance(pos, tuple):
pos = np.array(list(pos))
assert isinstance(pos, np.ndarray), "Position should be either a numpy.ndarray, list, or a tuple."

assert (isinstance(length, float) or isinstance(length, int)), "Length should be either a float, or int."

self.pos = pos
self.length = length
self.type = "EMPTY"
self.body = None
self.children = None
self.com = pos
self.mass = 0

def fit(self, body):
"""Fits a body into the node.
Recognized inputs:
ndbh.Body                - body object      + list of multiple ndbh.Body objects
[list, float]            - position, mass   + list of multiple [list, float] objects
(list, float)            - position, mass   + list of multiple (list, float) objects
[numpy.ndarray, float]   - position, mass   + list of multiple [numpy.ndarray, float] objects
(numpy.ndarray, float)   - position, mass   + list of multiple (numpy.ndarray, float) objects
"""

# input sanitation:
if isinstance(body, Body):
bodies = [body]
elif isinstance(body, tuple):
bodies = [Body(body, body)]
elif isinstance(body, list):
try:
if isinstance(body, float):
bodies = [Body(body, body)]
else:
bodies = []
for body in body:
if isinstance(body, Body):
bodies.append(body)
else:
bodies.append(Body(body, body))
except IndexError:
raise AssertionError("Body format not recognized.")
else:
raise AssertionError("Body format not recognized.")
# input sanitation END

for body in bodies:

assert len(body.pos) == len(self.pos), "Body and node dimensionality don't match."

# first, check for out of bounds
bounds_max = self.pos + self.length * 0.5
bounds_min = self.pos - self.length * 0.5
if any(body.pos > bounds_max) or any(body.pos < bounds_min):
raise AssertionError("Body is out of bounds!")

def child_node_index(body):
"""Returns an index of a child node from self.children to put body into"""

# evaluate position of body relative to node's center
ndim = len(self.pos)
relative_pos = np.array(body.pos > self.pos, dtype=int)
multiplier = np.array([2 ** (ndim - 1 - i) for i in range(ndim)])
index = sum(relative_pos * multiplier)
return index

if self.type == "EMPTY":
self.type = "EXTERNAL"
self.body = body

elif self.type == "EXTERNAL":

# first check if new body has the same position as the occupant
if np.array_equal(self.body.pos, body.pos):
self.body += body
return

# DIVIDE SELF
# calculate new centers
ndim = len(self.pos)
offset = self.length * 0.25
centers = []
for i in range(2 ** ndim):
s = np.binary_repr(i, width=ndim)  # creates strings like '000', '010', '111' (for ndim=3)
pos = self.pos + [(lambda c: offset if c == '1' else -offset)(c) for c in s]
centers.append(pos)

self.children = [Node(i, self.length * 0.5) for i in centers]

# find new place for occupant body
idx = child_node_index(self.body)
self.children[idx].fit(self.body)
self.body = None
self.type = "INTERNAL"

if self.type == "INTERNAL":
idx = child_node_index(body)
try:
self.children[idx].fit(body)
except RecursionError:
# just add to existing body
self.children[idx] += body

def summary(self, include_empty=False, _final=True):
"""Returns node and all its children in a dictionary form. For debugging / un-black-boxing purposes."""

return_dict = {'type': self.type, 'pos': str(self.pos.tolist())}

if self.type != "EMPTY":
return_dict['center_of_mass'] = str(self.com.tolist())
return_dict['mass'] = self.mass
return_dict['length'] = self.length

if self.type == "INTERNAL":
children = []
for child in self.children:
if child is None:
continue
if (not include_empty) & (child.type == "EMPTY"):
continue
children.append(child.summary(_final=False, include_empty=include_empty))
return_dict['children'] = children

if _final:
import json
return json.dumps(return_dict, indent=4)
else:
return return_dict

def calculate_coms(self):
"""Calculates centers of mass for all nodes."""

nodes = self._get_all_nodes()
from operator import attrgetter
sorted_nodes = sorted(nodes, key=attrgetter("length"))
for node in sorted_nodes:
node._calculate_center_of_mass()

def _get_all_nodes(self):
"""Used for calculate_coms(). Returns node and all its children's nodes."""

nodes = []
if self.type == "INTERNAL":
for child in self.children:
nodes += child._get_all_nodes()

nodes.append(self)
return nodes

def _calculate_center_of_mass(self):
"""Used for calculate_coms(). Calculates a center of mass of one node."""

if self.type == "EMPTY":
self.com = self.pos
self.mass = 0
elif self.type == "EXTERNAL":
self.com = self.body.pos
self.mass = self.body.mass
else:
sum_pos = np.zeros(len(self.pos))
sum_mass = 0
for child in self.children:
if child.type == "EMPTY": continue
if (child.mass == 0) & (child.type == "EXTERNAL"):
if child.occupant.mass != 0:
print("Error: Child seems to have wrongly calculated mass/center of mass. Recalculating.")
child._calculate_center_of_mass()
sum_pos += child.com * child.mass
sum_mass += child.mass
self.com = sum_pos / sum_mass
self.mass = sum_mass

def neighbors(self, body, theta=0.75):
"""Returns a list of (position = numpy.ndarray, mass = float) tuples of bodies/nodes affecting a given body.
Distance is controlled by theta. Lower theta = faster search = less accurate.
Recognized inputs:
ndbh.Body                - body object
list                     - position
numpy.ndarray            - position
[list, float]            - position, mass
(list, float]            - position, mass
[numpy.ndarray, float]   - position, mass
(numpy.ndarray, float]   - position, mass
"""

# input sanitation:
if (isinstance(body, list)) or (isinstance(body, tuple)):
if isinstance(body, list) or (isinstance(body, np.ndarray)):
assert len(body) == 2, "Body format not recognized."
body = Body(body, body)
else:
body = Body(body, 0)
elif isinstance(body, np.ndarray):
body = Body(body, 0)
assert isinstance(body, Body), "Body format not recognized."
# input sanitation END

neighbors = []
if self.type == "EXTERNAL":
if self.body == body:
pass
neighbors = [(self.com, float(self.mass))]
elif self.type == "INTERNAL":
dist = np.linalg.norm(body.pos - self.com)
if self.length / dist < theta:
neighbors = [(self.com, float(self.mass))]
else:
for child in self.children:
neighbors += child.neighbors(body=body, theta=theta)

return neighbors

def __repr__(self):
return "<ndbh.Node: %s at %s, length: %d>" % (self.type, self.pos, self.length)

class Body:
"""Body is an object populating Nodes. It is described by its:
a. position     - Position in n-dimensional space,
b. mass.
"""

def __init__(self, pos, mass):
if isinstance(pos, list):
pos = np.array(pos)
if isinstance(pos, tuple):
pos = np.array(list(pos))
assert isinstance(pos, np.ndarray), "Position should be either a numpy.ndarray, list, or a tuple."
assert (isinstance(mass, float) or isinstance(mass, int)), "Mass should be either a float, or int."

self.pos = pos
self.mass = mass

def __eq__(self, other):
return np.array_equal(self.pos, other.pos) and self.mass == other.mass

if isinstance(other, self.__class__):
return Body(self.pos, self.mass + other.mass)
else:
raise TypeError("unsupported operand type(s) for +: '{}' and '{}'").format(self.__class__, type(other))

def __repr__(self):
return "<ndbh.Body: %s, mass: %d>" % (self.pos, self.mass)


A few suggestions.

### Pep8:

Python has a strong idea of how the code should be styled, and it is expressed in pep8.

I suggest you get a style/lint checker. I use the pycharm ide which will show you style and compile issues right in the editor.

The primary violation was due to line length issues. Breaking things to 80 columns can seem a bit uncomfortable at first, but it makes the code easier to read.

### Input Sanitation

I would consider moving the input sanitation in Node.fit() and Node.neighbors() to class methods on Body. Since this is all about various representations of body or bodies, it make more sense to me to associate the code with that class.

I would suggest calling them something like:

@classmethod
def list_from(body_list_like):
""" Create list of Body's """

@classmethod
def create(body):
""" Create a Body from various forms """


They could then be used something like:

for i in range(100000):
tree.fit(Body.list_from([10 * np.random.random(3) - 10, np.random.random()]))


and,

tree.neighbors(Body.create([1, 2, 3]))


An additional benefit besides more closely associating the sanitation code with its class, is performing the sanitation on the input to the recursive code, and not require sanitation on bodies that have already been sanitized.

### isinstance can take a tuple of types

This:

isinstance(length, float) or isinstance(length, int))


can be simplified to the equivalent:

isinstance(length, (float, int))


### Comprehension opportunity:

neighbors = []
for child in self.children:
neighbors += child.neighbors(body=body, theta=theta)


Could be expressed as:

neighbors = sum([child.neighbors(body=body, theta=theta)
for child in self.children], [])

• Thanks! Changing neighbors = [] for child in self.children: neighbors += child.neighbors(body=body, theta=theta) to neighbors = [child.neighbors(body=body, theta=theta) for child in self.children], changes output to ...[(array([-4.08196309, 1.20883868, -1.57174241]), 1.40181811869958)], [[], [(array([-4.04851758, 3.55333593, -2.51482532]), 0.170661059987947)], [],...instead of neat ..(array([-4.36858193, 2.38901721, -0.28639853]), 1.1452272872358131), (array([-2.79813139, 1.5529723 , -2.2724586]), 0.49773052093412)... Should I leave it then? – Ankhe Jun 3 '17 at 17:58
• Sorry, this is what I get for not running the code. Should be fixed now... – Stephen Rauch Jun 3 '17 at 20:12