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

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 (2nd ed., CUP, 2023)

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 (2nd ed., CUP, 2023)

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 (2nd ed., CUP, 2023)

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 (2nd ed., CUP, 2023)

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))

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

from gauelim_pivot import gauelim_pivot
import numpy as np

def generatedata(N,a=0.,b=9,sts=(2,5,0.5,1)):
    sa, sb, sc, sd = sts
    np.random.seed(7921)
    data = np.zeros((3,N))
    data[0,:] = np.linspace(a,b,N)
    data[1,:] = sa + sb*np.sin(data[0,:])
    data[2,:] = sc + sd*np.random.random(N)
    data[1,:] += np.random.normal(0,data[2,:])
    return data

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(data,phi,n):
    N = data.shape[1]
    A = np.zeros((N,n))
    for k in range(n):
        A[:,k] = phi(n,k,data[0,:])/data[2,:]
    bs = data[1,:]/data[2,:]

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

if __name__ == '__main__':
    data = generatedata(8)
    for n in (5, 2):
        cs, chisq = normalfit(data, phi, n)
        print(cs)
        print(chisq/(data.shape[1]-cs.size))
        print("")

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

from gauelim_pivot import gauelim_pivot
from newnormal import generatedata, phi
import numpy as np

def inv(A):
    n = A.shape[0]
    invA = np.zeros((n,n))
    for i,bs in enumerate(np.identity(n).T):
        invA[:,i] = gauelim_pivot(A,bs)
    return invA

def bayes(data, batches, primus, priS):
    n = primus.size
    i = 0
    for N in batches:
        A = np.zeros((N,n))
        for k in range(n):
            A[:,k] = phi(n,k,data[0,i:i+N])/data[2,i:i+N]
        bs = data[1,i:i+N]/data[2,i:i+N]

        priSinv = inv(priS)
        postS = inv(A.T@A + priSinv)
        postmus = postS@(A.T@bs + priSinv@primus)
        
        primus, priS = postmus, postS
        i += N
        print(n, i, postmus, postS)
    return postmus, postS
    
if __name__ == '__main__':
    batches = 1, 4, 3
    data = generatedata(np.sum(batches))
    for n in (5, 2):
        primus = np.zeros(n)
        priS = np.zeros((n,n))
        np.fill_diagonal(priS, np.linspace(10,2,n))
        postmus, postS = bayes(data, batches, primus, priS)
        print(postmus)

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

from jacobi import termcrit
from newnormal import generatedata
from numpy import exp
import numpy as np

def soft(der, x):
    return np.log(1+exp(x)) if der==0 else exp(x)/(1+exp(x))

def getzspofx(W, x, act):
    zs = act(0, W[1,:]*x + W[2,:])
    pofx = act(0, W[0,:]@zs + W[3,0])
    return zs, pofx

def backprop(W, x, y, act):
    zs, pofx = getzspofx(W, x, act)
    ders = np.zeros((4,zs.size))
    ders[3,0] = 2*(pofx - y)*act(1, W[0,:]@zs + W[3,0])
    ders[0,:] = ders[3,0]*zs
    ders[1,:] = ders[3,0]*W[0,:]*x*act(1,W[1,:]*x+W[2,:])
    ders[2,:] = ders[3,0]*W[0,:]*act(1,W[1,:]*x+W[2,:])
    return ders

def sgd(W, x, y, act=soft, g=0.006, kmax=300, tol=1.e-4):
    for k in range(1,kmax):
        ders = backprop(W, x, y, act)
        Wn = np.array([W[i,:]-g*ders[i,:] for i in range(4)])
        _, pofxnew = getzspofx(Wn, x, act)
        conds = [termcrit(W[i,:],Wn[i,:]) for i in range(3)]
        if np.all([val<tol for val in conds[:3]]):
            break
        W = np.copy(Wn)
    return Wn, (pofxnew - y)**2

if __name__ == '__main__':
    sts = (1.5,1.2,0.15,0.08)
    data = generatedata(400,0.2,4.2,sts)
    np.random.seed(5179)
    W = np.random.uniform(-1,1,(4,10))
    W, s = sgd(W, data[0,155], data[1,155]); print(s)

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

from gauelim_pivot import gauelim_pivot
from jacobi import termcrit
import numpy as np

def generatedata():
    data = np.zeros((3,13))
    data[0,:] = np.array([373.1, 492.5, 733, 755, 799, 820,
            877, 1106, 1125, 1403, 1492, 1522, 1561])
    data[1,:] = np.array([156., 638, 3320, 3810, 4440, 5150,
            6910, 16400, 17700, 44700, 57400, 60600, 67800])
    data[2,:] = np.ones(data.shape[1])
    return data

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

def getKrs(data, cs):
    K = np.zeros((data.shape[1], cs.size))
    K[:,0] = 1/data[2,:]
    K[:,1] = data[0,:]**cs[2]/data[2,:]
    K[:,2] = cs[1]*data[0,:]**cs[2]*np.log(data[0,:])/data[2,:]
    rs = (data[1,:] - model(cs, data[0,:]))/data[2,:]
    return K, rs

def gaussnewton(data, colds, kmax=50, tol=1.e-8):
    for k in range(1,kmax):
        K, rs = getKrs(data, colds)
        cnews = colds + gauelim_pivot(K.T@K, K.T@rs)
        err = termcrit(colds,cnews)
        print(k,err)
        if err < tol:
            break
        colds = np.copy(cnews)
    return cnews

if __name__ == '__main__':
    data = generatedata()
    colds = np.array([-700, 1.26e-8, 4.5])
    cs = gaussnewton(data,colds); print(cs)