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

#  Use the forward Euler finite difference method to solve the backward
#  equation for Brownian motion with absorbing side boundary conditions.
#  Experiment with moving forward and backward in time to demonstrate
#  ill posedness.
#  See Week 3 for more information

#  The function backwardStep could have been done in a way that uses less
#  computer time and memory.  It could have used a solver for a tridiagonal
#  symmetric matrix.  It could have done the factorization just once rather
#  than at each time step.  It is done this way to make it easy to follow.

#  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
    
#-------------------- reverse a forward Euler step     --------------------

def backwardStep( Fk, dt, dx):
    """Take one backward Euler time step for the heat equation.
       Use linear algebra to undo a forward Euler step
       \partial_t F + (1./2) \partial_x^2 F = 0
       Warning: the backward solver assumes F[0] = F[n+1] = 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_kp1 = 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
    A = np.zeros([n,n])     #  A = time step matrix, Fkm1 = A*Fk,
    for j in range(n):
        A[j,j] = b          #  diagonal of A = prob(walk does not move)
    for j in range(n-1):
        A[j,j+1] = a        #  prob(walk move up) = prob(down)
        A[j+1,j] = c
        
    Ai = np.linalg.inv(A)   #  Ai = "A inverse"
    
    Fkp1      = np.zeros([n+2])
    Fkp1[0]   = Fk[0]
    Fkp1[n+1] = Fk[n+1]
    
    Fkp1[1:(n+1)] = Ai@Fk[1:(n+1)]   #  The "@" in new Python is matrix/vector mult
    
    return Fkp1
    
        
    
        
#------------------ The main program ---------------------------------------------------

#    Collect all constants and parameters in one place

n   = 8              # number of interior points, there are two boundary points
L   = 3.             # length of the space interval
T   = 1.2            # integrate from time T back to time 0
lam = .8             # CFL ratio = dt/(dx^2)

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


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.
    
#   evolve in the "right" time direction (backwards) to get the 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

#    evolve in the "wrong" time direction (forwards) to recover the final condition

Fk = F0                            #  start with the value function at time 0
for k in range(nt):                #  evolve toward the given final condition
    Fk = backwardStep( Fk, dt, dx)
FT = Fk                           #  FT = F at time T, should be V, but isn't

#   Plotting, mostly copy/pasted from Week 1 and Week 2 code

fig, ax = plt.subplots()                       # Create a figure containing a single axes.

ax.plot( x,  V,      label="Final condtion")
ax.plot( x, F0,      label="Value function")
ax.plot( x, FT, "o", label="Reconstructed final condition")

title = "V to F0 back to V, n = {n:3d}, T ={T:7.2f}, lam = {lam:6.2f}"
title = title.format(n = n, T=T, lam = lam)
ax.set_title(title)
ax.legend()
ax.grid

eMax = np.max(FT-V)                 #  This should be zero, but isn't
text = "max error is {eMax:9.2e}"   #  more information about the results ...
text = text.format(eMax = eMax)     #  ... in a text box in the plot
ax.text( .8, .22, text)
text = "{nt:3d} time steps"
text = text.format(nt = nt)
ax.text( .8, .28, text)

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