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

import numpy as np

def f(x,y):
    return - (30/(1-x**2)) + ((2*x)/(1-x**2))*y - y**2

def euler(f,a,b,n,yinit):
    h = (b-a)/(n-1)
    xs = a + np.arange(n)*h
    ys = np.zeros(n)

    y = yinit
    for j,x in enumerate(xs):
        ys[j] = y
        y += h*f(x, y)
    return xs, ys
        
def rk4(f,a,b,n,yinit):
    h = (b-a)/(n-1)
    xs = a + np.arange(n)*h
    ys = np.zeros(n)

    y = yinit
    for j,x in enumerate(xs):
        ys[j] = y
        k0 = h*f(x, y)
        k1 = h*f(x+h/2, y+k0/2)
        k2 = h*f(x+h/2, y+k1/2)
        k3 = h*f(x+h, y+k2)
        y += (k0 + 2*k1 + 2*k2 + k3)/6
    return xs, ys
        
if __name__ == '__main__':
    a, b, n, yinit = 0.05, 0.49, 12, 19.53
    xs, ys = euler(f,a,b,n,yinit); print(ys)
    xs, ys = rk4(f,a,b,n,yinit); print(ys)

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

import numpy as np

def fs(x,yvals):
    y0, y1 = yvals
    f0 = y1
    f1 = - (30/(1-x**2))*y0 + ((2*x)/(1-x**2))*y1 
    return np.array([f0, f1])

def rk4_gen(fs,a,b,n,yinits):
    h = (b-a)/(n-1)
    xs = a + np.arange(n)*h
    ys = np.zeros((n, yinits.size))

    yvals = np.copy(yinits)
    for j,x in enumerate(xs):
        ys[j,:] = yvals
        k0 = h*fs(x, yvals)
        k1 = h*fs(x+h/2, yvals+k0/2)
        k2 = h*fs(x+h/2, yvals+k1/2)
        k3 = h*fs(x+h, yvals+k2)
        yvals += (k0 + 2*k1 + 2*k2 + k3)/6
    return xs, ys
        
if __name__ == '__main__':
    a, b, n = 0.05, 0.49, 12
    yinits = np.array([0.0926587109375, 1.80962109375])
    xs, ys = rk4_gen(fs,a,b,n,yinits)
    print(ys)

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

from secant import secant
from ivp_two import fs, rk4_gen
import numpy as np

def shoot(sig):
    a, b, n = 0.05, 0.49, 100
    yinits = np.array([0.0926587109375, sig])
    xs, ws = rk4_gen(fs,a,b,n,yinits)  
    wfinal = 0.11177050858750004
    return ws[-1, 0] - wfinal

if __name__ == '__main__':
    wder = secant(shoot,0.,1.) 
    print(wder)

# bvp_matrix.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 matsetup(a,b,n):
    h = (b-a)/(n-1)
    xs = a + np.arange(n)*h

    A = np.zeros((n,n))
    np.fill_diagonal(A, -2 + 30*h**2/(1-xs**2))
    A[0,0] = 1; A[-1,-1] = 1
    np.fill_diagonal(A[1:,:], 1 + h*xs[1:]/(1-xs[1:]**2))   
    A[-1,-2] = 0
    np.fill_diagonal(A[:,1:], 1 - h*xs/(1-xs**2))
    A[0,1] = 0

    bs = np.zeros(n)
    bs[0] = 0.0926587109375
    bs[-1] = 0.11177050858750004
    return A, bs

def riccati(a,b,n):
    A, bs = matsetup(a, b, n)
    ws = gauelim_pivot(A, bs)
    return ws

if __name__ == '__main__':
    a, b, n = 0.05, 0.49, 400
    ws = riccati(a, b, n); print(ws)

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

from secant import secant
import numpy as np

def fs(x,yvals,s):
    q = 1.5
    y0, y1 = yvals
    f0 = y1
    f1 = (2*q*np.cos(2*x) - s)*y0
    return np.array([f0, f1])

def rk4_gen_eig(fs,a,b,n,yinits,s):
    h = (b-a)/(n-1)
    xs = a + np.arange(n)*h
    ys = np.zeros((n, yinits.size))

    yvals = np.copy(yinits)
    for j,x in enumerate(xs):
        ys[j,:] = yvals
        k0 = h*fs(x, yvals,s)
        k1 = h*fs(x+h/2, yvals+k0/2,s)
        k2 = h*fs(x+h/2, yvals+k1/2,s)
        k3 = h*fs(x+h, yvals+k2,s)
        yvals += (k0 + 2*k1 + 2*k2 + k3)/6
    return xs, ys

def shoot(s):
    a, b, n = 0, 2*np.pi, 500
    yinits = np.array([0., 5.])

    xs, ys = rk4_gen_eig(fs,a,b,n,yinits,s)  
    wfinal = 0.
    return ys[-1, 0] - wfinal

if __name__ == '__main__':
    for sinit in (-0.4, 3.3, 8.5):
        sval = secant(shoot,sinit,sinit+0.5) 
        print(sval, end=" ")

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

from qrmet import qrmet
import numpy as np

def matsetup(q,n):
    h = 2*np.pi/n
    xs = np.arange(n)*h

    A = np.zeros((n,n))
    np.fill_diagonal(A, -2 - 2*h**2*q*np.cos(2*xs))
    np.fill_diagonal(A[1:,:], 1)   
    np.fill_diagonal(A[:,1:], 1)
    A[0,-1] = 1
    A[-1,0] = 1
    return A

def mathieu(q,n):
    A = matsetup(q, n)
    qreigvals = qrmet(A,200)
    h = 2*np.pi/n
    qreigvals = np.sort(-qreigvals/h**2)
    return qreigvals

if __name__ == '__main__':
    q, n = 1.5, 200
    qreigvals = mathieu(q, n)
    print(qreigvals[:6])

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

from triang import forsub, backsub
from ludec import ludec
import numpy as np

def ftcs(fac, m, polds):
    pnews = np.copy(polds)
    for k in range(m-1):
        pnews[1:-1] = fac*(polds[:-2] + polds[2:] - 2*polds[1:-1])
        pnews[1:-1] += polds[1:-1]
        polds = np.copy(pnews)
    return pnews

def btcs(fac, m, polds):
    n = polds.size
    A = np.zeros((n,n))
    np.fill_diagonal(A, 1.+2*fac)
    A[0,0] = 1; A[-1,-1] = 1
    np.fill_diagonal(A[1:,:], -fac)   
    A[-1,-2] = 0
    np.fill_diagonal(A[:,1:], -fac)
    A[0,1] = 0
    LA, UA = ludec(A)

    for k in range(m-1):
        ys = forsub(LA,polds)
        pnews = backsub(UA,ys)
        polds = np.copy(pnews)

    return pnews

if __name__ == '__main__':
    al, L, T, n, m = 0.1, 1., 0.4, 51, 200
    hx = L/(n-1); ht = T/(m-1)
    fac = al*ht/hx**2
    pinits = 30*np.ones(n)
    pinits[:n//2] = 10
    phis = ftcs(fac, m, pinits); print(phis)
    phis = btcs(fac, m, pinits); print(phis)

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

from fft import fft
import numpy as np

def fft2(A):
    B = A.astype(complex)
    for i, row in enumerate(A):
        B[i,:] = fft(row)
    for j, col in enumerate(B.T):
        B[:,j] = fft(col)
    return B

def inversefft2(A):
    n2 = A.size
    newA = fft2(np.conjugate(A))/n2
    return np.conjugate(newA)

def func(x,y):
    return np.cos(3*x+4*y) - np.cos(5*x-2*y)

def poisson(a,b,n):
    h = (b-a)/(n-1)
    xs = a + np.arange(n)*h
    Xs, Ys = np.meshgrid(xs,xs)

    F = func(Xs, Ys)
    Ftil = fft2(F)

    ks = np.arange(n)
    Kxs, Kys = np.meshgrid(ks,ks)
    Denom = np.cos(2*np.pi*Kxs/n) + np.cos(2*np.pi*Kys/n) - 2
    Phitil = 0.5*Ftil*h**2/Denom
    Phitil[0,0] = 0

    Phi = np.real(inversefft2(Phitil))
    return Phi

if __name__ == '__main__':
    Phi = poisson(0, 2*np.pi, 128); print(Phi)