#  Python 3.7
#  For the course Stochastic Calculus, Fall semester 2018, Courant Institute, NYU
#  Author: course instructor, Jonathan Goodman
#  https://www.math.nyu.edu/faculty/goodman/teaching/StochCalc2018/StochCalc.html

#  filename: PathDemo.py
#  See Assignment 7 (on class web site above) for details and notation.

#  Generate M paths of a one dimensional diffusion process and visualize
#  the results using histograms and plots of sample paths.

#  This demo is for geometric Brownian motion

#      dX = mu*X*dt + sig*X*dW

#  The code uses Python vector instructions to make it run faster

import numpy             as np      # numerical computing library
import matplotlib.pyplot as plt     # library of plot routines
import numpy.random      as ran     # random number generators

print("Python file PathDemo.py")
print("How to make and display a diffusion paths")

pathsPlotFile     = "paths.pdf"
histogramPlotFile = "hist.pdf"


M  = 5000   # number of random paths
Mp = 40     # number of paths to plot as lines
T  = .5     # The length of the simulation
L  = 41     # number of bins for histograms
n  = 4000   # number of time steps per path
dt = T/n    # the time step, Delta t

mu  = .2     # growth rate
sig = .5     # volatility
X0  = 1.     # starting price


#              Make the paths

paths = np.ndarray([n+1,M]) # hold all the paths at all times, starting at t=0
X     = np.ndarray(M)       # the locations of the paths at time t

X[0:M]= X0                  # all the paths start at X0
                            # [0:M] means 0, 1, ..., M-1
                            # start at 0, end just before M

t = 0.
paths[0,0:M] = X            # copy in the initial condition
                            # This makes paths[0,k]=X[k] for k = 0, ..., M-1
ran.seed(12345)             # set a seed to the random numbers are the same each time
for k in range(n):
    Z = ran.randn(M)        # M standard normals, one for each path
    
    X[0:M] += sig*X[0:M]*np.sqrt(dt)*Z[0:M]     # Add the noise part of the increment
                                                # this is a vector verson of the formula
                                                #  X[k] = X[k] + sig*X[k]*sqrt(dt)*Z[k]
    X[0:M] += mu *X[0:M]*dt                     # Add the drift part of the increment

    t += dt                                     # advance time by one time step
    paths[ (k+1), 0:M] = X                      # copy the present position into
                                                # the path array

#                Plot some paths

times = np.linspace( start = 0., stop = T, num=n+1)
for path in range(Mp):
    plt.plot( times, paths[0:(n+1), path], linewidth = .2)
plt.grid()
plt.xlabel("t")
plt.ylabel("X")
title = "Geometric Brownian motion with sig=" + "{:3.1f}".format(sig)
title = title +                      ", mu="  + "{:3.1f}".format(mu)
title = title +                      ", dt="  + "{:7.2e}".format(dt)
plt.title(title)
plt.savefig(pathsPlotFile)
plt.show()
plt.close()


#              Plot a histogram of the locations at the final time

a = np.min(X)     # make bins go from the smallest to largest actual value
b = np.max(X)

dx = (b-a)/L      # dx = bin size = (length being binned)/(number of bins)
xc = np.linspace( start = a+.5*dx, stop = b-.5*dx, num = L)   # bin centers

N  = np.zeros( shape = [L], dtype =  int )  # bin counts, initialized to zero

for Xp in X:                          # look at each X in the X array
    k = int(np.floor( ( Xp-a )/dx ) ) # the bin X goes into
    if ( ( k >= 0 ) and ( k < L ) ):  # check that k is in the bin range
        N[k] += 1                     # increment the count for bin k
uh = N/(M*dx)                         # uh = u hat = estimate of PDF, M = # of points
plt.plot( xc, uh)

title = "GBM density sig=" + "{:3.1f}".format(sig)
title = title +                      ", mu="  + "{:3.1f}".format(mu)
title = title +                      ", T="  + "{:3.1f}".format(T)
title = title +                      ", dt="  + "{:7.2e}".format(dt)
plt.title(title)
plt.xlabel("x")
plt.ylabel("u")
plt.grid()
plt.savefig(histogramPlotFile)
plt.show()