aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorRahiel Kasim <rahielkasim@gmail.com>2016-06-04 13:49:15 +0200
committerRahiel Kasim <rahielkasim@gmail.com>2016-06-04 13:49:15 +0200
commitc48d93453e9bcca6763dba94e95ce68f1376de16 (patch)
treea678457262b98c9ec3daaefb89551de4c2fe0401
parent8c3520a0ba08c7e8d4850fa3926729ea82793bdc (diff)
finish schrodinger parabolic + plots
-rw-r--r--Set 3/schrödinger.py103
1 files changed, 66 insertions, 37 deletions
diff --git a/Set 3/schrödinger.py b/Set 3/schrödinger.py
index b97fd70..1442b35 100644
--- a/Set 3/schrödinger.py
+++ b/Set 3/schrödinger.py
@@ -1,7 +1,8 @@
# -*- coding: utf-8 -*-
+from math import sqrt, sin, pi
+
import numpy as np
import matplotlib.pyplot as plt
-from math import sqrt, sin, pi
N = 200
@@ -9,41 +10,47 @@ assert N % 2 == 0, "N should be even"
# potential is given for -a <= x <= a
a = 1
dx = 2 * a / N
-V0 = 1e6
+V0 = 10
hbar = m = 1
-p = "infinite"
def x(i):
return -a + i * dx
-A = np.zeros((N, N))
-def potential(i):
+def potential(i, p):
if p == "infinite":
return 0
elif p == "finite":
return V0 if i in (0, N) else 0
+ elif p == "finite_scaled":
+ return 0 if N // 10 <= i <= 9 * N // 10 else V0
elif p == "parabolic":
- b = dx
+ b = 500
return b * x(i) ** 2
-A[0, 0] = 2 + 2 * dx ** 2 * potential(0)
-A[0, 1] = -1
+def solve(p):
+ """Solve the time-independent Schrödinger equation with potential p.
+ Returns the energies and wave functions.
+ """
+ A = np.zeros((N, N))
-A[-1, -1] = 2 + 2 * dx ** 2 * potential(N)
-A[-1, -2] = -1
+ A[0, 0] = 2 + 2 * dx ** 2 * potential(0, p)
+ A[0, 1] = -1
-for i in range(1, N - 1):
- A[i, i - 1] = A[i, i + 1] = -1
- A[i, i] = 2 + 2 * dx ** 2 * potential(i)
+ A[-1, -1] = 2 + 2 * dx ** 2 * potential(N, p)
+ A[-1, -2] = -1
+ for i in range(1, N - 1):
+ A[i, i - 1] = A[i, i + 1] = -1
+ A[i, i] = 2 + 2 * dx ** 2 * potential(i, p)
-w, v = np.linalg.eig(A)
-eigenvectors = [v[:, i] for i in range(len(v))]
-w, v = zip(*sorted(zip(w, eigenvectors), key=lambda x: x[0])) # sort by eigenvalues
-energies = [e / (2 * dx ** 2) for e in w]
+ w, v = np.linalg.eig(A)
+ eigenvectors = [v[:, i] for i in range(len(v))]
+ w, eigenvectors = zip(*sorted(zip(w, eigenvectors), key=lambda x: x[0])) # sort by eigenvalues
+ energies = [e / (2 * dx ** 2) for e in w]
+ return energies, eigenvectors
def psi_inf(n):
@@ -62,39 +69,61 @@ def energy(n):
return (n * pi * hbar / (2 * a)) ** 2 / (2 * m)
-def plot(n):
+def plot(n, p, psi):
"""Plot probability density for quantum number n."""
- plt.plot(psi_inf(n) ** 2, label="analytic")
- plt.plot(v[n - 1] ** 2, label="discrete")
- pot = np.array([potential(i) for i in range(N)])
- pot = pot / np.linalg.norm(pot)
- plt.plot(pot, label="potential")
- plt.legend(ncol=3)
+ # plt.plot(psi_inf(n) ** 2, label="analytic")
+ c1 = "black"
+ fig, ax1 = plt.subplots()
+ ax1.plot(psi[n - 1] ** 2, label=r"$n$ = %d" % n, color=c1)
+ ax1.set_xlabel(r"$i$")
+ ax1.set_ylabel(r"$|\psi(x)|^2$", color=c1)
+ for t in ax1.get_yticklabels():
+ t.set_color(c1)
+
+ ax2 = ax1.twinx()
+ c2 = "#5b07ed"
+ pot = np.array([potential(i, p) for i in range(N)])
+ ax2.plot(pot, label="potential", color=c2, linewidth=4)
+ ax2.set_ylabel("potential", color=c2)
+ for t in ax2.get_yticklabels():
+ t.set_color(c2)
+
+ ncols = 1 if n > 2 else 2
+ # ask matplotlib for the plotted objects and their labels, from http://stackoverflow.com/a/10129461
+ lines, labels = ax1.get_legend_handles_labels()
+ lines2, labels2 = ax2.get_legend_handles_labels()
+ ax2.legend(lines + lines2, labels + labels2, loc="upper center", ncol=ncols)
+
+ ylim = {1: 0.037, 2: 0.027}
+ if n in ylim:
+ ax1.set_ylim([0, ylim[n]])
+
plt.title(r"Time-independent Schrödinger: $n = %d$" % n)
- plt.xlabel(r"$i$")
- plt.ylabel(r"$|\psi(x)|^2$")
- # plt.show()
- plt.savefig("%s_%d" % (p, n))
+ plt.show()
+ # plt.savefig("%s_%d" % (p, n))
plt.close()
def plot_energies(n=N):
- real = [energy(n) for n in range(1, n + 1)]
- calculated = energies[:n]
- assert len(real) == len(calculated)
- diff = [real[i] - calculated[i] for i in range(n)]
+ analytic = [energy(n) for n in range(1, n + 1)]
+ infinite = solve("infinite")[0][:n]
+ # finite = solve("finite_scaled")[0][:n]
+ parabolic = solve("parabolic")[0][:n]
+
quantum = range(1, n + 1)
s = 1
fig = plt.figure(figsize=(8, 5))
- plt.scatter(quantum, real, color='b', s=s, label="analytic")
- plt.scatter(quantum, calculated, color='r', s=s, label="calculated")
+ # plt.scatter(quantum, finite, color="#5d8715", s=s, label="finite")
+ plt.scatter(quantum, parabolic, color="#155787", s=s, label="parabolic")
+ plt.scatter(quantum, analytic, color='b', s=s, label="analytic infinite")
+ plt.scatter(quantum, infinite, color='r', s=s, label="infinite")
plt.grid(c="grey")
- plt.title("Energies for particle in the infinite potential well")
+ plt.title("Energies for particle in an infinite potential well")
plt.xlabel(r"$n$")
plt.ylabel("Energy")
plt.xlim([0, n])
plt.ylim([0, 60000])
plt.legend(loc="upper left")
- # plt.show()
- plt.savefig("energies", dpi=400)
+ plt.show()
+ # plt.savefig("energies", dpi=400)