Trinomial tree in Python

Question

I am struggling with implementing trinomial tree in Python. I have found very nice solution (and vectorized version) for binomial tree and I am trying to change it for a trinomial case.
Here is what I've got:

def Trinomial(type, S0, K, r, sigma, T, N=2000):
import numpy as np
t = float(T) / N
#fixed lambda
lam = np.sqrt(5)
# up and down factor will be constant for the tree so we calculate outside the loop
u = np.exp(lam * sigma * np.sqrt(t))
d = 1.0 / u

#to work with vector we need to init the arrays using numpy
fs =  np.asarray([0.0 for i in xrange(2*N + 1)])

#we need the stock tree for calculations of expiration values
fs2 = np.asarray([(S0 * u**j) for j in xrange(-N, N+1)])

#we vectorize the strikes as well so the expiration check will be faster
fs3 =np.asarray( [float(K) for i in xrange(2*N + 1)])

#formulas, that can be found in the pdf document
a = np.exp(r*t/2.0)
b = np.exp(-sigma * np.sqrt(t/2.0))
c = np.exp(sigma * np.sqrt(t/2.0))
p_u = ( ( a - b ) / ( c - d ) )**2 
p_d = ( ( c - a ) / ( c - b ) )**2
p_m = 1 - p_u - p_d

# Compute the leaves, f_{N, j}
if type =="C":
    fs[:] = np.maximum(fs2-fs3, 0.0)
else:
    fs[:] = np.maximum(-fs2+fs3, 0.0)

#calculate backward the option prices
for i in xrange(N-1, -1, -1):
   fs[:-1] = np.exp(-r * t) * (p_u * fs[1:] + p_d * fs[:-1] + p_m*fs[-1])

return fs[0]

Of course, it dosen't work as intended, for example a call

print Trinomial("C",100, 100, 0.1, 0.1, 5, 3)

Should output something between 39 and 40.
I am most concerned about lines:

fs2 = np.asarray([(S0 * u**j) for j in xrange(-N, N+1)])

and

fs[:-1] = np.exp(-r * t) * (p_u * fs[1:] + p_d * fs[:-1] + p_m*fs[-1])

I'm not sure if I'm populating initial tree correctly, and I'm 100% sure, that reverse computing option price is not working. I don't know how to implement formula [10] from pdf I linked at the beginning.
Maybe it can't be done on a vectorized version, but I tried with simple tree and failed aswell.
I'm not using binomial or BS price in this case, because my assigment is to do it with trinomial tree.

Answer 1

I've just started picking up Python and have built binomial and trinomial models just to test my understanding, especially about arrays.

For terminal stock price array, I have stepped down from InitialStock * u**(iSteps - i) for i going from zero to 2*iSteps+ 1. ie. I start at top-most terminal stock price InitialStock * u**(iSteps) and travel down until I get to InitalStock * u**(-iSteps) and populate the TerminalStock array from 0 to 2*iSteps+1 (ie. last array element is InitialStock * u**(-iSteps). Because I come from the land of VBA, I haven't yet picked up the eloquent (and hence, brief and sometimes difficult-to-read Python style), my loops look like this:

for i in range(0, 2*iSteps+1) :

    # Terminal stock is initial with up staps and down steps
    dblStockTerminal = dblStock * u ** float(iSteps - i)

    # compute intrinsic values at terminal stock value
    dblTrinOpt[i][iSteps] = max( dblSwitch * (dblStockTerminal - dblStrike), 0 )

I have initialised dblTrinOpt array as dblTrinOpt = np.ndarray( ( 2*iSteps+1, iSteps+1), float) so that it has 2*iSteps price elements at terminal stock price and iStep time steps, where iSteps = OptionTerm / NumberOfSteps. Yes, apology for my clunky python code! (but I have to read it).

after that, my option calc is as follows (yes, again the code is inelegant):

#------------------------------------------
# steps in time from Terminal to Initial Stock price
#------------------------------------------
for i in range(iSteps-1, -1, -1) :

    #------------------------------------------
    # steps in price range from highest to lowest
    #------------------------------------------
    for j in range(0, 2*i+1) :

        #------------------------------------------
        # discount average of future option value
        # as usual, trinomial averages three values
        #------------------------------------------
        dblTrinOpt[j][i] = dblDisc * (pu * dblTrinOpt[j][i+1] + pm * \
            dblTrinOpt[j+1][i+1] + pd * dblTrinOpt[j+2][i+1])

after that, my Trinomial function passes dblTrinOpt[0][0] as final discounted option price.

I wanted to write this in a more readable form so that it became easier to add two components: (a) an American option facility, hence tracking stock prices at each tree node (and so then things like barriers and Bermudans would be an easy addition, and (b) add two additional terminal prices one at each end and so that at time zero, I'd have three stock and option values so that I could compute delta and gamma without needing to recompute the option tree (i.e.. avoid the bump-and-recompute approach). I hope that helps, even though I didn't fix your specific code, but explained the logic.