Draft: Master

Course material

src/code/landau_zener.py

 import matplotlib.pyplot as plt
 from matplotlib.animation import FuncAnimation
 from timeprop import propagate
-from adiabatic import _make_hamiltonian
 import numpy as np
 
+def make_hamiltonian(L=100, pot_func=None, time=0, t=1):
+
+    ham = np.zeros(shape=(L, L), dtype=complex)
+
+    if pot_func is not None:
+        pot = np.array([pot_func(i, time) for i in range(L)], dtype=float)
+    else:
+        pot = np.zeros(shape=(L,), dtype=float)
+
+    np.fill_diagonal(ham, 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 ham, pot
+
 def landau_zener_data(v, L=3, npts=100, gap=1):
     """
     Evolve a wavefunction over an adiabatic Hamiltonian of size L x L at velocity v
     """
 
-    v = 1/v
     if L == 2:
         def pot_r(i, t):
             return t*v*(-1)**i
@@ -32,9 +48,9 @@ def landau_zener_data(v, L=3, npts=100, gap=1):
     wfs.append(psi)
 
     for time in times:
-        ham = _make_hamiltonian(L=L, pot_func=pot_r, time=time, t=gap)[0]
+        ham = make_hamiltonian(L=L, pot_func=pot_r, time=time, t=gap)[0]
         psi = propagate(ham, psi, step)
-        en, _ = np.linalg.eigh(ham.toarray())
+        en, _ = np.linalg.eigh(ham)
         ens.append(en)
         wfs.append(psi)
 
@@ -44,8 +60,7 @@ def landau_zener_data(v, L=3, npts=100, gap=1):
 
     return times, ens, probs
 
-
-def animate_landau_zener_2plots(times, ens, probs, interval, title):
+def animate_landau_zener_2plots(times, ens, probs, interval, title, subtitles=None):
     """
     Generate animation of two systems with different velocities
     """
@@ -84,6 +99,9 @@ def animate_landau_zener_2plots(times, ens, probs, interval, title):
         ax.set_xticks([]);
         ax.set_yticks([]);
         ax.legend(fontsize=12)
+
+        if subtitles is not None:
+            ax.set_title(subtitles[k])
         k += 1
 
     def update(i):
@@ -92,11 +110,63 @@ def animate_landau_zener_2plots(times, ens, probs, interval, title):
         es1 = ens[0][i, :]
         es2 = ens[1][i, :]
         scatter1.set_offsets(np.vstack([ts1, es1]).T)
-        scatter1.set_sizes(probs[0][i]**2*size)
+        scatter1.set_sizes(probs[1][i]**2*size)
         scatter2.set_offsets(np.vstack([ts2, es2]).T)
-        scatter2.set_sizes(probs[1][i]**2*size)
+        scatter2.set_sizes(probs[0][i]**2*size)
         return scatter1, scatter2,
 
     anim = FuncAnimation(fig, update, interval=interval)
 
-    return anim
\ No newline at end of file
+    return anim
+
+
+def animate_landau_zener(times, ens, probs, L, interval, title):
+    """
+    Generate single animation
+    """
+
+    labels=[r'$|c_0(t)|^2$', r'$|c_1(t)|^2$', r'$|c_2(t)|^2$']
+    colors=['darkblue', 'red', 'green']
+    labels=labels[:L]
+    colors=colors[:L]
+    size = 200
+
+    fig, ax = plt.subplots(figsize=(5, 5))
+    plt.close()
+    j = 0
+
+    for level in ens.T:
+        ax.plot(times, level, c=colors[j])
+        j += 1
+    ax.set_xlabel('time')
+    ax.set_ylabel('Energy')
+
+    scatter = ax.scatter(times[0]*np.ones(L), ens[0, :], s=probs[0, :]**2*size, c=colors);
+
+    ax.set_xticks([]);
+    ax.set_yticks([]);
+    ax.set_title(title);
+
+    ytop = ens[0, -1]
+    ybot = ens[0, 0]
+    ax.set_ylim(ybot, ytop);
+
+    ax.set_xlabel('time')
+    ax.set_ylabel('Energy')
+
+    for i in range(L):
+        ax.scatter(0, 1000, s=size, label=labels[:L][i], c=colors[:L][i])
+
+    ax.legend(fontsize=12);
+
+    def update(i):
+
+        ts = times[i]*np.ones(L)
+        es = ens[i, :]
+        scatter.set_offsets(np.vstack([ts, es]).T);
+        scatter.set_sizes(probs[i]**2*size);
+        return scatter,
+
+    anim = FuncAnimation(fig, update, interval=interval);
+
+    return anim