def equal_arc_layout(T, B):
"""
@param T: tree topology
@param B: branch lengths
@return: a map from vertex to location
"""
# arbitrarily root the tree
R = Ftree.T_to_R_canonical(T)
r = Ftree.R_to_root(R)
# map vertices to subtree tip count
v_to_sinks = Ftree.R_to_v_to_sinks(R)
v_to_count = {}
for v in Ftree.R_to_postorder(R):
sinks = v_to_sinks.get(v, [])
if sinks:
v_to_count[v] = sum(v_to_count[sink] for sink in sinks)
else:
v_to_count[v] = 1
# create the equal arc angles
v_to_theta = {}
_force_equal_arcs(
v_to_sinks, v_to_count, v_to_theta,
r, -math.pi, math.pi)
# convert angles to coordinates
v_to_source = Ftree.R_to_v_to_source(R)
v_to_location = {}
_update_locations(
R, B,
v_to_source, v_to_sinks, v_to_theta, v_to_location,
r, (0, 0), 0)
return v_to_location
def RB_to_v_to_age(R, B):
"""
@param R: directed topology
@param B: branch lengths in time units
@return: map from vertex to age
"""
sources, sinks = zip(*R)
leaves = set(sinks) - set(sources)
v_to_age = dict((v, 0) for v in leaves)
v_to_source = Ftree.R_to_v_to_source(R)
for v in Ftree.R_to_postorder(R):
p = v_to_source.get(v, None)
if p is not None:
v_to_age[p] = v_to_age[v] + B[frozenset([v, p])]
return v_to_age