drawing a binomial tree in latex with python

Question

I want to have some code that draws a binomial tree (for finance) for how ever many periods I give it.

\documentclass{article}
\usepackage{tikz}
\usetikzlibrary{matrix}
\begin{document}
\begin{tikzpicture}[>=stealth,sloped]
            \matrix (tree) [%
            matrix of nodes,
            minimum size=1cm,
            column sep=3.5cm,
            row sep=1cm,
            ]
        { &  & $S_0U^2 $ \\
 & $S_0U $ &  \\
$S_0$ &  & $S_0U D$ \\
 & $S_0 D$ &  \\
 &  & $S_0 D^2$ \\
};
\draw[->] (tree-4-2) -- (tree-3-3) node [midway,above] {$ p (1-p) $};
\draw[->] (tree-4-2) -- (tree-5-3) node [midway,above] {$  (1-p)^2 $};
\draw[->] (tree-3-1) -- (tree-2-2) node [midway,above] {$ p  $};
\draw[->] (tree-3-1) -- (tree-4-2) node [midway,above] {$  (1-p) $};
\draw[->] (tree-2-2) -- (tree-1-3) node [midway,above] {$ p^2  $};
\draw[->] (tree-2-2) -- (tree-3-3) node [midway,above] {$ p (1-p) $};
\end{tikzpicture}
\end{document} 

Something like this but that then can be extended to however my periods want ?

Answer 1

I made the following function

def make_binomial_tree(number_of_periods):
    columns = {}
    columns2 = {}
    if number_of_periods < 4:
        st = """\\begin{tikzpicture}[>=stealth,sloped]
            \\matrix (tree) [%
            matrix of nodes,
            minimum size=1cm,
            column sep=3.5cm,
            row sep=1cm,
            ]
        {"""
    else:
        st = '''\\begin{tikzpicture}[>=stealth,sloped]
            \\matrix (tree) [%
            matrix of nodes,minimum size=0.5cm,
        column sep=1cm,
        row sep=0.5cm,]{'''
    st2 = ''
    for j in range_(-number_of_periods, 1, -1):
        for i in range_(0, j+1):
            try:
                columns[j].append('')
                columns2[j].append('')
            except:
                columns[j] = ['' for i in range(number_of_periods-j)]
                columns2[j] = ['' for i in range(number_of_periods-j)]
            columns2[j].append(('U'*i+'D'*(j-i)))
            columns[j].append(
                (f"""{'$S_0U' if i==1 else f'$S_0U^{i}'if i>0 else '$S_0'} {"D$" if j-i==1 else f'D^{(j-i)}$'if j-i>0 else '$'}"""
                 if not (i == j and j == 0) else '$S_0$')
            )
        for i in range(number_of_periods-j):
            columns[j].append('')
            columns2[j].append('')
    for i in range_(-2*number_of_periods, 1, -1):
        for j in range(len(columns)):
            if j != len(columns)-1:
                st += (columns[j][i] + ' & ')
                if columns[j][i] != '':
                    U1, D1, U2, D2 = columns2[j+1][i-1].count("D"), columns2[j+1][i-1].count(
                        "U"), columns2[j+1][i+1].count("D"), columns2[j+1][i+1].count("U")
                    # adding the ligns and their values between nodes on the tree
                    st2 += (
                        f'\draw[->] (tree-{i+1}-{j+1}) -- (tree-{i}-{j+2}) node [midway,above] ') + \
                        f'{"{"}$ {"p" if U1 ==1 else f"p^{U1}" if U1 > 0 else ""} {"(1-p)" if D1 == 1 else f"(1-p)^{D1}" if D1 > 0 else ""} ${"}"};\n' + \
                        (
                        f'\draw[->] (tree-{i+1}-{j+1}) -- (tree-{i+2}-{j+2}) node [midway,above] ') +\
                        f'{"{"}$ {"p" if U2 ==1 else f"p^{U2}" if U2 > 0 else ""} {"(1-p)" if D2 == 1 else f"(1-p)^{D2}" if D2 > 0 else ""} ${"}"};\n'
            else:
                st += (columns[j][i])
        st += ' \\\\\n'
    st += '};\n' + st2 + '\end{tikzpicture}'
    return st

Hopefully this helps someone else!

Answer 2

thank you very much. Really helpful. Though I had some issues with your code (the range_ is not a standard function so it doesn't work). I rewrote the code and this one seems to work fine

def make_binomial_tree(number_of_periods):
    columns = {}
    columns2 = {}
    if number_of_periods < 4:
        st = """\\begin{tikzpicture}[>=stealth,sloped]
            \\matrix (tree) [%
            matrix of nodes,
            minimum size=1cm,
            column sep=3.5cm,
            row sep=1cm,
            ]
        {"""
    else:
        st = '''\\begin{tikzpicture}[>=stealth,sloped]
            \\matrix (tree) [%
            matrix of nodes,minimum size=0.5cm,
        column sep=1cm,
        row sep=0.5cm,]{'''
    st2 = ''
    for j in range(number_of_periods, -1, -1):
        for i in range(0, j+1):
            if j in columns:
                columns[j].append('')
                columns2[j].append('')
            else:
                columns[j] = ['' for i in range(number_of_periods-j)]
                columns2[j] = ['' for i in range(number_of_periods-j)]
                
            columns2[j].append(('U'*(j-i)+'D'*(i)))
            tmp = (f"""{'$S_0U' if j-i==1 else f'$S_0U^{j-i}'if j-i>0 else '$S_0'} {"D$" if i==1 else f'D^{(i)}$'if i>0 else '$'}"""
                 if not (i == j and j == 0) else '$S_0$')
            columns[j].append(tmp)
        for i in range(number_of_periods-j):
            columns[j].append('')
            columns2[j].append('')

    columns = dict(reversed(list(columns.items())))

    for i in range(2*number_of_periods+1):
        for j in range(len(columns)):
            if j != len(columns)-1:
                st += (columns[j][i] + ' & ')
                if columns[j][i] != '':
                    U1, D1, U2, D2 = columns2[j+1][i-1].count("U"), columns2[j+1][i-1].count(
                        "D"), columns2[j+1][i+1].count("U"), columns2[j+1][i+1].count("D")
                    # adding the ligns and their values between nodes on the tree
                    st2 += (
                        f'\draw[->] (tree-{i+1}-{j+1}) -- (tree-{i}-{j+2}) node [midway,above] ') + \
                        f'{"{"}$ {"p" if U1 ==1 else f"p^{U1}" if U1 > 0 else ""} {"(1-p)" if D1 == 1 else f"(1-p)^{D1}" if D1 > 0 else ""} ${"}"};\n' + \
                        (
                        f'\draw[->] (tree-{i+1}-{j+1}) -- (tree-{i+2}-{j+2}) node [midway,above] ') +\
                        f'{"{"}$ {"p" if U2 ==1 else f"p^{U2}" if U2 > 0 else ""} {"(1-p)" if D2 == 1 else f"(1-p)^{D2}" if D2 > 0 else ""} ${"}"};\n'
            else:
                st += (columns[j][i])
        st += ' \\\\\n'
    st += '};\n' + st2 + '\end{tikzpicture}'
    return st

When you call print(make_binomial_tree(4)) you get:

\begin{tikzpicture}[>=stealth,sloped]
        \matrix (tree) [%
        matrix of nodes,minimum size=0.5cm,
    column sep=1cm,
    row sep=0.5cm,]{ &  &  &  & $S_0U^4 $ \\
&  &  & $S_0U^3 $ &  \\
 &  & $S_0U^2 $ &  & $S_0U^3 D$ \\
 & $S_0U $ &  & $S_0U^2 D$ &  \\
$S_0$ &  & $S_0U D$ &  & $S_0U^2 D^2$ \\
 & $S_0 D$ &  & $S_0U D^2$ &  \\
 &  & $S_0 D^2$ &  & $S_0U D^3$ \\
 &  &  & $S_0 D^3$ &  \\
 &  &  &  & $S_0 D^4$ \\
};
\draw[->] (tree-2-4) -- (tree-1-5) node [midway,above] {$ p^4  $};
\draw[->] (tree-2-4) -- (tree-3-5) node [midway,above] {$ p^3 (1-p) $};
\draw[->] (tree-3-3) -- (tree-2-4) node [midway,above] {$ p^3  $};
\draw[->] (tree-3-3) -- (tree-4-4) node [midway,above] {$ p^2 (1-p) $};
\draw[->] (tree-4-2) -- (tree-3-3) node [midway,above] {$ p^2  $};
\draw[->] (tree-4-2) -- (tree-5-3) node [midway,above] {$ p (1-p) $};
\draw[->] (tree-4-4) -- (tree-3-5) node [midway,above] {$ p^3 (1-p) $};
\draw[->] (tree-4-4) -- (tree-5-5) node [midway,above] {$ p^2 (1-p)^2 $};
\draw[->] (tree-5-1) -- (tree-4-2) node [midway,above] {$ p  $};
\draw[->] (tree-5-1) -- (tree-6-2) node [midway,above] {$  (1-p) $};
\draw[->] (tree-5-3) -- (tree-4-4) node [midway,above] {$ p^2 (1-p) $};
\draw[->] (tree-5-3) -- (tree-6-4) node [midway,above] {$ p (1-p)^2 $};
\draw[->] (tree-6-2) -- (tree-5-3) node [midway,above] {$ p (1-p) $};
\draw[->] (tree-6-2) -- (tree-7-3) node [midway,above] {$  (1-p)^2 $};
\draw[->] (tree-6-4) -- (tree-5-5) node [midway,above] {$ p^2 (1-p)^2 $};
\draw[->] (tree-6-4) -- (tree-7-5) node [midway,above] {$ p (1-p)^3 $};
\draw[->] (tree-7-3) -- (tree-6-4) node [midway,above] {$ p (1-p)^2 $};
\draw[->] (tree-7-3) -- (tree-8-4) node [midway,above] {$  (1-p)^3 $};
\draw[->] (tree-8-4) -- (tree-7-5) node [midway,above] {$ p (1-p)^3 $};
\draw[->] (tree-8-4) -- (tree-9-5) node [midway,above] {$  (1-p)^4 $};
\end{tikzpicture}

and putting it into a latex document like this:

\documentclass{article}
\usepackage{tikz}
\usetikzlibrary{matrix}
\begin{document}
\begin{tikzpicture}[>=stealth,sloped]
            \matrix (tree) [%
            matrix of nodes,minimum size=0.5cm,
        column sep=1cm,
        row sep=0.5cm,]{ &  &  &  & $S_0U^4 $ \\
 &  &  & $S_0U^3 $ &  \\
 &  & $S_0U^2 $ &  & $S_0U^3 D$ \\
 & $S_0U $ &  & $S_0U^2 D$ &  \\
$S_0$ &  & $S_0U D$ &  & $S_0U^2 D^2$ \\
 & $S_0 D$ &  & $S_0U D^2$ &  \\
 &  & $S_0 D^2$ &  & $S_0U D^3$ \\
 &  &  & $S_0 D^3$ &  \\
 &  &  &  & $S_0 D^4$ \\
};
\draw[->] (tree-2-4) -- (tree-1-5) node [midway,above] {$ p^4  $};
\draw[->] (tree-2-4) -- (tree-3-5) node [midway,above] {$ p^3 (1-p) $};
\draw[->] (tree-3-3) -- (tree-2-4) node [midway,above] {$ p^3  $};
\draw[->] (tree-3-3) -- (tree-4-4) node [midway,above] {$ p^2 (1-p) $};
\draw[->] (tree-4-2) -- (tree-3-3) node [midway,above] {$ p^2  $};
\draw[->] (tree-4-2) -- (tree-5-3) node [midway,above] {$ p (1-p) $};
\draw[->] (tree-4-4) -- (tree-3-5) node [midway,above] {$ p^3 (1-p) $};
\draw[->] (tree-4-4) -- (tree-5-5) node [midway,above] {$ p^2 (1-p)^2 $};
\draw[->] (tree-5-1) -- (tree-4-2) node [midway,above] {$ p  $};
\draw[->] (tree-5-1) -- (tree-6-2) node [midway,above] {$  (1-p) $};
\draw[->] (tree-5-3) -- (tree-4-4) node [midway,above] {$ p^2 (1-p) $};
\draw[->] (tree-5-3) -- (tree-6-4) node [midway,above] {$ p (1-p)^2 $};
\draw[->] (tree-6-2) -- (tree-5-3) node [midway,above] {$ p (1-p) $};
\draw[->] (tree-6-2) -- (tree-7-3) node [midway,above] {$  (1-p)^2 $};
\draw[->] (tree-6-4) -- (tree-5-5) node [midway,above] {$ p^2 (1-p)^2 $};
\draw[->] (tree-6-4) -- (tree-7-5) node [midway,above] {$ p (1-p)^3 $};
\draw[->] (tree-7-3) -- (tree-6-4) node [midway,above] {$ p (1-p)^2 $};
\draw[->] (tree-7-3) -- (tree-8-4) node [midway,above] {$  (1-p)^3 $};
\draw[->] (tree-8-4) -- (tree-7-5) node [midway,above] {$ p (1-p)^3 $};
\draw[->] (tree-8-4) -- (tree-9-5) node [midway,above] {$  (1-p)^4 $};
\end{tikzpicture}
\end{document} 

you get: