Solution

#Define a function to build nodes given the starting stock price, the number of periods T
#and the rate for continiously compounded returns r
def create_nodes(S, T, r):
    up = np.exp(r)
    down = 1/up
    node_levels = []
    for t in range(T+1):
        nodes = [S * (up ** k) * (down ** (t-k))  for k in range(t+1)]
        nodes = ["T={}-${:.2f}".format(t, x) for x in nodes]
        node_levels.append(nodes)
    return node_levels
node_levels = create_nodes(100, 3, .1)
print(node_levels)

The first step in creating the visualization of a binomial tree is to actually add our nodes. This is simple, like follows:

#Build the first steps towards creating a binary tree

def draw_binary_tree(node_levels):
    G = nx.Graph()
    for nodes in node_levels:
        G.add_nodes_from(nodes)
    nx.draw(G, with_labels=True, font_weight='bold', font_size=15)
    plt.show()
draw_binary_tree(node_levels)

The next steps are going to be a little more complex. Essentially, we start with our first node and iterate through every level after that. Each node will have either one or two connections based on whether or not it is an outer node. The outer nodes only have one connection so we will handle them outside of the inner nodes. Don’t worry if you can’t totally follow the comments in the code.

#Update the function to add edges
def draw_binary_tree(node_levels):
    G = nx.Graph()
    #Add the first node
    G.add_nodes_from(node_levels[0])

    #Iterate through all the next i-1 levels
    for i in range(len(node_levels)-1):
        #Add the nodes
        G.add_nodes_from(node_levels[i+1])

        #Deal with outer nodes, they only can be reached by prior outer nodes
        #Connect the node from the prior level at the lowest position to the lowest position on the current level
        G.add_edge(node_levels[i][0], node_levels[i+1][0])
        #Connect the node from the prior level at the highest position to the highest position on the current level
        G.add_edge(node_levels[i][-1], node_levels[i+1][-1])

        #Now we have to go through the remaining nodes (minus the two outer nodes)
        for i2 in range(len(node_levels[i+1])-2):
            #Add the bottom node connection
            G.add_edge(node_levels[i][i2], node_levels[i+1][i2+1])
            #Add the top node connection
            G.add_edge(node_levels[i][i2+1], node_levels[i+1][i2+1])
    nx.draw(G, with_labels=True, font_weight='bold', font_size=15)
    plt.show()
draw_binary_tree(node_levels)

The plot which this draws is not in the layout that we want to see. For that part we are going to have to assign positions. The following will create a nice tree structure that we are looking for.

#Now we can define a function that defines the node positions
def create_node_positions(node_levels):
    pos_dictionary = {}

    for i in range(len(node_levels)):
        #Define the y positions
        y_positions = [i/2 - x for x in range(i+1)]
        for i2 in range(len(node_levels[i])):
            #Set the positions, they all will have i/2 as the x positions
            pos_dictionary[node_levels[i][i2]] = (i/2, y_positions[i2])
    return pos_dictionary
pos_dictionary = create_node_positions(node_levels)
print(pos_dictionary)

With the positions, add it to the function for drawing the binomial tree.

def draw_binary_tree(node_levels, pos_dictionary):
    G = nx.Graph()
    G.add_nodes_from(node_levels[0])
    for i in range(len(node_levels)-1):
        G.add_nodes_from(node_levels[i+1])
        G.add_edge(node_levels[i][0], node_levels[i+1][0])
        G.add_edge(node_levels[i][-1], node_levels[i+1][-1])
        for i2 in range(len(node_levels[i+1])-2):
            G.add_edge(node_levels[i][i2], node_levels[i+1][i2+1])
            G.add_edge(node_levels[i][i2+1], node_levels[i+1][i2+1])
    #Add in the position dictionary
    nx.draw(G, pos=pos_dictionary, with_labels=True, font_weight='bold', font_size=15)
    #Add in axis limits so we can see everything
    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)
pos_dictionary = create_node_positions(node_levels)
draw_binary_tree(node_levels, pos_dictionary)

Finally, clean up the functions to make it more easily readable including taking the edge drawing into a separate function.

#Clean up the code a bit

def create_nodes(S, T, r):
    up = np.exp(r)
    down = 1/up
    node_levels = []
    for t in range(T+1):
        nodes = [S * (up ** k) * (down ** (t-k))  for k in range(t+1)]
        nodes = ["T={}-${:.2f}".format(t, x) for x in nodes]
        node_levels.append(nodes)
    return node_levels

def create_node_positions(node_levels):
    pos_dictionary = {}

    for i in range(len(node_levels)):
        y_positions = [i/2 - x for x in range(i+1)]
        for i2 in range(len(node_levels[i])):
            pos_dictionary[node_levels[i][i2]] = (i/2, y_positions[i2])
    return pos_dictionary


#Separate a function that adds the edges

def draw_nodes_edges(G, node_levels):
    G.add_nodes_from(node_levels[0])
    for i in range(len(node_levels)-1):
        G.add_nodes_from(node_levels[i+1])
        G.add_edge(node_levels[i][0], node_levels[i+1][0])
        G.add_edge(node_levels[i][-1], node_levels[i+1][-1])
        for i2 in range(len(node_levels[i+1])-2):
            G.add_edge(node_levels[i][i2], node_levels[i+1][i2+1])
            G.add_edge(node_levels[i][i2+1], node_levels[i+1][i2+1])

def draw_binary_tree(node_levels, pos_dictionary):
    G = nx.Graph()
    draw_nodes_edges(G, node_levels)
    nx.draw(G, pos=pos_dictionary, with_labels=True, font_weight='bold', font_size=15)
    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)
pos_dictionary = create_node_positions(node_levels)
draw_binary_tree(node_levels, pos_dictionary)