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.


  • 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:

  • 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 example implementation:

    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[0], body[1])]
        elif isinstance(body, list):
                if isinstance(body[1], float):
                    bodies = [Body(body[0], body[1])]
                    bodies = []
                    for body in body:
                        if isinstance(body, Body):
                            bodies.append(Body(body[0], body[1]))
            except IndexError:
                raise AssertionError("Body format not recognized.")
            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

                # 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]

                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.body = None
                self.type = "INTERNAL"

            if self.type == "INTERNAL":
                idx = child_node_index(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:
                if (not include_empty) & (child.type == "EMPTY"):
                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)
            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:

    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()

        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
            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.")
                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[0], list) or (isinstance(body[0], np.ndarray)):
                assert len(body) == 2, "Body format not recognized."
                body = Body(body[0], body[1])
                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:
            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))]
                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

    def __add__(self, other):
        if isinstance(other, self.__class__):
            return Body(self.pos, self.mass + other.mass)
            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.


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:

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

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()]))


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


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], [])
| improve this answer | |
  • \$\begingroup\$ 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? \$\endgroup\$ – Ankhe Jun 3 '17 at 17:58
  • \$\begingroup\$ Sorry, this is what I get for not running the code. Should be fixed now... \$\endgroup\$ – Stephen Rauch Jun 3 '17 at 20:12

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.