An issue with using this method for binomial trees is that we can’t have different nodes with the same label. If two labels were T=2-$0 there would not be two nodes but one for both of them. We can fix this by manually choosing labels, but first check what happens if we do not fix it.

#An issue with using this method for binomial trees is that we can't have different nodes with the same label
#So if two labels were T=2-$0 there would not be two nodes but one
#We can fix this by manually choosing labels, but see what happens otherwise
#For example, if we only wanted it to be the time labels....

node_levels = create_nodes(100, 3, .1)
for i in range(len(node_levels)):
    node_levels[i] = [x[:3] for x in node_levels[i]]
print(node_levels)

pos_dictionary = create_node_positions(node_levels)
draw_binary_tree(node_levels, pos_dictionary)

The first part of fixing this is to make a separate labels object.

#The solution is to make node labels and assign them manually

#Create a separate object for node_labels
node_levels = create_nodes(100, 3, .1)
node_labels = []
for i in range(len(node_levels)):
    node_labels.append([x[:3] for x in node_levels[i]])
print(node_labels)

Next we need to map each node to the node label.

#Then create a dictionary which maps each label to the position of the node

node_label_dictionary = {}
for i1 in range(len(node_levels)):
    for i2 in range(len(node_levels[i1])):
        node_label_dictionary[node_levels[i1][i2]] = node_labels[i1][i2]
print(node_label_dictionary)

Let’s formalize this into a function.

#Turn it into a function
def build_label_dictionary(node_levels, node_labels):
    node_label_dictionary = {}
    for i1 in range(len(node_levels)):
        for i2 in range(len(node_levels[i1])):
            node_label_dictionary[node_levels[i1][i2]] = node_labels[i1][i2]
    return node_label_dictionary
print(build_label_dictionary(node_levels, node_labels))

Add a function to draw labels first, then we can update our function for drawing binary trees.

#Function to draw labels
def draw_node_labels(G, node_label_dictionary, label_positions):
    nx.draw_networkx_labels(G,label_positions, node_label_dictionary, font_weight='bold', font_size=15)

def draw_binary_tree(node_levels, pos_dictionary, node_labels):
    G = nx.Graph()
    draw_nodes_edges(G, node_levels)
    #Switch with_labels to False
    nx.draw(G, pos=pos_dictionary, with_labels=False, font_weight='bold', font_size=15)
    #Create the label dictionary
    node_label_dictionary = build_label_dictionary(node_levels, node_labels)
    #Label the nodes with the position dictionary and label dictionary
    draw_node_labels(G, node_label_dictionary, pos_dictionary)
    plt.xlim([-.5,.25+.5*len(node_levels)])
    plt.ylim([-.5-.5*len(node_levels),.5+.5*len(node_levels)])
    plt.show()

node_levels = create_nodes(100, 3, .1)
node_labels = []
for i in range(len(node_levels)):
    node_labels.append([x[:3] for x in node_levels[i]])
pos_dictionary = create_node_positions(node_levels)
draw_binary_tree(node_levels, pos_dictionary,node_labels)

As a final update, we could shift labels upwards by creating a position dictionary specific to the labels (where we shift up the node positions).

#We can also move the labels up
label_positions = {}
for key in pos_dictionary.keys():
    pos = pos_dictionary[key]
    #Add .5 to the y position
    label_positions[key] = (pos[0], pos[1]+.25)
print(label_positions)

Now update the function for building label dictionaries to also return a position dictionary for the labels. Also update the binary tree drawing function.

#Update function to also give back label_positions
def build_label_dictionary(node_levels, node_labels, pos_dictionary):
    node_label_dictionary = {}
    label_positions = {}
    for i1 in range(len(node_levels)):
        for i2 in range(len(node_levels[i1])):
            node_label_dictionary[node_levels[i1][i2]] = node_labels[i1][i2]

            #Get the position
            key = node_levels[i1][i2]
            pos = pos_dictionary[key]
            label_positions[key] = (pos[0], pos[1]+.25)
    return node_label_dictionary, label_positions

def draw_binary_tree(node_levels, pos_dictionary, node_labels):
    G = nx.Graph()
    draw_nodes_edges(G, node_levels)
    nx.draw(G, pos=pos_dictionary, with_labels=False, font_weight='bold', font_size=15)
    #The function returns label positions now
    node_label_dictionary, label_positions = build_label_dictionary(node_levels, node_labels, pos_dictionary)
    #Change the last argument to be label positions
    draw_node_labels(G, node_label_dictionary, label_positions)
    plt.xlim([-.5,.25+.5*len(node_levels)])
    plt.ylim([-.5-.5*len(node_levels),.5+.5*len(node_levels)])
    plt.show()

#Run the version
node_levels = create_nodes(100, 3, .1)
node_labels = []
for i in range(len(node_levels)):
    node_labels.append([x[:3] for x in node_levels[i]])
pos_dictionary = create_node_positions(node_levels)
draw_binary_tree(node_levels, pos_dictionary,node_labels)