# Untitled

unknown
python
2 years ago
2.0 kB
2
Indexable
Never
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__':
#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)