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

import numpy as np

def f(x):
    return 1/np.sqrt(x**2 + 1)

def rectangle(f,a,b,n):
    h = (b-a)/(n-1)
    xs = a + np.arange(n-1)*h
    fs = f(xs)
    return h*np.sum(fs)

def trapezoid(f,a,b,n):
    h = (b-a)/(n-1)
    xs = a + np.arange(n)*h
    cs = np.ones(n); cs[0] = 0.5; cs[-1] = 0.5
    contribs = cs*f(xs)
    return h*np.sum(contribs)

def simpson(f,a,b,n):
    h = (b-a)/(n-1)
    xs = a + np.arange(n)*h
    cs = 2*np.ones(n)
    cs[1::2] = 4; cs[0] = 1; cs[-1] = 1
    contribs = cs*f(xs)
    return (h/3)*np.sum(contribs)

if __name__ == '__main__':
    ans = np.log(1 + np.sqrt(2))
    print(ans)

    for integrator in (rectangle, trapezoid, simpson):
        print(integrator(f, 0., 1., 51), end=" ")

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

from newtoncotes import f, rectangle, trapezoid, simpson
import numpy as np

def adaptive(f,a,b,integrator,kmax = 20,tol = 1.e-12):
    functodenom = {rectangle:1, trapezoid:3, simpson:15}
    denom = functodenom[integrator]

    n = 2
    val = integrator(f,a,b,n) 
    for k in range(kmax):
        nprime = 2*n-1
        valprime = integrator(f,a,b,nprime)
        err = abs(valprime-val)/denom
        err /= abs(valprime)
        print(nprime, valprime, err)
        if err<tol:
            break

        n, val = nprime, valprime
    else:
        valprime = None
    return valprime

if __name__ == '__main__':
    ans = np.log(1 + np.sqrt(2))
    print(ans); print("")
    for integrator in (rectangle, trapezoid, simpson):
        print(adaptive(f, 0., 1., integrator)); print("")

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

from newtoncotes import f, trapezoid
import numpy as np

def prettyprint(row):
    for elem in row:
        print("{0:1.11f} ".format(elem),end="")
    print("")

def richardson(Rprev, Rincurr0, i):
    Rcurr = [Rincurr0]
    for j in range(1, i+1):
        val = (4**j*Rcurr[j-1] - Rprev[j-1])/(4**j - 1)
        Rcurr.append(val)
    return Rcurr

def romberg(f,a,b,imax = 20,tol = 1.e-8):
    n = 2
    val = trapezoid(f,a,b,n) 
    Rprev = [val]
    #prettyprint(Rprev)
    for i in range(1,imax):
        nprime = 2*n-1
        Rincurr0 = trapezoid(f,a,b,nprime)
        Rcurr = richardson(Rprev, Rincurr0, i)
        #prettyprint(Rcurr)
        err = abs(Rprev[-1] - Rcurr[-1])/abs(Rcurr[-1])
        valprime = Rcurr[-1]
        if err < tol:
            break
        n = nprime
        Rprev = Rcurr[:]
    else:
        valprime = None
    return valprime

if __name__ == '__main__':
    ans = np.log(1 + np.sqrt(2))
    print(ans)
    print(romberg(f,0.,1.))

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

from legendre import legendre
from legroots import legroots
from newtoncotes import f
import numpy as np

def gauleg_params(n):
    xs = legroots(n)
    cs = 2/((1-xs**2)*legendre(n,xs)[1]**2)
    return xs, cs

def gauleg(f,a,b,n):
    xs, cs = gauleg_params(n)
    coeffp = 0.5*(b+a) 
    coeffm = 0.5*(b-a)
    ts = coeffp + coeffm*xs
    contribs = cs*f(ts)
    return coeffm*np.sum(contribs)

if __name__ == '__main__':
    ans = np.log(1 + np.sqrt(2))
    print(ans)
    for n in range(2,10):
        print(n, gauleg(f,0.,1,n))

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

from newtoncotes import f
import numpy as np

def montecarlo(f,a,b,n,option="uniform"):
    np.random.seed(314159)
    us = np.random.uniform(a, b, n)

    if option=="uniform":
        fs = f(us)
    else:
        c0 = 4 - 2*np.sqrt(2)
        c1 = -6 + 4*np.sqrt(2)
        xs = (-c0 + np.sqrt(2*c1*us + c0**2))/c1
        fs = f(xs)/(c0 + c1*xs)

    fbar, err = stats(fs)
    return (b-a)*fbar, (b-a)*err

def stats(fs):
    n = fs.size
    fbar = np.sum(fs)/n
    fsq = np.sum(fs**2)/n
    varfbar = (fsq - fbar**2)/(n - 1)
    return fbar, np.sqrt(varfbar)

if __name__ == '__main__':
    for n in 10**np.arange(2,7):
        avu, erru = montecarlo(f, 0., 1., n)
        avi, erri = montecarlo(f, 0., 1., n, option="is")
        rowf = "{0:7d}   {1:1.9f} {2:1.9f}   {3:1.9f} {4:1.9f}"
        print(rowf.format(n, avu, erru, avi, erri))

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

import numpy as np

def psi(al, oms, X):
    rexp = np.sum(oms*X**2)
    return np.exp(-0.5*al*rexp)

def ekin(al, oms, X, h=0.01, hom = 1.):
    npart, ndim = X.shape
    psiold = psi(al, oms, X)
    kin = 0.
    for (j,el), r in np.ndenumerate(X):
        X[j,el] = r + h
        psip = psi(al, oms, X)
        X[j,el] = r - h
        psim = psi(al, oms, X)
        X[j,el] = r
        lapl = (psip + psim - 2.*psiold)/h**2
        kin += -0.5*hom*lapl/psiold
    return kin

def epot(oms, X, strength = 3, m = 1., q = 1.):
    npart, ndim = X.shape
    pot = 0.5*m*np.sum(oms**2*X**2)
    for k in range(1,npart):
        for j in range(k):
            r2 = np.sum((X[j,:] - X[k,:])**2)
            pot += strength*np.exp(-q*r2)
    return pot

if __name__ == '__main__':
    npart, ndim, al = 4, 3, 0.6
    oms = np.arange(1, 1 + ndim)
    X = 1/np.arange(1, 1 + npart*ndim).reshape(npart, ndim)
    print(ekin(al, oms, X), epot(oms, X))

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

from eloc import psi, ekin, epot
from montecarlo import stats
import numpy as np

def vmc(npart, ndim, al, oms, inseed=8735):
    ncurly, nm, th = 10**4, 100, 0.8
    np.random.seed(inseed)
    Xold = np.random.uniform(-1, 1, (npart, ndim))
    psiold = psi(al, oms, Xold)
    iacc, imeas = 0, 0
    eners = np.zeros(ncurly)

    for itot in range(nm*ncurly):
        Xnew = Xold + th*np.random.uniform(-1,1,(npart, ndim))
        psinew = psi(al, oms, Xnew)
        psiratio = (psinew/psiold)**2

        if np.random.uniform(0,1) <= psiratio:
            Xold = np.copy(Xnew)
            psiold = psinew
            iacc +=1
        if (itot%nm)==0:
            eners[imeas] = ekin(al,oms,Xold) + epot(oms,Xold)
            imeas += 1

    return iacc/(nm*ncurly), eners

if __name__ == '__main__':
    npart, ndim, al = 4, 3, 0.6
    oms = np.arange(1, 1 + ndim)
    accrate, eners = vmc(npart, ndim, al, oms)
    av, err = stats(eners)
    print(accrate, av, err)