Update tunneling section & adiabatic part

src/code/wkb.py

@@ -137,3 +137,91 @@ def plot_patching_region(V, Erange, start, stop, m, wf_scaling_factor, y_range):
     fig.update_yaxes(range=y_range)
 
     py.iplot(fig)
+
+def _wkb_free(x, m, E, V):
+
+    def p(x):
+        return np.sqrt(2 * m * (E - V(x)))
+
+    return 2/np.sqrt(p(x)) * np.sin(quad(p, 0, x)[0] + np.pi/4)
+
+wkb_free = np.vectorize(_wkb_free)
+
+def pot_f(A):
+    def V(x):
+        return A*np.exp(-(x-1500.0)**2/400**2)
+    return V
+
+def wkb_static_animation(x, E=1.5):
+
+    fig, ax = plt.subplots()
+    plt.close()
+
+    mags = np.linspace(-1.4, 1.4, 50)
+    V = pot_f(A=mags[0])
+    pot, = ax.plot(x, V(x), label=r'$V(x)$', color='black')
+    wkb, = ax.plot(x, np.real(wkb_free(x, m=0.01, E=E, V=V))/15+E, label=r'$\psi(x)$', color='red')
+    ax.set_xticks([]);
+    ax.set_yticks([]);
+
+    def animate(i):
+        V = pot_f(A=mags[i])
+        wkb.set_data(x, np.real(wkb_free(x, m=0.01, E=E, V=V))/15+E)
+        pot.set_data(x, V(x))
+        ax.set_ylim(-E, E+1)
+        ax.legend()
+        return wkb, pot,
+
+    anim = matplotlib.animation.FuncAnimation(fig, animate, frames=len(mags), interval=50)
+
+    return anim
+
+def _make_hamiltonian_tunel(L=100, pot_func=None):
+    t = 1
+    
+    ham = np.zeros(shape=(L, L), dtype=complex)
+    
+    if pot_func is not None:
+        pot = np.array([pot_func(i) for i in range(L)], dtype=float)
+    else:
+        pot = np.zeros(shape=(L,), dtype=float)
+
+    np.fill_diagonal(ham, 2*t+pot)
+    
+    offdiag = np.zeros(shape=(L-1,), dtype=complex)
+    offdiag[:] = -t
+    np.fill_diagonal(ham[1:, :-1], offdiag)
+    np.fill_diagonal(ham[:-1, 1:], offdiag)
+    
+    return scipy.sparse.csr_matrix(ham), pot
+
+def tunnel_animation(mags, potential_f):
+    
+    pots = []
+    wfs = []
+
+    for mag in mags:
+        V = potential_f(mag)
+        ham, pot = _make_hamiltonian_tunel(pot_func=V)
+        evals, evec = scipy.sparse.linalg.eigsh(ham, which='SM')
+        wfs.append(np.abs(evec.T[4]**2))
+        pots.append(pot)
+
+    fig, ax = plt.subplots()
+    plt.close()
+    x = np.linspace(0, 100, 100)
+    pot, = ax.plot(x, pots[0]*10, label=r'$V(x)$', color='black')
+    wf, = ax.plot(x, wfs[0], label=r'$\psi(x)$', color='red')
+    #ax.set_xticks([]);
+    #ax.set_yticks([]);
+    ax.set_ylim(min(wfs[-1]), max(wfs[-1])+0.01)
+
+    def animate(i):
+        pot.set_data(x, pots[i])
+        wf.set_data(x, wfs[i])
+        ax.legend()
+        return wf, pot,
+
+    anim = matplotlib.animation.FuncAnimation(fig, animate, frames=len(mags), interval=100)
+
+    return anim
\ No newline at end of file