Optimize option pricing code I am new to Python and am trying to optimize the following code:

import sys
import numpy as np

def LSMPut(T, r, sigma, K, S0, TimeSteps, Paths, k):
    dt = T/TimeSteps
    t = np.arange(0, T+dt, dt).tolist()

    z=[np.random.standard_normal() for _ in range(Paths)]
    w = (r-sigma**2/2)*T + sigma*np.sqrt(T)*np.array(z)

    S = S0*np.exp(np.array(w))
    P=np.maximum(K-np.array(S),0)

    for i in range(TimeSteps-1, -1, -1):
        z=[np.random.standard_normal() for _ in range(Paths)]
        w = t[i]*np.array(w)/t[i+1] + sigma*np.sqrt(dt*t[i]/t[i+1])*np.array(z)

        S = S0*np.exp(np.array(w))
        itmPaths = [index for index,value in enumerate(K-np.array(S)) if value &gt; 0]

        itmS = S[itmPaths]
        Pt = K - np.array(itmS)

        itmDiscP = P[itmPaths]*np.exp(-r*dt)

        A = BasisFunct(itmS, k)
        beta = np.linalg.lstsq(A,itmDiscP)[0]
        C = np.dot(A,beta)

        exPaths = [itmPaths[i] for i, value in enumerate(zip(Pt, C)) if value[0] &gt; value[1]]
        restPaths = np.setdiff1d(np.arange(0, Paths-1, 1).tolist(), exPaths) # Rest of the paths

        P[exPaths] = [Pt[i] for i, value in enumerate(zip(Pt, C)) if value[0] &gt; value[1]]
        P[restPaths] = np.array(P[restPaths])*np.exp(-r*dt)

        u=np.mean(P*np.exp(-r*dt))

    return u

def BasisFunct(X, k):
    Ones=[1 for _ in range(len(X))]
    if k == 1:
        A = np.column_stack((Ones,1 - np.array(X)))
    elif k == 2:
        A = np.column_stack((Ones,1 - np.array(X),1/2*(2-4*np.array(X) + np.array(X)**2)))
    elif k == 3:
        A = np.column_stack((Ones,1 - np.array(X),1/2*(2-4*np.array(X) + np.array(X)**2), 1/6*(6-18*np.array(X) + 9*np.array(X)**2-np.array(X)**3)))
    elif k == 4:
        A = np.column_stack((Ones,1 - np.array(X),1/2*(2-4*np.array(X) + np.array(X)**2), 1/6*(6-18*np.array(X) + 9*np.array(X)**2-np.array(X)**3),1/24*(24 - 96*np.array(X) + 72*np.array(X)**2 - 16*np.array(X)**3 + np.array(X)**4)))
    elif k == 5:
        A = np.column_stack((Ones,1 - np.array(X),1/2*(2-4*np.array(X) + np.array(X)**2), 1/6*(6-18*np.array(X) + 9*np.array(X)**2-np.array(X)**3),1/24*(24 - 96*np.array(X) + 72*np.array(X)**2 - 16*np.array(X)**3 + np.array(X)**4),1/120*(120-600*np.array(X)+600*np.array(X)**2-200*np.array(X)**3+25*np.array(X)**4-np.array(X)**5)))
    else:
            sys.exit("Too many basis functions requested")
    return A

print(LSMPut(1, 0.06, 0.15, 100, 90, 20, 1000000, 5))


The purpose of the code is to calculate the price of an American put option. The time to execute if Paths &gt; 1,000,000 takes a long time, especially if I perform a sensitivity analysis. I would like to find out if there is any way for me to optimize the code or accelerate the processing time. Thank you.
I changed the main part of the script to:

from time import time
start = time()
val = LSMPut(1, 0.06, 0.15, 100, 90, 20, 40_000, 5)
print(f'Done in {time() - start:.3} seconds')
print(val)


The point is that Python will print how long the calculation took to
run. Before modifications, the script takes 3.05 seconds to run on my
computer.

Now I can modify lsmput (I renamed LSMPut and some other functions
and variables because they names shouldn't contain capital letters
according to the Python style guide) and see which changes makes it
run faster.

Here follows the changes that impacted the runtime the most. Changing

z=np.array([np.random.standard_normal() for _ in range(paths)])


to

z = np.random.standard_normal(paths)


brings the runtime down to 2.43 seconds. Changing

ones=[1 for _ in range(len(X))]


to

ones = np.ones(len(X))


reduces the runtime further to 2.21 seconds. Precalculating the powers
of X in basis_funct like this

X2 = X**2
X3 = X**3
X4 = X**4
X5 = X**5
...
A = np.column_stack((ones, 1 - X, 1/2 * (2 - 4*X + X2), ...


saves an additional 100 milliseconds. Changing

itmPaths = [index for index,value in enumerate(K - S) if value &gt; 0]


to

itmPaths = np.nonzero((K - S) &gt; 0)[0]


brings the runtime down to about 1.60 seconds. I think you see the
pattern now; anytime you loop over a numpy array there's a numpy
builtin function to do the job better and faster.

Then in the calculation of rest_paths:

rest_paths = np.setdiff1d(np.arange(0, paths - 1, 1), exPaths)


Is the paths - 1 intentional here? Looks like a bug to me. Assuming
it is a bug, updating P can be done more efficiently, like this:

mask = np.zeros(P.shape, dtype = bool)
...
mask.fill(False)
mask[itmPaths[Pt &gt; C]] = True

P[mask] = Pt[Pt &gt; C]
P[~mask] *= e_r_dt


Final code

The modified code runs in about 420 milliseconds on my computer. Note
that I have removed all branches except for the k == 5 one from the
basis_funct function. It should be straightforward for you to add it
back.

import sys
import numpy as np

# Uncomment this line during development to always get the same random
# numbers.
# np.random.seed(1234)

def lsmput(T, r, sigma, K, S0, TimeSteps, paths, k):
    dt = T/TimeSteps
    t = np.arange(0, T+dt, dt).tolist()

    z = np.random.standard_normal(paths)
    w = (r-sigma**2/2)*T + sigma*np.sqrt(T)*z

    S = S0*np.exp(w)
    P = np.maximum(K - S,0)

    e_r_dt = np.exp(-r * dt)
    mask = np.zeros(P.shape, dtype = bool)
    for i in range(TimeSteps-1, -1, -1):
        z = np.random.standard_normal(paths)
        w = t[i]*w/t[i+1] + sigma*np.sqrt(dt*t[i]/t[i+1])*z

        S = S0 * np.exp(w)
        itmPaths = np.nonzero(K &gt; S)[0]
        itmS = S[itmPaths]
        Pt = K - itmS

        itmDiscP = P[itmPaths] * e_r_dt

        A = basis_funct(itmS, k)
        beta = np.linalg.lstsq(A, itmDiscP)[0]
        C = np.dot(A, beta)

        mask.fill(False)
        mask[itmPaths[Pt &gt; C]] = True

        P[mask] = Pt[Pt &gt; C]
        P[~mask] *= e_r_dt
    return np.mean(P * e_r_dt)

def basis_funct(X, k):
    ones = np.ones(len(X))
    assert k == 5
    X2 = X**2
    X3 = X**3
    X4 = X**4
    X5 = X**5
    A = np.column_stack((ones, 1 - X,
                         1/2 * (2 - 4*X + X2),
                         1/6 * (6 - 18*X + 9*X2 - X3),
                         1/24 * (24 - 96*X + 72*X2 - 16*X3 + X4),
                         1/120 * (120 - 600*X + 600*X2 - 200*X3 + 25*X4 - X5)))
    return A

from time import time
start = time()
val = lsmput(1, 0.06, 0.15, 100, 90, 20, 40_000, 5)
print(f'Done in {time() - start:.3} seconds')
print(val)