# barycentric.py
# Author: Alex Gezerlis
# Numerical Methods in Physics with Python (CUP, 2020)

import numpy as np

def f(x):
    return 1/(1 + 25*x**2)

def generatedata(n,f,nodes="cheb"):
    if nodes=="cheb":
        dataxs = -np.cos(np.linspace(0,np.pi,n))
    else:
        dataxs = np.linspace(-1,1,n)
    datays = f(dataxs)
    return dataxs, datays

def weights(dataxs):
    n = dataxs.size
    ws = np.ones(n)
    for k in range(n):
        for j in range(n):
            if j == k:
                continue
            ws[k] *= (dataxs[k]-dataxs[j])
    return 1/ws

def bary(dataxs,datays,ws,x):
    k = np.where(x == dataxs)[0]
    if k.size == 0:
        nume = np.sum(ws*datays/(x-dataxs))
        denom = np.sum(ws/(x-dataxs))
        val = nume/denom
    else:
        val = datays[k[0]]
    return val

if __name__ == '__main__':
    dataxs, datays = generatedata(15, f)
    ws = weights(dataxs)
    x = 0.3; pofx = bary(dataxs, datays, ws, x)
    print(x, pofx, f(x))

# splines.py
# Author: Alex Gezerlis
# Numerical Methods in Physics with Python (CUP, 2020)

from barycentric import f, generatedata
from gauelim_pivot import gauelim_pivot
import numpy as np

def computecs(dataxs,datays):
    n = dataxs.size
    A = np.zeros((n-2,n-2))
    np.fill_diagonal(A, 2*(dataxs[2:]-dataxs[:-2]))
    np.fill_diagonal(A[1:,:], dataxs[2:-1]-dataxs[1:-2])
    np.fill_diagonal(A[:,1:], dataxs[2:-1]-dataxs[1:-2])

    b1 = (datays[2:]-datays[1:-1])/(dataxs[2:]-dataxs[1:-1])
    b2 = (datays[1:-1]-datays[:-2])/(dataxs[1:-1]-dataxs[:-2])
    bs = 6*(b1 - b2)

    cs = np.zeros(n)
    cs[1:-1] = gauelim_pivot(A, bs)
    return cs

def splineinterp(dataxs,datays,cs,x):
    k = np.argmax(dataxs>x)
    xk = dataxs[k]; xk1 = dataxs[k-1]
    yk = datays[k]; yk1 = datays[k-1]
    ck = cs[k]; ck1 = cs[k-1]

    val = yk1*(xk-x)/(xk-xk1) + yk*(x-xk1)/(xk-xk1)
    val -= ck1*((xk-x)*(xk-xk1) - (xk-x)**3/(xk-xk1))/6
    val -= ck*((x-xk1)*(xk-xk1) - (x-xk1)**3/(xk-xk1))/6
    return val

if __name__ == '__main__':
    dataxs, datays = generatedata(15, f, "equi")
    cs = computecs(dataxs, datays)
    x = 0.95; pofx = splineinterp(dataxs, datays, cs, x) 
    print(x, pofx, f(x))

# triginterp.py
# Author: Alex Gezerlis
# Numerical Methods in Physics with Python (CUP, 2020)

from math import pi
import numpy as np

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

def generatedata(n,f):
    dataxs = 2*pi*np.arange(n)/n
    datays = f(dataxs)
    return dataxs, datays

def computeparams(dataxs,datays):
    n = dataxs.size
    m = n//2
    aparams = np.zeros(m+1)
    bparams = np.zeros(m-1)
    
    for k in range(m+1):
        aparams[k] = datays@np.cos(k*dataxs)/m
    for k in range(1,m):
        bparams[k-1] = datays@np.sin(k*dataxs)/m
    return aparams, bparams

def triginterp(aparams,bparams,x):
    n = aparams.size + bparams.size
    m = n//2
    val = 0.5*(aparams[0] + aparams[-1]*np.cos(m*x))
    for k in range(1,m):
        val += aparams[k]*np.cos(k*x)
        val += bparams[k-1]*np.sin(k*x)
    return val

if __name__ == '__main__':
    dataxs, datays = generatedata(8, f)
    aparams, bparams = computeparams(dataxs, datays)
    x = 0.3; pofx = triginterp(aparams, bparams, x)
    print(x,pofx)

# fft.py
# Author: Alex Gezerlis
# Numerical Methods in Physics with Python (CUP, 2020)

from triginterp import f, generatedata
from math import pi
import numpy as np

def fft(ys):
    n = ys.size
    m = n//2
    if n==1:
        ytils = ys
    else:
        evens = fft(ys[::2])
        odds = fft(ys[1::2])
        coeffs = np.exp(-2*pi*np.arange(m)*1j/n)
        first = evens + coeffs*odds
        second = evens - coeffs*odds
        ytils = np.concatenate((first, second))
    return ytils

def fftinterp(ytils,x):
    n = ytils.size
    m = n//2
    val = ytils[:m]@np.exp(np.arange(m)*x*1j)
    val += ytils[m]*np.cos(m*x)
    val += ytils[m+1:]@np.exp(np.arange(-m+1,0)*x*1j)
    return val/n

if __name__ == '__main__':
    n = 8
    dataxs, datays = generatedata(n, f)
    ytils = fft(datays)
    x = 0.3; pofx = fftinterp(ytils, x)
    print(x,pofx.real)

# linefit.py
# Author: Alex Gezerlis
# Numerical Methods in Physics with Python (CUP, 2020)

import numpy as np

def helpers(dataxs,datays,datasigs):
    S = np.sum(1/datasigs**2)
    Sx = np.sum(dataxs/datasigs**2)
    Sy = np.sum(datays/datasigs**2)
    Sxx = np.sum(dataxs**2/datasigs**2)
    Sxy = np.sum(dataxs*datays/datasigs**2)
    Del = S*Sxx - Sx**2
    return S, Sx, Sy, Sxx, Sxy, Del

def computecs(dataxs,datays,datasigs):
    S,Sx,Sy,Sxx,Sxy,Del = helpers(dataxs,datays,datasigs)
    cs = np.zeros(2); dcs = np.zeros(2)
    cs[0] = (Sxx*Sy - Sx*Sxy)/Del
    cs[1] = (S*Sxy - Sx*Sy)/Del
    dcs[0] = np.sqrt(Sxx/Del)
    dcs[1] = np.sqrt(S/Del)
    return cs, dcs

def computechisq(dataxs,datays,datasigs,cs):
    chisq = np.sum((datays-cs[0]-cs[1]*dataxs)**2/datasigs**2)
    return chisq

dataxs = np.linspace(0,1,6)
datays = np.array([3.085, 3.123, 3.224, 3.360, 3.438, 3.569])
datasigs = np.array([0.048, 0.053, 0.02, 0.005, 0.023, 0.07])

cs, dcs = computecs(dataxs, datays, datasigs)
print(cs); print(dcs)
chisq = computechisq(dataxs, datays, datasigs, cs)
print(chisq/(dataxs.size - cs.size))

# normalfit.py
# Author: Alex Gezerlis
# Numerical Methods in Physics with Python (CUP, 2020)

from gauelim_pivot import gauelim_pivot
import numpy as np

def generatedata(N):
    np.random.seed(45379)
    dataxs = np.linspace(0,9,N)
    datays = 2 + 5*np.sin(dataxs) + 0.3*np.random.randn(N)
    datasigs = 0.2*np.abs(np.random.randn(N))
    return dataxs, datays, datasigs

def phi(n,k,x):
    if n==5:
        val = x**k
    elif n==2:
        val = 1. if k==0 else np.sin(x)
    return val

def normalfit(dataxs,datays,datasigs,n):
    N = dataxs.size
    A = np.zeros((N,n))
    for k in range(n):
        A[:,k] = phi(n,k,dataxs)/datasigs
    bs = datays/datasigs

    cs = gauelim_pivot(A.T@A, A.T@bs)
    chisq = np.sum((bs - A@cs)**2)
    return cs, chisq

if __name__ == '__main__':
    dataxs, datays, datasigs = generatedata(8)
    for n in (5, 2):
        cs, chisq = normalfit(dataxs, datays, datasigs, n)
        print(cs)
        print(chisq/(dataxs.size-cs.size)); print("")

# blackbody.py
# Author: Alex Gezerlis
# Numerical Methods in Physics with Python (CUP, 2020)

from multi_newton import jacobian, multi_newton
import numpy as np

def generatedata():
    dataxs = np.array([373.1, 492.5, 733, 755, 799, 820,
            877, 1106, 1125, 1403, 1492, 1522, 1561])
    datays = np.array([156., 638, 3320, 3810, 4440, 5150,
            6910, 16400, 17700, 44700, 57400, 60600, 67800])
    datasigs = np.ones(dataxs.size)
    return dataxs, datays, datasigs

def model(cs,x):
    return cs[0] + cs[1]*x**cs[2]

def fs(cs):
    dataxs, datays, datasigs = generatedata()
    c0, c1, c2 = cs
    resids = datays - model(cs, dataxs)
    f0 = np.sum(resids/datasigs**2)
    f1 = np.sum(dataxs**c2*resids/datasigs**2)
    numers = c1*dataxs**c2*np.log(dataxs)*resids
    f2 = np.sum(numers/datasigs**2)
    return np.array([f0,f1,f2])

def computechisq(cs):
    dataxs, datays, datasigs = generatedata()
    chisq = np.sum((datays - model(cs,dataxs))**2/datasigs**2)
    return chisq

if __name__ == '__main__':
    colds = np.array([-700, 1.26e-8, 6])
    cs = multi_newton(fs,jacobian,colds,kmax=500,tol=2.e-6)
    chisq = computechisq(cs); print(cs); print(chisq)