Skip to content

Instantly share code, notes, and snippets.

@oir
Last active Feb 9, 2016
Embed
What would you like to do?
# Given a binary tree, this routine computes new coordinates for visualization purposes
# (reduces overlap between nodes).
# ix (paix) should be np.arrays that have integer indices for nodes (parent nodes).
# paix of root should be -1.
# depth is a np.array with integer value of the depth, and
# coords is 2d initial coordinates of each node.
# It works as follows, as a physics simulation:
# - y coordinates are fixed by looking at the depth (not free).
# - Every node is a charged particle that pushes nearby particles away (at the same depth).
# - Connection between a parent and children is like a string: parent is trying to pull its children.
# - This system is simulated until it reaches equilibrium: Forces turn into acceleration, which turns
# into velocity, which turns into translation.
# - When the system is in equilibrium (i.e. velocities are small enough) it usually gives a decent tree.
# Sometimes particles can move too fast and pass each other, which can give weird trees with parents
# on the right and children on the left. To avoid that, the order with respect to each depth value is
# explicitly preserved.
def compute_coords(ix, paix, depth, coords):
rangex = 100.
cpush = 2e-1
cspring = 1e-4
N = np.shape(ix)[0]
acc = np.zeros((N))
vel = np.zeros((N))
# will fix the order per depth, i don't want any inversions
# when particles are moving too fast
sorted_per_depth = {}
for i in range(N):
d = depth[i]
if d not in sorted_per_depth:
sorted_per_depth[d] = [i]
else:
l = sorted_per_depth[d]
for j in range(len(l)):
if (j+1) == len(l):
sorted_per_depth[d].append(i)
break
if coords[l[j],0] > coords[i,0]:
sorted_per_depth[d].insert(j,i)
break
print 'Starting to compute coords...'
while True:
#compute push forces
for i in range(N):
for j in range(N):
if i != j:
if depth[i] == depth[j] and abs(coords[i,0]-coords[j,0]) < rangex:
assert coords[i,0] != coords[j,0]
k = -1. if coords[j,0] > coords[i,0] else 1.
acc[i] += k * min(cpush / (coords[i,0]-coords[j,0])**2, 200)
#compute pull forces
for i in range(N):
if paix[i] != -1:
j = paix[i]
dx = abs(coords[i,0]-coords[j,0])
assert dx > 0
k = 1. if coords[j,0] > coords[i,0] else -1.
acc[i] += k * cspring * (dx ** 2) # pow2 might be too powerful?
acc[j] -= k * cspring * (dx ** 2)
#update velocities
vel += acc
vel = np.maximum(vel, -200)
vel = np.minimum(vel, 200)
sys.stdout.write('%.4f' % (np.max(vel)) + ' ' + '%.4f' % (np.max(acc)) + ' \r')
sys.stdout.flush()
#translate
coords[:,0] += vel
vel[:] *= 0.9 #lose velocity to friction
acc[:] = 0.
# if there are inversions because the particles were too fast, fix
# depthwise left-to-right order should be preserved
for d, l in sorted_per_depth.items():
#print l
inversion = True
if len(l) == 1:
continue
while inversion:
inversion = False
for i in range(len(l)-1):
if (coords[l[i],0] > coords[l[i+1],0]):
inversion = True
coords[l[i],0],coords[l[i+1],0] = coords[l[i+1],0],coords[l[i],0]
vel[l[i]] *= -0.1
vel[l[i+1]] *= -0.1 # as if they collided and bounced back
break
#test for stopping condition
if np.max(vel) <= 1e-2:
break
print ''
print 'Done.'
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment