Finish ps16-1

ps16-1.py

@@ -1,12 +1,12 @@
 import numpy as np
 import pylab as pl
 
-a = 5
+a = 10
 D = 1
 hbar = 1
 m = 1
-z = 2.851
-l = z / a
+zs = np.array([5.963, 2.851, 1.745, 1.124])
+l = zs / a
 V0 = 2
 
 
@@ -14,26 +14,31 @@ def cot(x):
     return 1 / np.tan(x)
 
 
-x = np.linspace(-10, 10, 1000)
+x = np.linspace(-20, 20, 1000)
 x_pos = x[x > 0]
 z0 = (a / hbar) * np.sqrt(2 * m * V0)
 print(f'z0 = {z0}')
 k = -l * cot(l * a)
 F = np.sin(l * a) / np.exp(-k * a)
 
-psi_conditions = [
-    (0 < x_pos) & (x_pos <= a),
-    (x_pos > a)
-]
-psi_funcs = [
-    lambda x: D * np.sin(l * x),
-    lambda x: F * np.exp(-k * x),
-]
-psi_half = np.piecewise(x_pos, psi_conditions, psi_funcs)
-psi = np.concatenate((-np.flip(psi_half), psi_half))
+psis = []
+for (i, z) in enumerate(zs):
+    psi_conditions = [
+        (0 < x_pos) & (x_pos <= a),
+        (x_pos > a)
+    ]
+    psi_funcs = [
+        lambda x: D * np.sin(l[i] * x),
+        lambda x: F[i] * np.exp(-k[i] * x),
+    ]
+    psi_half = np.piecewise(x_pos, psi_conditions, psi_funcs)
+    psi = np.concatenate((-np.flip(psi_half), psi_half))
+    psis.append(psi)
 
 V0_graph = np.piecewise(x, [(x < -a), (-a <= x) & (x <= a), (x > a)], [0, -V0, 0])
 
-pl.plot(x, psi)
+pl.rcParams['figure.dpi'] = 300
 pl.plot(x, V0_graph)
+for (i, psi) in enumerate(psis):
+    pl.plot(x, psi, label=f'psi_odd_{i} z={zs[i]}')
 pl.show()