import numpy import scipy import scipy.linalg import matplotlib.pyplot as plt def schroedinger_matrix(V, h): N = len(V) A = numpy.zeros((N,N)) for i in range(1,N): A[i,i-1] = -1.0/(h*h) for i in range(0,N): A[i,i] = +2.0/(h*h) + V[i] for i in range(0, N-1): A[i,i+1] = -1.0/(h*h) return A def solve_schroedinger_scipy(V, h, n): A = schroedinger_matrix(V, h) N = len(V) print 'Computing eigenvalues...' E, psi = scipy.linalg.eigh(A, eigvals=(0,n-1)) return E, psi def qr_eigenvalues(A, tolerance): n = A.shape[0] I = numpy.identity(n) Ak = A.copy() error = tolerance+1 iteration = 0 while error > tolerance: shift = Ak[-1,-1] Q, R = scipy.linalg.qr(Ak - shift*I) Ak = numpy.dot(R, Q) + shift*I error = max([abs(Ak[i,k]) \ for i in range(n) \ for k in range(i)]) iteration += 1 print " iteration %d: error = %lf" % (iteration, error) return numpy.array([Ak[i,i] for i in range(n)]) def inverse_iteration(A, l, tolerance): n = A.shape[0] x = numpy.ones(n) Ashift = A - l*numpy.identity(n) error = tolerance + 1 while error > tolerance: x = scipy.linalg.solve(Ashift, x) x = x / scipy.linalg.norm(x) error = scipy.linalg.norm(l*x - numpy.dot(A, x)) return x def solve_schroedinger_qr(V, h, n, tolerance=1e-1): A = schroedinger_matrix(V, h) N = len(V) A = numpy.zeros((N,N)) for i in range(1,N): A[i,i-1] = -1.0/(h*h) for i in range(0,N): A[i,i] = +2.0/(h*h) + V[i] for i in range(0, N-1): A[i,i+1] = -1.0/(h*h) print 'Computing eigenvalues...' E = qr_eigenvalues(A, tolerance)[:n] psi = numpy.empty((N, n)) print 'Computing inverse iteration...' for i in range(n): psi[:, i] = inverse_iteration(A, E[i], tolerance) return E, psi N = 50 n = 5 L = 10.0 h = L/N x = numpy.linspace(0.0, L, N) V = numpy.zeros(N) # # finite well V[:N/5] = 2.0 V[-N/5:] = 2.0 # harmonic well #V = 0.1*(x-L/2)**2 # double well #V[N*0.4:N*0.6] = 2.0 # add "infinite" walls #V[0] = 1000.0 #V[-1] = 1000.0 # plot potential p1 = plt.subplot(221) p1.plot(x[1:-1], V[1:-1], 'o-') p1.set_title('psi') p2 = plt.subplot(222) p2.plot(x[1:-1], V[1:-1], 'o-') p2.set_title('psi**2') p3 = plt.subplot(223) p3.plot(x[1:-1], V[1:-1], 'o-') p3.set_title('psi') p4 = plt.subplot(224) p4.plot(x[1:-1], V[1:-1], 'o-') p4.set_title('psi**2') # V legends = ['V'] print "SciPy Solving..." Es, psis = solve_schroedinger_scipy(V, h, n) for i in range(n): E = Es[i] psi = psis[:,i] p1.plot(x, E+psi, 'o-') p2.plot(x, E+psi**2, 'o-') legends.append('E = %5.2f' % E) p2.legend(legends) # V legends = ['V'] print "QR Solving..." Es, psis = solve_schroedinger_qr(V, h, n, tolerance=1e-1) for i in range(n): E = Es[i] psi = psis[:,i] p3.plot(x, E+psi, 'o-') p4.plot(x, E+psi**2, 'o-') legends.append('E = %5.2f' % E) p4.legend(legends) E = [(i*numpy.pi/L)**2 for i in range(n)] print "E = ", E plt.show()