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

import numpy as np

def forsub(L,bs):
    n = bs.size
    xs = np.zeros(n)
    for i in range(n):
        xs[i] = (bs[i] - L[i,:i]@xs[:i])/L[i,i]
    return xs

def backsub(U,bs):
    n = bs.size
    xs = np.zeros(n)
    for i in reversed(range(n)):
        xs[i] = (bs[i] - U[i,i+1:]@xs[i+1:])/U[i,i]
    return xs

def testcreate(n,val):
    A = np.arange(val,val+n*n).reshape(n,n)
    A = np.sqrt(A)
    bs = (A[0,:])**2.1
    return A, bs

def testsolve(f,A,bs):
    xs = f(A,bs); print(xs)
    xs = np.linalg.solve(A,bs); print(xs)

if __name__ == '__main__':
    A, bs = testcreate(4,21)
    L = np.tril(A)
    testsolve(forsub,L,bs)
    print(" ")
    U = np.triu(A)
    testsolve(backsub,U,bs)

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

from triang import backsub, testcreate, testsolve
import numpy as np

def gauelim(inA,inbs):
    A = np.copy(inA)
    bs = np.copy(inbs)
    n = bs.size

    for j in range(n-1):
        for i in range(j+1,n):
            coeff = A[i,j]/A[j,j]
            A[i,j:] -= coeff*A[j,j:]
            bs[i] -= coeff*bs[j]

    xs = backsub(A,bs)
    return xs

if __name__ == '__main__':
    A, bs = testcreate(4,21)
    testsolve(gauelim,A,bs)

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

from triang import forsub, backsub, testcreate, testsolve
import numpy as np

def ludec(A):
    n = A.shape[0]
    U = np.copy(A)
    L = np.identity(n)

    for j in range(n-1):
        for i in range(j+1,n):
            coeff = U[i,j]/U[j,j]
            U[i,j:] -= coeff*U[j,j:]
            L[i,j] = coeff

    return L, U

def lusolve(A,bs):
    L, U = ludec(A)
    ys = forsub(L,bs)
    xs = backsub(U,ys)
    return xs

if __name__ == '__main__':
    A, bs = testcreate(4,21)
    testsolve(lusolve,A,bs)

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

from triang import backsub, testcreate, testsolve
import numpy as np

def gauelim_pivot(inA,inbs):
    A = np.copy(inA)
    bs = np.copy(inbs)
    n = bs.size

    for j in range(n-1):
        k = np.argmax(np.abs(A[j:,j])) + j
        if k != j:
            A[j,:], A[k,:] = A[k,:], A[j,:].copy()
            bs[j], bs[k] = bs[k], bs[j]

        for i in range(j+1,n):
            coeff = A[i,j]/A[j,j]
            A[i,j:] -= coeff*A[j,j:]
            bs[i] -= coeff*bs[j]

    xs = backsub(A,bs)
    return xs

if __name__ == '__main__':
    A, bs = testcreate(4,21)
    testsolve(gauelim_pivot,A,bs)

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

from triang import testcreate, testsolve
import numpy as np

def termcrit(xolds,xnews):
    errs = np.abs((xnews - xolds)/xnews)
    return np.sum(errs)

def jacobi(A,bs,kmax=50,tol=1.e-6):
    n = bs.size
    xnews = np.zeros(n)

    for k in range(1,kmax):
        xs = np.copy(xnews)

        for i in range(n):
            slt = A[i,:i]@xs[:i]
            sgt = A[i,i+1:]@xs[i+1:]
            xnews[i] = (bs[i] - slt - sgt)/A[i,i]

        err = termcrit(xs,xnews)
        print(k, xnews, err)
        if err < tol:
            break
    else:
        xnews = None

    return xnews

if __name__ == '__main__':
    n = 4; val = 21
    A, bs = testcreate(n,val)
    A += val*np.identity(n)
    testsolve(jacobi,A,bs)

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

from triang import testcreate
import numpy as np

def mag(xs):
    return np.sqrt(np.sum(xs*xs))

def power(A,kmax=6):
    zs = np.ones(A.shape[0])
    qs = zs/mag(zs)
    for k in range(1,kmax):
        zs = A@qs
        qs = zs/mag(zs)
        print(k,qs)

    lam = qs@A@qs
    return lam, qs

def testeigone(f,A,indx=0):
    eigval, eigvec = f(A)
    print(" "); print(eigval); print(eigvec)
    npeigvals, npeigvecs = np.linalg.eig(A)
    print(" ")
    print(npeigvals[indx]); print(npeigvecs[:,indx])

if __name__ == '__main__':
    A, _ = testcreate(4,21)
    testeigone(power,A)

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

from triang import forsub, backsub, testcreate
from ludec import ludec
from jacobi import termcrit
from power import mag, testeigone
import numpy as np

def invpowershift(A,shift=20,kmax=200,tol=1.e-8):
    n = A.shape[0]
    znews = np.ones(n)
    qnews = znews/mag(znews)
    Astar = A - np.identity(n)*shift
    L, U = ludec(Astar)

    for k in range(1,kmax):
        qs = np.copy(qnews)
        ys = forsub(L,qs)
        znews = backsub(U,ys)
        qnews = znews/mag(znews)

        if qs@qnews<0:
            qnews = -qnews

        err = termcrit(qs,qnews)
        #print(k, qnews, err)

        if err < tol:
            lam = qnews@A@qnews
            break
    else:
        lam = qnews = None

    return lam, qnews

if __name__ == '__main__':
    A, _ = testcreate(4,21)
    testeigone(invpowershift,A)

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

from triang import testcreate
from power import mag
import numpy as np

def qrdec(A):
    n = A.shape[0]
    Ap = np.copy(A)
    Q = np.zeros((n,n))
    R = np.zeros((n,n))
    for j in range(n):
        for i in range(j):
            R[i,j] = Q[:,i]@A[:,j]
            Ap[:,j] -= R[i,j]*Q[:,i]

        R[j,j] = mag(Ap[:,j])
        Q[:,j] = Ap[:,j]/R[j,j]
    return Q, R

def testqrdec(A):
    n = A.shape[0]
    Q, R = qrdec(A)
    diffa = A - Q@R
    diffq = np.transpose(Q)@Q - np.identity(n) 
    print(n, mag(diffa), mag(diffq))

if __name__ == '__main__':
    for n in range(4,10,2):
        A, _ = testcreate(n,21)
        testqrdec(A)

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

from triang import testcreate
from qrdec import qrdec
import numpy as np

def qrmet(inA,kmax=100):
    A = np.copy(inA)
    for k in range(1,kmax):
        Q, R = qrdec(A)
        A = R@Q
        #print(k, np.diag(A))

    qreigvals = np.diag(A)
    return qreigvals

if __name__ == '__main__':
    A, _ = testcreate(4,21)
    qreigvals = qrmet(A,6)
    print(" ")
    npeigvals, npeigvecs = np.linalg.eig(A); print(npeigvals)

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

from triang import testcreate
from invpowershift import invpowershift
from qrmet import qrmet
import numpy as np

def eig(A,eps=1.e-12):
    n = A.shape[0]
    eigvals = np.zeros(n)
    eigvecs = np.zeros((n,n))
    qreigvals = qrmet(A)
    for i, qre in enumerate(qreigvals):
        eigvals[i], eigvecs[:,i] = invpowershift(A,qre+eps)
    return eigvals, eigvecs

def testeigall(f,A):
    eigvals, eigvecs = f(A)
    npeigvals, npeigvecs = np.linalg.eig(A)
    print(eigvals); print(npeigvals)
    print(" ")
    for eigvec, npeigvec in zip(eigvecs.T,npeigvecs.T):
        print(eigvec); print(npeigvec)
        print(" ")
    
if __name__ == '__main__':
    A, _ = testcreate(4,21)
    testeigall(eig,A)

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

from triang import testcreate
from invpowershift import invpowershift
from qrmet import qrmet
import numpy as np

def makeH(A):
    n = A.shape[0]
    H = np.zeros((2*n,2*n))
    H[n:,:n] = A
    H[:n,n:] = A.T
    return H

def svd(A, solver=qrmet):
    H = makeH(A)
    shift = 1.
    Hstar = H - np.identity(H.shape[0])*shift
    Hvals = solver(Hstar) + shift
    indices = (Hvals>=0).nonzero()[0]
    S = Hvals[indices]

    n = A.shape[0]
    vals = np.zeros(2*n)
    vecs = np.zeros((2*n,2*n))
    for i, qre in enumerate(Hvals):
        vals[i], vecs[:,i] = invpowershift(H,qre+1.e-20,tol=1.e-8)
    Vs = [vecs[:n,i] for i in indices]
    V = np.sqrt(2)*np.column_stack(Vs)
    Us = [vecs[n:,i] for i in indices]
    U = np.sqrt(2)*np.column_stack(Us)
    return U, np.diag(S), V.T 

if __name__ == '__main__':
    A, _ = testcreate(4,21)
    A -= np.identity(4)
    U, S, VT = svd(A)
    print(np.all(np.linalg.norm(A - U@S@VT)<1.e-6), S)
    npU, npS, npVT = np.linalg.svd(A)
    diffA =  A - npU@np.diag(npS)@npVT
    print(np.all(np.linalg.norm(diffA)<1.e-6), npS)

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

import numpy as np

def paulimatrices():
    sigx = np.array([0.,1,1,0]).reshape(2,2)
    sigy = np.array([0.,-1j,1j,0]).reshape(2,2)
    sigz = np.array([1.,0,0,-1]).reshape(2,2)
    return sigx, sigy, sigz

def kron(U,V):
    n = U.shape[0]
    p = V.shape[0]
    W = np.zeros((n*p,n*p), dtype=np.complex64)
    for i in range(n):
        for k in range(n):
            for j in range(p):
                for l in range(p):
                    W[p*i+j,p*k+l] = U[i,k]*V[j,l]
    return W

if __name__ == '__main__':
    sigx, sigy, sigz = paulimatrices()
    allones = np.ones((3,3))
    kronprod = kron(sigx,allones); print(kronprod.real)

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

from kron import paulimatrices, kron
from qrmet import qrmet
import numpy as np

def twospins(omI,omII,gam):
    hbar = 1.
    paulis = paulimatrices()
    iden = np.identity(2)

    SIs = [hbar*kron(pa,iden)/2 for pa in paulis]
    SIIs = [hbar*kron(iden,pa)/2 for pa in paulis]
    SIdotII = sum([SIs[i]@SIIs[i] for i in range(3)])

    H = -omI*SIs[2] - omII*SIIs[2] + gam*SIdotII
    H = H.real
    return H

if __name__ == '__main__':
    H = twospins(1.,2.,0.5)
    qreigvals = qrmet(H); print(qreigvals)

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

from qrmet import qrmet
from kron import paulimatrices, kron
import numpy as np

def threespins(omI,omII,omIII,gam):
    hbar = 1.
    paulis = paulimatrices()
    iden = np.identity(2)

    SIs = [hbar*kron(kron(pa,iden),iden)/2 for pa in paulis]
    SIIs = [hbar*kron(kron(iden,pa),iden)/2 for pa in paulis]
    SIIIs = [hbar*kron(kron(iden,iden),pa)/2 for pa in paulis]

    SIdotII = sum([SIs[i]@SIIs[i] for i in range(3)])
    SIdotIII = sum([SIs[i]@SIIIs[i] for i in range(3)])
    SIIdotIII = sum([SIIs[i]@SIIIs[i] for i in range(3)])

    H = -omI*SIs[2] - omII*SIIs[2] - omIII*SIIIs[2]
    H += gam*(SIdotII+SIdotIII+SIIdotIII)
    H = H.real
    return H

if __name__ == '__main__':
    np.set_printoptions(precision=3)
    H = threespins(1.,2.,3.,0.5)
    qreigvals = qrmet(H); print(qreigvals)