#  Code for Stochastic Calculus class, Jonathan Goodman
#  http://www.math.nyu.edu/faculty/goodman/teaching/StochCalc20/
#  FiniteDifference.py

#  Use the forward Euler finite difference method to solve the backward
#  equation for Brownian motion with absorbing side boundary conditions.
#  Demonstrate convergence as the number of points and time steps increases.
#  See Week 3 for more information

#  The author gives permission for anyone to use this publicly posted 
#  code for any purpose.  The code was written for teaching, not research
#  or commercial use.  It has not been tested thoroughly and probably has
#  serious bugs.  Results may be inaccurate, incorrect, or just wrong.

# illustrate using the Python random number generator from numpy version 19

import numpy             as np
import matplotlib.pyplot as plt

#-------------------- The normal forward step with forward Euler  --------------------

def forwardStep( Fk, dt, dx):
    """Take one forward Euler time step for the heat equation.
       \partial_t F + (1./2) \partial_x^2 F = 0
       Fk   = solution at time t_k, a numpy array with Fk[0] and Fk[n+1] boundary values
       dx   = space step size
       dt   = time step size.  Must have dt <= dx^2
       return values:
       F_km1 = solution at time t_{k-1}= t_k-dt
    """
    
    a = .5*dt/(dx**2)      # random walk probabilities ...
    b = 1. - dt/(dx**2)    # ... = finite difference weights.
    c = a                  # Week 3 notes derive these formulas
    if ( b < 0 ):
        print("violated the time step constraint")
        return
        
    Fkm1      = np.zeros([n+2])    # n internal values plus two boundary values
    Fkm1[0]   = Fk[0]              # copy boundary conditions
    Fkm1[n+1] = Fk[n+1]
    for j in range(1,(n+1)):
        Fkm1[j] = a*Fk[j-1] + b*Fk[j] + c*Fk[j+1]
    return Fkm1
        
    
        
#------------------ The main program ---------------------------------------------------

#    Collect all constants and parameters in one place

#  number of grid points and plot color or symbol for each run

run1    = { 'n':  8, 'symbol': "o"}
run2    = { 'n': 14, 'symbol': "g"}
run3    = { 'n': 29, 'symbol': "r"}
runList = [ run1, run2, run3]

L     = 12.               # length of the space interval
T     = 1.25               # integrate from time T back to time 0
lam   = .5                # CFL ratio = dt/(dx^2)

PlotFile = "FiniteDifference.pdf"   # filename for plot file

#  Do a separate solve for each n in the list

fig, ax = plt.subplots()    # Put the results into a single plot.
title = "Value function for payout at time {T:7.2f} with lambda = {lam:7.2f}"
title = title.format(T=T, lam = lam)
ax.set_title(title)

for run in runList:       # do all these runs

    n      = run["n"]           # the specs for this run
    symbol = run["symbol"]

    dx = L/(n+1)          # n+1 = number of intervals for n internal points
    dt = lam*dx*dx        # CFL formula for the time step, see Week 3 notes
    nt = int(T/dt) + 1    # number of time steps between T and 0, rounded up
    dt = T/nt             # possibly adjust dt down to have nt equal size time steps

    x  = np.linspace( 0., L, n+2)   # x values, for plotting only.

#   Set up the final condition

    V = np.zeros([n+2])      #  zero boundary values
    for j in range(1,(n+1)):
        V[j] = 1.            #  all ones in the interior.
    
#   time step loop to compute the time zero value function
    
    Fk = V                             #  Fk = value function at time t_k = k*dt
    for ntmk in range(nt):             #  ntmk = nt - k = "nt minus k"
        Fk = forwardStep( Fk, dt, dx)
    F0 = Fk                            #  value funcion at t=0, save for plotting

#   Add this computation to the figure
    legendText = "n = {n:4d}"
    legendText = legendText.format( n = n)
    ax.plot( x, F0, symbol, label= legendText, )
    #ax.plot( x, F0, label= legendText, symbol)

ax.legend()
ax.grid()
ax.set_ylabel('F')
ax.set_xlabel('x')


plt.savefig(PlotFile)       # Save a good copy of the plot
plt.show()