From 4f403ded0e9595332bf2774f7494b30feb3a11bf Mon Sep 17 00:00:00 2001
From: Sathish Kumar RK <rksathish09@gmail.com>
Date: Mon, 21 Jan 2019 16:14:27 +0000
Subject: [PATCH] harmonic_oscillator_plot.md

---
 src/1_einstein_model.md | 96 +++++++++++++++++++++++++++++++++++++++++
 1 file changed, 96 insertions(+)

diff --git a/src/1_einstein_model.md b/src/1_einstein_model.md
index 945b37e..151c177 100644
--- a/src/1_einstein_model.md
+++ b/src/1_einstein_model.md
@@ -65,6 +65,102 @@ This can be explained by considering a _quantum_ harmonic oscillator:
 
 ![Wave functions and energies of a harmonic oscillator](figures/harmonic.svg)
 
+```python
+import math
+from numpy.polynomial.hermite import Hermite
+
+def ho_evec(x, n, no_states):
+    
+    """
+    Calculate the wavefunction of states confined in the harmonic oscillator
+    
+    Input:
+    ------
+    x: numpy array of x coordinates (in units of hbar.omega)
+    n: n^th bound state in the oscillator
+    no_states: no of states confined
+    
+    Returns:
+    --------
+    Wavefunctions
+    
+    """
+    
+    # calculate hermite polynomial
+    vec = [0] * no_states
+    vec[n] = 1/2
+    Hn = Hermite(vec)
+    
+    return ((1/np.sqrt(math.factorial(n)*2**n))* 
+            pow(np.pi,-1/4)* 
+            np.exp(-pow(x, 2)/2)* 
+            Hn(x))
+
+def h0_ener(n):
+    """
+    Calculate the energy of nth bound state
+    """
+    return (n + 1/2)
+    
+x = np.linspace(-4, 4, 100) #units of hbar.omega
+no_states = 7 #no of bound states confined in the quantum well
+
+omega = 1.0 #frequency of harmonic oscillator
+V = 0.5*(omega**2)*(x**2)
+
+fig, ax = pyplot.subplots(figsize=(10, 7))
+
+ax.plot(x, V) #plot harmonic potential
+
+for i in range(no_states):
+    
+    ax.hlines(h0_ener(i), x[0], x[len(x)-1], linestyles='dotted', colors='k')
+    
+    ax.plot(x, ho_evec(x, i, no_states) + h0_ener(i)) #plot wavefunctions
+    
+    
+    # annotate plot
+    ax.text(x[len(x)-1], h0_ener(i)+1/4, '$\Psi_%2i (x)$' %(i), 
+             horizontalalignment='center', fontsize=14)
+    
+    ax.text(1/4, h0_ener(i)+1/4, '$E_%2i$' %(i), 
+             horizontalalignment='center', fontsize=14)
+        
+    if i==0: 
+        ax.text(x[0]+1/4, h0_ener(i)/4, r'$\frac{\hbar\omega}{2}$', 
+                 horizontalalignment='center', fontsize=14)
+        
+        ax.annotate("", xy=(x[0]+1/2, h0_ener(i)-1/2), 
+                    xytext=(x[0]+1/2, h0_ener(i)), 
+                    arrowprops=dict(arrowstyle="<->"))
+    elif i==1:
+        ax.text(x[0]+1/4, h0_ener(i-1)+1/3, '$\hbar\omega$', 
+                 horizontalalignment='center', fontsize=14)
+        
+        ax.annotate("", xy=(x[0]+1/2, h0_ener(i)), 
+                    xytext=(x[0]+1/2, h0_ener(i-1)), 
+                    arrowprops=dict(arrowstyle="<->"))
+        
+    ax.fill_between(x, h0_ener(i), ho_evec(x, i, no_states) + h0_ener(i))
+        
+# Move left y-axis and bottim x-axis to centre, passing through (0,0)
+ax.spines['left'].set_position('center')
+ax.spines['bottom'].set_position(('data', 0.0))
+
+# Eliminate upper and right axes
+ax.spines['right'].set_color('none')
+ax.spines['top'].set_color('none')
+
+# Eliminate x and y axes labels
+ax.set_yticklabels([])
+ax.set_xticklabels([])
+
+# Set x and y labels
+ax.set_xlabel('X '+'($\sqrt{\hbar/m\omega}$)', fontsize=12)
+ax.set_ylabel('E '+'($\hbar\omega$)', fontsize=12)
+ax.yaxis.set_label_coords(0.5,1)
+```
+
 $$\varepsilon_n=\left(n+\frac{1}{2}\right)\hbar\omega$$
 
 Phonons are bosons $\Rightarrow$ they follow Bose-Einstein statistics.
-- 
GitLab