Untitled

import matplotlib
import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import expm
from scipy import integrate
from scipy.sparse.linalg import spsolve
import scipy.sparse.linalg
from scipy.sparse import linalg as sla
import time
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.linalg.spsolve.html

def twopBVP(fvec, alpha, beta, L, N):
    dx = L/(N+1)
    fvec[0] -= alpha/dx**2 
    fvec[-1] -= beta/dx**2 #(f - b) from lecture notes
    Tdx = scipy.sparse.csr_matrix((1/(dx**2))*(np.diag(-2*np.ones(N)) + np.diag(np.ones(N-1),1) +  np.diag(np.ones(N-1),-1) ))
    y = spsolve(Tdx, fvec)
    return np.array([alpha, *y, beta])

def task1p1():
    L = 10
    alpha = 1
    beta = 2
    N = 100
    x = np.linspace(0, L, N+2)
    inx = x[1:-1]
    f = lambda x: np.exp(-x)
    f = np.vectorize(f)
    g = np.exp(-x)
    
    ys = twopBVP(f(inx), g[0], g[-1], L, N )
    
    Nbr = [2**4, 2**6, 2**8, 2**10, 2**12]
    err = []
    h = []
    for i in Nbr:
        delx = L/(i+1)
        ingrid = np.arange(0, L , delx)
        fvec = f(ingrid)
        alpha = f(ingrid[0])
        beta = f(L)
        fvec = np.delete(fvec, 0)
        y = twopBVP(fvec, alpha, beta, L, i)
        ingrid = np.append(ingrid, L)
        err.append(np.linalg.norm(y - f(ingrid) )*np.sqrt(delx))
        h.append(1/i)
    
    plt.figure(1)
    plt.loglog(h, err, 'm', label='error')
    plt.loglog(h, np.multiply(h,h), label='h^2')
    plt.title('Error for Custom d^2/dx^2y = f(x,y)')
    plt.legend()
    plt.grid()
    plt.show()


    plt.figure(2)
    plt.plot(x, ys, 'r--', label='approx.', )
    plt.plot(x,g, 'b--', label ='analytical sol.')
    plt.title("Approx. vs. analytical sol.")
    plt.legend()
    plt.grid()
    plt.show()


def task1p2():
    N = 999
    alpha = 0
    beta = 0
    L = 10
    x = np.linspace(0, L, N+2)
    inx = x[1:-1] 
    
    E = 1.9 * 1e11
    f = lambda x: ((-25000*x**2 + 25000*L*x)/(E*1e-3 * (3-2*np.cos(np.pi*x/L)**12)))
    f = np.vectorize(f)
    us = twopBVP(f(inx), alpha, beta, L , N)
    print('Deflection at midpoint (x,y):')
    np.printoptions(precision = 8) #does not work with correct usage of printoptions. 
    answr = [inx[500], 1e3*us[500]] #499th step will land exactly on L = 5.00 meters, yet index 500 has min value
    print(answr)

    plt.figure(1)
    plt.plot(x, us, 'r--', label = 'approx.')
    plt.ylabel('Deflection (m)')
    plt.title('Approximated Solution of The Beam Equation')
    plt.legend()
    plt.grid()
    plt.show()



def eigLaplace(L, N ):
    dx = L/(N+1)
    c = np.zeros(N)
    c[0] = -2
    c[1] = 1
    A = scipy.linalg.toeplitz(c)
    Tdx = (1/(dx**2))*(A)
    E = np.zeros((N,N))
    E[-1][-2] = 4/3
    E[-1][-1] = -1/3
    z  = Tdx + (1/dx**2)*E
    eigz, eigzvec = scipy.linalg.eig(z)
    return [eigz, eigzvec]


def Schrödinger(V, N, x, neigs):
    t0 = time.time()
    ingrid = x[1:-1]
    dx = 1/(N+1)
    Tdx = scipy.sparse.diags([1, -2, 1], [-1, 0, 1], shape=(N, N)) / dx**2
    T = Tdx - scipy.sparse.csr_matrix(V(ingrid)).multiply(scipy.sparse.identity(N))
    # help(sla.eigs)
    v, w = sla.eigsh(T, k=neigs, which ='SM')
    t1 = time.time()
    print(f'time: {t1 - t0}')
    z = np.zeros((1, neigs))
    S = np.vstack((z, w, z)) #Satisfying Dirichlet-conditions
    # print(np.size(S, 0), np.size(S, 1))

    return v, S

def task2p1():
    N = 499
    L = 1
    eigz, eigzvec = eigLaplace(L,N)
    eigenvaluef = lambda k: -1*(np.pi*(0.5 + k))**2

    plt.figure(1)
    plt.plot(eigzvec[:,np.argsort(abs(eigz))[0]], label = 'k=1')
    plt.plot(eigzvec[:,np.argsort(abs(eigz))[1]], label = 'k=2')
    plt.plot(eigzvec[:,np.argsort(abs(eigz))[2]], label = 'k=3')
    plt.legend()
    plt.title("First 3 Eigenvectors")
    plt.show()
    
    Z = np.arange(10,500,10)
    err = []
    h = []
    _, (ax1, ax2, ax3) = plt.subplots(3, 1, sharex=True)
    for i in Z:
        
        for k in range(1,4):
            y = eigLaplace(L, i)
            y[0].sort()
            err.append(abs(y[0][-1] - eigenvaluef(k)))
            h.append(i)
        
    def choice(kvalue):
        switch = {
        1:  ax1.loglog(h, err, 'm', label = 'k = {k}'),
        2:  ax2.loglog(h, err, 'm', label = 'k = {k}'),
        3:  ax3.loglog(h, err, 'm', label = 'k = {k}')
        }
        choice(k)
    plt.show()
        
       
    

def task2p2(vname, V):
    N = 1000
    neigs = 7
    V = np.vectorize(V)
    xgrid = np.linspace(0,1, N+2)
       
    d, S = Schrödinger(V, N, xgrid, neigs)
    # print(f'First {neigs} energy levels: {-1*d}')
    # print(f'First {neigs} wave functions: ')
    # print(S)
    
    pnorm = 1/N**2
    _, (ax1, ax2, ax3) = plt.subplots(3, 1, sharex=True)        
    ax1.plot(xgrid, V(xgrid), 'k', label = vname)
    ax1.legend(loc = 'upper right')
    for i in range(neigs):
        E = -d[i]
        psi = S[:,i]
        
        psi /= pnorm
        probDens = np.real(np.conj(psi)*psi) 

        ax2.plot(xgrid, psi + E)
        ax3.plot(xgrid, probDens + E)
    psiname = f'$\hat\psi$'
    psisquare = f'$|\hat\psi|^2$'
    ax1.set(title = f'Smallest {neigs} $\hat\psi$ and $|\hat\psi|^2 $ \n in ascending energy magnitude, \n with N = {N} and norm  = ' + r'$N^{-2}$.')
    ax1.set(ylabel = '$V$')
    ax2.set(ylabel = psiname)
    ax3.set(ylabel = psisquare)
    ax3.set(xlabel ='$x$')

    plt.show()
    
    print("Energy Eigenvalues, V(x) = " + vname)
    E = np.sort(-d)
    for i,j in enumerate(E):
        print("{}: {:.2f}".format(i+1,j))
    

    
if __name__ == '__main__':
    # task1p1()
    # task1p2()   
    task2p1()
    # task2p2("$800\sin(\pi x)^2$", V = lambda x: 800*np.multiply(np.sin(np.pi*x), np.sin(np.pi*x)) )
    # task2p2("$  0.5 (x - 0.5)^2  $", V = lambda x: 0.5*(x - 0.5)**2) 
    # task2p2("0", V = lambda x: 0)
    # task2p2("700*(0.5- abs(x - 0.5))", V = lambda x: 700*(0.5- abs(x-0.5)))