Plot lattice tree in Python

I'm using this piece of code which gives a pretty good result:

from matplotlib import pyplot as plt
import numpy as np

fig = plt.figure(figsize=[5, 5])
for i in range(3):
    x = [1, 0, 1]
    for j in range(i):
        x.append(0)
        x.append(1)
    x = np.array(x) + i
    y = np.arange(-(i+1), i+2)[::-1]
    plt.plot(x, y, 'bo-')
plt.show()

This is an older post, but I recently needed the same solution. This is my proposal with the major improvements: the code is reusable, we have more control on the parameters, and we use a true scale of values on the plot.

We need to import from numpy, matplotlib, and scipy:

import numpy as np
import matplotlib.pyplot as plt
from scipy.sparse import lil_matrix

To plot a tuple tree like t = ((4,), (3, 5,), (2, 4, 6,), (1, 3, 5, 7,)) we start from the linear case where the up and down steps are equally spaced and can step up and down by vola=1.

We start from defining the base (bottom) of the tree and build up the values up that are spaced by 2*vola on each time step. Then we store the values of the tree in a sparse matrix. The address of the matrix represents the grid. At the address we store the value.

Note: mat.toarray() is just a visual representation of the sparse matrix. The base of the tree is visible on the diagonal starting from the center going to the bottom left. Each row is a time step. Up and down are flipped right to left.

nb_steps = 3

vola = 1
r = 4 - (np.arange(nb_steps+1) * vola) # the tree root starts at 4
print(r, "\n")
# [4 3 2 1] 

mat = lil_matrix((nb_steps+1, nb_steps*2+1), dtype = np.float32)
for i in range(nb_steps+1):
    value = r[i] + 2*vola * np.arange(0,i+1)
    rows = np.zeros(i+1, dtype=np.int16) + i
    cols = np.arange(i+1)*2+nb_steps-i
    mat[rows, cols] = value

print(mat.toarray())
# [[0. 0. 0. 4. 0. 0. 0.]
#  [0. 0. 3. 0. 5. 0. 0.]
#  [0. 2. 0. 4. 0. 6. 0.]
#  [1. 0. 3. 0. 5. 0. 7.]]

Once we have defined the tree in the sparse matrix, the following snippet, plots the segments that join the nodes of the tree, where x is the steps axis (or time axis if you will) and v is the value at that node.

fig = plt.figure(figsize=[5, 5])
for i in range(nb_steps):
    x = (np.arange(3 + i*2) % 2 == 0) + i
    y = np.arange(0, 2*(i+1)+1) -1 + (nb_steps-i) #[::-1]
    v = mat[x, y].toarray()[0]
    plt.plot(x, v, 'bo-')

This will produce the following figure:

We can generalize to a non-linear example. Let's use the interest rates binomial tree model (often used for bond pricing - a similar model is used for equity options pricing). The steps up and down are not equally spaced, their ratio is, i.e. the steps are spaced by exp(2*vola).

As before, we need to define the base of the tree first. Then, we store the values of the tree steps in a sparse matrix.

In the following snippet we change these lines:

• the definition of r = 0.03 *1/(1 + (np.arange(nb_steps+1) * 0.15) )
• and the definition of value = r[i] * np.exp(vola*2*np.arange(0,i+1))
• increase the number of steps nb_steps = 4

That is:

nb_steps = 4

vola = 0.15
r = 0.03 *1/(1 + (np.arange(nb_steps+1) * 0.15) )
print(r, "\n")
# [0.03       0.02608696 0.02307692 0.02068966 0.01875   ] 

mat = lil_matrix((nb_steps+1, nb_steps*2+1), dtype = np.float32)
for i in range(nb_steps+1):
    value = r[i] * np.exp(vola*2*np.arange(0,i+1))
    rows = np.zeros(i+1, dtype=np.int16) + i
    cols = np.arange(i+1)*2+nb_steps-i
    mat[rows, cols] = value

print(np.round(mat.toarray(), 4))

# [[0.     0.     0.     0.     0.03   0.     0.     0.     0.    ]
#  [0.     0.     0.     0.0261 0.     0.0352 0.     0.     0.    ]
#  [0.     0.     0.0231 0.     0.0312 0.     0.042  0.     0.    ]
#  [0.     0.0207 0.     0.0279 0.     0.0377 0.     0.0509 0.    ]
#  [0.0188 0.     0.0253 0.     0.0342 0.     0.0461 0.     0.0623]]

To do the plot, we use the very same snippet as before.

fig = plt.figure(figsize=[5, 5])
for i in range(nb_steps):
    x = (np.arange(3 + i*2) % 2 == 0) + i
    y = np.arange(0, 2*(i+1)+1) -1 + (nb_steps-i) #[::-1]
    v = mat[x, y].toarray()[0]
    plt.plot(x, v, 'bo-')

which produces: