Untitled

def task2p2(vname, V):
    N = 1000
    neigs = 6
    V = np.vectorize(V)
    xgrid = np.linspace(0,1, N+2)
       
    d, S = Schrödinger(V, N, xgrid, neigs)

    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("{}: {:.5f}".format(i+1,j))


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))
    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
    return v, S

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)
    print([np.power(np.pi*k, 2) for k in range(1,7, 1)] ) # Energy eigenvalues for Vx) = 0
    #task2p2("700*(0.5- abs(x - 0.5))", V = lambda x: 700*(0.5- abs(x-0.5)))