# finitediff.py
# Author: Alex Gezerlis
# Numerical Methods in Physics with Python (2nd ed., CUP, 2023)

from math import exp, sin, cos, log10

def f(x):
    return exp(sin(2*x))

def fprime(x):
    return 2*exp(sin(2*x))*cos(2*x)

def calc_fd(f,x,h):
    fd = (f(x+h) - f(x))/h
    return fd

def calc_cd(f,x,h):
    cd = (f(x+h/2) - f(x-h/2))/h
    return cd

if __name__ == '__main__':
    x = 0.5
    an = fprime(x)

    hs = [10**(-i) for i in range(1,12)]
    fds = [abs(calc_fd(f,x,h) - an) for h in hs]
    cds = [abs(calc_cd(f,x,h) - an) for h in hs]

    rowf = "{0:1.0e} {1:1.16f} {2:1.16f}"
    print("h     abs. error in fd   abs. error in cd")
    for h,fd,cd in zip(hs,fds,cds):
        print(rowf.format(h,fd,cd))

# richardsondiff.py
# Author: Alex Gezerlis
# Numerical Methods in Physics with Python (2nd ed., CUP, 2023)

from finitediff import f, fprime, calc_fd, calc_cd

x = 0.5
an = fprime(x)

hs = [10**(-i) for i in range(1,7)]

rowf = "{0:1.0e} {1:1.16f} {2:1.16f}"
print("h     abs. err. rich fd  abs. err. rich cd")
for h in hs:
    fdrich = 2*calc_fd(f,x,h/2) - calc_fd(f,x,h)
    fd = abs(fdrich-an)
    cdrich = (4*calc_cd(f,x,h/2) - calc_cd(f,x,h))/3
    cd = abs(cdrich-an)
    print(rowf.format(h,fd,cd))

# psis.py
# Author: Alex Gezerlis
# Numerical Methods in Physics with Python (2nd ed., CUP, 2023)

from math import sqrt, pi, factorial, exp
import cmath

def hermite(n,x):
    val0 = 1.; val1 = 2*x
    for j in range(1,n):
        val2 = 2*x*val1 - 2*j*val0
        val0, val1 = val1, val2
    dval2 = 2*n*val0
    return val2, dval2

def psiqho(x,nametoval):
    n = nametoval["n"]
    momohbar = nametoval["momohbar"]
    al = nametoval["al"]
    psival = momohbar**0.25*exp(-0.5*al*momohbar * x**2) 
    psival *= hermite(n,sqrt(momohbar)*x)[0]
    psival /= sqrt(2**n * factorial(n) * sqrt(pi))
    return psival

def psibox(x,nametoval):
    n = nametoval["n"]
    boxl = nametoval["boxl"]
    return cmath.exp(2*pi*n*x*1j/boxl) / sqrt(boxl)

if __name__ == '__main__':
    x = 1.
    nametoval = {"n": 100, "momohbar": 1., "al": 1.}
    psiA = psiqho(x, nametoval)
    nametoval = {"n": -2, "boxl": 2*pi}
    psiB = psibox(x, nametoval)
    print(psiA, psiB)

# kinetic.py
# Author: Alex Gezerlis
# Numerical Methods in Physics with Python (2nd ed., CUP, 2023)

from psis import psiqho, psibox
from math import pi

def kinetic(psi,x,nametoval,h=0.005):
    hom = 1.
    psiold = psi(x,nametoval)
    psip = psi(x+h,nametoval)
    psim = psi(x-h,nametoval)

    lapl = (psip + psim - 2.*psiold)/h**2
    kin = -0.5*hom*lapl/psiold
    return kin

def test_kinetic():
    x = 1.
    hs = [10**(-i) for i in range(1,6)]

    nametoval = {"n": 100, "momohbar": 1., "al": 1.}
    qhos = [kinetic(psiqho,x,nametoval,h) for h in hs]
    nametoval = {"n": -2, "boxl": 2*pi}
    boxs = [kinetic(psibox,x,nametoval,h) for h in hs]

    rowf = "{0:1.0e} {1:1.16f} {2:1.16f}"
    print("h     qho                 box")
    for h,qho,box in zip(hs,qhos,boxs):
        print(rowf.format(h,qho,box))

if __name__ == '__main__':
    test_kinetic()