While graphs are inherently difficult to layout pleasantly, algorithms based on the physical concept of force often yield good results. In this post, I’m going to implement such an algorithm from scratch while motivating each step in the process.

Consider the following screenshot.

Colorful rectangles scrambled in the center
Colorful rectangles scrambled in the center

First, we need to get the rectangles (called nodes from here on) apart. To do that, we’re going to make the nodes repel each other by applying a repulsion force to each pair of nodes.

# Repulsion force
for i in range(len(self.nodes) - 1):
    for j in range(i+1, len(self.nodes)):
        delta = self.nodes[j].position - self.nodes[i].position
        if np.any(delta):
            distance_squared = np.sum(delta**2)
            distance = np.sqrt(distance_squared)
            force = self.repulsive_force_constant / distance_squared
            f = force * delta / distance
        else:
            f = np.random.rand(2)
        self.nodes[i].force -= f
        self.nodes[j].force += f

Here we simply calculate the distance between each pair of nodes and move the two nodes of the pair apart by adding a force to each node that points away from the other node. The bigger the distance between two nodes, the smaller the force applied.

If the nodes are at exactly the same position though, both delta and distance will be 0 thus this will not work. Instead, we simply add a small force into a random direction to drive the nodes apart.

Nodes repelling each other
Nodes repelling each other

This is better, but you can see that related (linked) nodes are all over the place. We want nodes that are linked to be closer together than to other nodes. Thus, we need a force that makes linked nodes attract each other.

# Attraction force
for i in range(len(self.nodes)):
    for j in self.nodes[i].neighbors:
        if i < j:
            delta = self.nodes[j].position - self.nodes[i].position
            if np.any(delta):
                distance = np.sqrt(np.sum(delta**2))
                force = self.attraction_constant * distance
                f = force * delta / distance
                self.nodes[i].force += f
                self.nodes[j].force -= f
Linked nodes attracting each other
Linked nodes attracting each other

Wow, they are getting cozy! This isn’t bad but if I was a node I would prefer to have a little more space for myself. Let’s replace the attraction force with a spring force.

# Spring force
for i in range(len(self.nodes)):
    for j in self.nodes[i].neighbors:
        if i < j:
            delta = self.nodes[j].position - self.nodes[i].position
            if np.any(delta):
                distance = np.sqrt(np.sum(delta**2))
                force = self.spring_constant * (distance - self.spring_rest_length)
                f = force * delta / distance
                self.nodes[i].force += f
                self.nodes[j].force -= f
Nodes keeping a healthy distance
Nodes keeping a healthy distance

Much better, now this looks almost beautiful!

However, the brown node in the top is trying to escape. The purple one below it is not far behind. Not so fast! We apply a center force that makes every node wanna be the center of attention.

# Center force
for node in self.nodes:
    delta = self.center - node.position
    if np.any(delta):
        distance = np.sqrt(np.sum(delta**2))
        force = distance * self.center_constant
        f = force * delta / distance
        node.force += f

The bigger the distance between the node and the center of the window, the larger the force urging it back.

Escapees rushing back
Escapees rushing back

Excellent! Now that we have gathered our forces, how do we use them to move objects around? Here’s how to update the objects’ positions.

# Update positions
for node in self.nodes:
    delta = self.delta_t * node.force
    displacement_squared = np.sum(delta**2)
    if displacement_squared > self.max_displacement_squared:
        s = np.sqrt(self.max_displacement_squared / displacement_squared)
        delta *= s
    node.position += delta

As you can see, we limit how far a node can be displaced at a time, otherwise a large force might cause it to overshoot the node’s desired position.

All this talk about nodes, but what are they?

class ForceNode:
    def __init__(self, position, obj, neighbors):
        self.position = position
        self.obj = obj
        self.neighbors = neighbors
        self.force = None

Who creates them?

class ForceLayout:
    def __init__(self, nodes,
                 position_getter=itemgetter('position'),
                 position_setter=itemsetter('position'),
                 neighbor_getter=itemgetter('neighbors'),
                 center=(0, 0), spring_rest_length=40, repulsive_force_constant=6666,
                 spring_constant=2, delta_t=0.003, max_displacement_squared=20,
                 center_constant=0.6):
        self.spring_rest_length = spring_rest_length
        self.repulsive_force_constant = repulsive_force_constant
        self.spring_constant = spring_constant
        self.delta_t = delta_t
        self.max_displacement_squared = max_displacement_squared
        self.center_constant = center_constant

        self.position_getter = position_getter
        self.position_setter = position_setter
        self.center = np.array(center, dtype=float)

        self.nodes = [ForceNode(np.array(position_getter(x), dtype=float),
                                obj=x,
                                neighbors=neighbor_getter(x))
                      for x in nodes]
    
    def apply_forces(self):
        # Init forces
        for node in self.nodes:
            node.force = np.array([0.0, 0.0], dtype=float)
        ...
        
    def update(self):
        self.apply_forces()
        for node in self.nodes:
            self.position_setter(node.obj, node.position)

This class is designed to be reusable and makes no assumptions about the objects passed to it. Instead, you need to pass functions that tell the class how to deal with your objects.

And how do you use it with your own objects? Here’s an example with pygame.

pygame.init()

window = (600, 400)
center = (window[0]/2, window[1]/2)
screen = pygame.display.set_mode(window)

num_nodes = 20
num_links = 20

nodes = [dict(position=(0,0), size=(10,10), color=(255,0,0))
         for _ in range(num_nodes)]
links = [(random.randint(0, num_nodes-1), random.randint(0, num_nodes-1))
         for _ in range(num_links)]


def neighbor_getter(x):
    idx = nodes.index(x)
    return [i for i, n in enumerate(nodes)
            if (i, idx) in links or (idx, i) in links]


force = ForceLayout(nodes, neighbor_getter=neighbor_getter, center=center)


def node_center(node):
    return [p + s/2 for p, s in zip(node['position'], node['size'])]


while True:
    for event in pygame.event.get():
        pass

    screen.fill((255, 255, 255))

    for node in nodes:
        pygame.draw.rect(screen, node['color'], (*node['position'], *node['size']))

    for a, b in links:
        pygame.draw.line(screen, nodes[a]['color'], node_center(nodes[a]),
                         node_center(nodes[b]))

    force.update()
    pygame.display.flip()

Let’s not forget the imports.

import pygame
import random
import numpy as np
from operator import itemgetter

And a little helper function.

def itemsetter(name):
    def setter(obj, val):
        obj[name] = val
    return setter

That’s it! I leave it to you to assemble the code in the correct order.