File:Barnes-Hut tree construction.png

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.cm as cm

def cube_plot(ax, pos, level, idx):
    d = 1./2**level
    idx = [ idx[0]*d, idx[1]*d ]
    count = np.sum((pos[:,0] - idx[0] >= 0)*(pos[:,0] - idx[0] < d)*(pos[:,1] - idx[1] >= 0)*(pos[:,1] - idx[1] < d))
    if count == 0:
        return count
    
    color = cm.nipy_spectral(level/7.)
    ax.plot([idx[0]  , idx[0]+d], [idx[1]  , idx[1]  ], lw=1, color=color, alpha=1)
    ax.plot([idx[0]  , idx[0]  ], [idx[1]  , idx[1]+d], lw=1, color=color, alpha=1)
    ax.plot([idx[0]+d, idx[0]+d], [idx[1]  , idx[1]+d], lw=1, color=color, alpha=1)
    ax.plot([idx[0]  , idx[0]+d], [idx[1]+d, idx[1]+d], lw=1, color=color, alpha=1)
    return count

def plot_rec(ax, pos, Lmax, level, idx):
    count = cube_plot(ax, pos, level, idx)
    if level +1 == Lmax:
        return
    if count > 1:
        plot_rec(ax, pos, Lmax, level+1, [idx[0]*2, idx[1]*2])
        plot_rec(ax, pos, Lmax, level+1, [idx[0]*2+1, idx[1]*2])
        plot_rec(ax, pos, Lmax, level+1, [idx[0]*2, idx[1]*2+1])
        plot_rec(ax, pos, Lmax, level+1, [idx[0]*2+1, idx[1]*2+1])

n = 48
np.random.seed(123)

pos = np.random.rand(n, 2)

for Lmax in range(1, 7):
    fig = plt.figure( figsize=(3, 3) )
    plt.subplots_adjust(left=0.05, right=0.95, bottom=0.05, top=0.95)
    ax = fig.add_subplot(111, aspect=1)
    ax.set_xlim([0, 1])
    ax.set_ylim([0, 1])
    plt.tick_params(axis='x',which='both',bottom=False,top=False,labelbottom=False)
    plt.tick_params(axis='y',which='both',left=False,right=False,labelleft=False)

    ax.scatter(pos[:,0], pos[:,1], c="indianred", s=5)
    plot_rec(ax, pos, Lmax, 0, [0,0])

    plt.savefig("bh{}.png".format(Lmax))
    plt.close()