8
\$\begingroup\$

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

  1. Avoid complexity
  2. Avoid library imports
  3. 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

  1. Every node and body is identified by a global ID.
  2. A body is added to the quadtree by registering its mass, position and velocity. It is assigned an ID (called bid).
  3. A node is made by specifying the South-West vertex and box length (an ID is issued here too - called nid).
  4. updatecom is called for updating the quadtree's (or individual nodes') centre of mass.
  5. walk is called throughout the script for navigating the quadtree.
  6. 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)
\$\endgroup\$

2 Answers 2

5
\$\begingroup\$

You may well have a good reason for using Python 2.x. If you don't, I'd suggest that using the latest version is more Pythonic. In particular, as you describe yourself as a beginner, it may be easier to use the latest version now rather than having to unlearn certain quirks of 2.x at some point in future. I would also describe myself as a beginner, and of course I'm biased by having learned Python 3.3 first...

If you prefer to keep it as Python 2.x it might be worth adding the tag to attract answers from people expert in it.

"""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.
"""

You are right, this comprehensive docstring does indeed make the code longer. I don't see that as a bad thing. Anyone in a hurry can skip over it easily as it is self contained, and anyone confused by the code can refer back to this rather than having to work it all out for themselves. Helpful docstrings are definitely Pythonic. In particular, I wouldn't have been able to write this review without the insight provided by the docstrings.

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.
    """

The misspelling of 'Consequtive' (correct spelling Consecutive) will only have a minor effect on someone reading the code, but I point it out since Python is all about readability.

    global ID
    if reset:
        ID = 1
    while any(ID in maps for maps in (masses, childnodes)):
        ID += 1
    return ID

This function is called when assigning a new nid and also when assigning a new bid. I'm not sufficiently familiar with your algorithm to say for certain, but I'm guessing that the numbering of nodes (nid) and the numbering of bodies (bid) are unrelated. If so, then there is no need to check for an existing node id in masses, and no need to check for an existing body id in childnodes. I would be inclined to replace the function newid with two similar functions newnid and newbid. This will make assigning new ids slightly more efficient, and allow all numbers to be available (currently any given number can only be used as either a nid or a bid, not both, if I understand correctly).

This will also be more readable, as both functions will be simpler and more intuitive than newid. Currently the question of why a new nid can't have the same number as an existing bid makes the algorithm appear more complex than it needs to.

If you split the function this way, you will also need to have two global variables so that they don't share ID. For example, nextBodyID and nextNodeID. I would avoid all upper case as this is usually used for constants.

def initiate(vertex, length):
    """initiate(vertex, length)

    Initiates the quadtree by specifying the SW vertex and side length.
    """
    lengths[0] = length
    vertices[0] = vertex

You don't need to include the name of the function at the start of the docstring. The suggested approach is to start with a one line description of the function. Changing this will save you two lines per function, making your code more readable while still keeping all of the useful content of the docstrings. I mention this to save you some space. If you want more detail on suggested docstring structure, see PEP 257.

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

These three functions are very simple and easy to read. They also keep the rest of the code uncluttered and easier to read.

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)

The only thing I would change here is to replace newid with newbid, as described earlier.

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'

The code in this function is clear, but the usage of the function is potentially misleading. Since the function is called canfit I'm expecting it to return True if the body can fit, but instead it returns EMPTY. The return values that you use are obviously a good fit to the rest of your program, so I would suggest simply changing the name of this function to avoid any confusion. For example, you could call it node_state since all the return values are describing the state of the node. Or nodestate for consistency with your naming style. Personally I find separated words slightly more readable, but a consistent approach is also useful.

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

This is all readable. The only thing I would change here is to replace newid with newnid, as described earlier.

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

The only thing I would change here is to replace canfit with node_state, as described earlier.

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

If you want to avoid holding the queue in memory (as your comment suggests), you can use recursion but this will give you a top down depth first walk rather than a top down breadth first walk which you have here. It will still be top down, if that's all you require. The order based on your numbering in the docstring would be: 0, 1, 2, 3, 5, 6, 7, 13, 14, 15, 16, 8, 4, 9, 10, 11, 12. If you require the exact order shown in your docstring then I can't think of a more straightforward way than the queue you use.

    # This part is recursive.
    else:
        for child in childnodes[top]:
            for node in walk(child, topdown, gettop=True):
                yield node
        if gettop:
            yield top

The recursive part only happens if topdown is False. You don't need to say walk(child, topdown ... - you can just say walk(child, False ....

Similarly you don't need to say gettop=True as this is the default anyway. You could simply say walk(child, False). This is entirely up to you. If you want to leave the default showing explicitly that is just as correct. Most readable might be walk(child, topdown=False).

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)

This all looks good. Will you ever have need to recalculate all of the nodes above a given child node (its parent nodes rather than its child nodes)? If not then this function covers all you need.

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

This function isn't used in your code (you said you had pasted in the whole program), so I can't review its usage. It seems to do what the docstring describes perfectly well. I would question why a single function is calculating both kinds of distance. If you ever need to use just one or other of the distance types, you may improve efficiency by having two separate functions, for example xy_distances and euclid_distance.

If you intend to only call this from another function that uses both distances together, then the combined function is probably more efficient, as it avoids having to calculate the vertical and horizontal distances twice.

I think it's perfectly readable either way.

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
    """

If you use more than one line for a docstring, leaving the second line blank (as you have in all the others) will allow ease of automation. For example, Python's help function will still work without the blank line but the output won't be quite as tidy.

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

A minor point in case it's of interest: writing if i%2 == 0: is equivalent to writing if not i%2:. If you swap the order of the results you can simply say if i % 2:.

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)

I would describe your code as pretty Pythonic already. If you want any further detail on suggested style I'd recommend PEP 8.

\$\endgroup\$
2
  • 2
    \$\begingroup\$ Thank you very much! As you've mentioned, my prolonged use of python 2.7 made me uneducated about v 3.x, but I must begin to use it anyway. You've thoroughly reviewed this code and to answer some of the questions you've asked: the distance function is for future implementation of acceleration. updatecom doesn't have update functionality for parent nodes (might require an entire makeover, but not needed at the moment). Your suggestions on changing names, spellings, newnid and newbid and others are excellent! I'll modify the code accordingly. \$\endgroup\$ Mar 13, 2014 at 9:00
  • 2
    \$\begingroup\$ @user3058846 glad to hear it helped. It's still perfectly valid to use Python 2.7. If you wanted to try out Python 3.3 all you'd need to change in this code is to use range instead of xrange (identical usage here) and to use print as a function: print(x) instead of print x. Everything else works in either version, except you'll no longer need to import from future... \$\endgroup\$ Mar 13, 2014 at 9:50
3
\$\begingroup\$

As general comments, I would find more pythonic to avoid the use of unique identifier. You can implement the same algorithm using the quadrants' informations ((x,y),width) directly. This would produce a more readable code by avoiding indirections and the need to generate indices.

Additionally, I would find more pythonic to propose an interface similar to existing data structure in python. I would find easier to use qt = QuadTree(points) to build the tree and attach a set of points; to be able to use qt.append(point) instead of attach; to be able to iterate over points (like [point for point in quadtree]) or even to permits query as an item getter: points = quadtree[my_quadrant].

Incidentally, this last example shows the advantage of using quadrants themselves instead of indices.

To do so, either you build a proxy class that just call the correct functions while showing a pythonesque API, either you directly use a class embedding your data structures. I would find the later better than the existing global variables, because most of your functions does just depend on your very specific data structures and can probably not be used in other algorithms.

def canfit(body, node):
    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'

Here you may use python's enum.Enum for the returned status. Also, an 'OUT' status may help you when you will do queries (as in the attach function).

The way you test if the position is within the quadrant may be more sensible to floating point errors. Maybe it would be better to test if nx <= bx <= nx+l … (to be checked).

def attach(body, node=0):
    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

Here you may avoid walking the whole tree to find a node that fit the given body. You should avoid going through the nodes for which the points is out of scope. Just test if the body can fit within the node and proceed to its children only if it's the case.

    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)

As a user, I would want to be able to initialize the tree without having to pre-compute the quadrant. To give a set of points should be sufficient for the algorithm to compute the root quadrant by itself.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

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