#!/usr/bin/env python
# coding: utf-8

# ## Perturbation Theory: Energy Shifts and Perturbed Wavefunctions 
# 
# #### Advanced Quantum Physics content
# 
# In this exercise we will be investigating first order perturbation theory on a harmonic potential and a square well. Detailed below are the results derived within the AQP course which we shall be using.
# 
# Total Hamiltonian:
# \begin{equation}
#     \hat{H} = \hat{H}^{(0)}+\hat{H}'
# \end{equation}
# 
# First order energy shift:
# \begin{equation}
#     \Delta E_n^{(1)} = \langle n^{(0)} | \hat{H}' | n^{(0)} \rangle
# \end{equation}
# 
# First order perturbed wavefunctions:
# \begin{equation}
#     |n^{(1)}\rangle \approx | n^{(0)} \rangle + \sum_{m \neq n} | m^{(0)} \rangle \frac{\langle m^{(0)} | \hat{H}' | n^{(0)} \rangle}{E_n^{(0)}-E_m^{0}}
# \end{equation}
# 
# In the following we are working in units where $\hbar = 1$
# 
# Importing libraries which we will be using in this exercise, <i>pycav.quantum</i> contains functions specifically designed for this program. The following will outline the basic use of this program so that it can be used with ease in the exercise at the end.

# In[1]:


import numpy as np
from scipy.special import hermite
from scipy.misc import factorial
from scipy.integrate import quad

import matplotlib.pyplot as plt

import pycav.quantum as pm

get_ipython().run_line_magic('matplotlib', 'inline')


# Define a potential for use in plotting. The function should take an array argument, <i>params</i>, containing the necessary parameters which define the potential. The function should then return a function of position, x.
# 
# E.g. Harmonic potential: params[0] is the mass and params[1] is the angular frequency $V(x) = \frac{1}{2} m \omega^2 x^2$

# In[2]:


def harmonic_potential(params):
    def V(x):
        return 0.5*params[0]*params[1]**2*x**2
    return V


# Define the form of the unperturbed wavefunctions. The function should take the <i>params</i> argument as well as the primary quantum number, <i>n</i>. The function should then return the wavefunction as a function of position, like with the potential. The harmonic oscillator wavefunctions are of the form given here
# 
# Also define the form of the unperturbed energies. The function just needs to return the value of the energy of state $n$ by using the <i>params</i> array.

# In[3]:


def harmonic_unperturb_wf(params,n):
    alpha = params[0]*params[1]
    def psi_n(x):
        n_fact = factorial(n)
        y = np.sqrt(alpha)*x
        n_herm = hermite(n)
        return (alpha/np.pi)**(0.25)*(1./(2**n*n_fact)**0.5*n_herm(y)*np.exp(-y**2/2.))
    return psi_n
    
def harmonic_unperturb_erg(params,n):
    return (n+0.5)*params[1]


# With these functions in place we can plot the unperturbed system. This is done by the <i>plot_levels</i> function in the <i>perturb_module</i>. It will plot the first 4 wavefunctions, with their origins placed at their energy level. The function requires the arguments: <i>unperturb_wf, unperturb_erg, potential</i> and <i>params</i> with an additional optional arguments, <i>x_lims</i> which takes a list with the $x$ axis limits e.g. <i>x_lims</i> = [-0.1,0.1] and <i>plot_zoom</i> which takes a list which contains one element which is the factor by which wavefunctions are multiplied on the plot so they are visible.

# In[4]:


h_params = [0.5,2000.]
pm.plot_levels(harmonic_unperturb_wf,harmonic_unperturb_erg,harmonic_potential,
               h_params,x_lims = [-0.1,0.1], plot_zoom = [200])


# Now it is time to define the perturbing potential $\hat{H}'$. This is simply a function of $x$ which returns the potential value at $x$. If statements can be used to define potentials which have different behaviour in different regions. 

# In[5]:


def perturbation(x):
    return 1000.*x**4


# First order energy shifts and perturbed wavefunctions for individual levels can be called via <i>first_order_energy_sft</i> and <i>first_order_wf</i> respectively. These return the energy shift as a float and perturbed wavefunction as a function of position. The first argument for both is the principle quantum number of the energy level in question. For <i>first_order_energy_sft</i> a list of principle quantum numbers can also be passed and a list of energy shifts will be output. The mandatory arguments are familiar. 
# 
# There are the additional optional arguments <i>limits</i> and <i>tolerance</i>. <i>limits</i> puts the limits on the expectation value and overlap integrals done for the first order energy shift and perturbed wavefunction calculations respectively. Tolerance sets the level at which additional terms in the perturbed wavefunction sum are discarded. If the fraction accompanying $|m^{(0)}\rangle$ is less than the tolerance the sum is truncated.

# In[6]:


dE_0 = pm.first_order_energy_sft(0,perturbation,harmonic_unperturb_wf,h_params,limits = [-np.inf,np.inf])

perturb_1 = pm.first_order_wf(1,perturbation,harmonic_unperturb_wf,harmonic_unperturb_erg,
                              h_params,tolerance = 0.01,limits = [-np.inf,np.inf])


# Performs a similar plot to <i>plot_levels</i> with the unperturbed wavefunctions with faint lines and the perturbed wavefunctions with solid lines. The perturbed wavefunctions origin is also shifted by the first order energy shift.
# 
# <i>plot_perturb</i> has the optional arguments <i>x_lims, tolerance, limits</i> and <i>plot_zoom</i>. The first 3 have the usual meaning given above. <i>plot_zoom</i> is a list with 2 element, the first of which scales the size of the wavefunction on the plot (default value is 200). The second element scales the factor by which the first order energy shift is increased so that it can be viewed on the plot (default value is 100)

# In[7]:


pm.plot_perturb(perturbation,harmonic_unperturb_wf,harmonic_unperturb_erg,harmonic_potential,
                h_params,tolerance = 0.01,plot_zoom = [200,10000])


# ### Exercises
# 
# 1. For the harmonic potential, investigate linear and quadratic perturbations and higher order $x^n$. What do you notice about the first order energy shift as a function of $(n+1/2)$?
# (Theoretical Exercise: Why do the shifts have the shown function form? Hint: use ladder operators)
# 
# 2. For a square well with infintely high walls, write down the unperturbed wavefunctions and the unperturbed energies and create the required functions
# 
# 3. Investigate a small step in potential near the centre of the well such that perturbation has the form:
# \begin{align}
#     V'(x) &= \epsilon \ \mbox{for} \ |x| < b/2
# \end{align}
# where $b \ll a$ where $a$ is the width of the well. Give the first order energy shift to the ground state and first excited state
# 
# 4. Investigate a linear perturbation, $V'(x) \propto -x$, and give the first order energy shift to the ground state and first excited state. Now write a function to calculate the second order energy shift:
# \begin{equation}
#     \Delta E_n^{(2)} =  \sum_{m \neq n} \frac{\langle m^{(0)} | \hat{H}' | n^{(0)} \rangle}{E_n^{(0)}-E_m^{0}}
# \end{equation}
# Use [SciPy Quad](http://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.quad.html) to calculate the matrix elements. You should truncate the sum when the terms become negligibly small. Why should you perform seperate sums for even and odd $m$? Find the second order energy shift for the ground and first excited states.
# 
# 5. Use the optional argument <i>return_list = True</i> in <i>first_order_wf</i> to return a tuple containing (<i>perturbed_wf, k_list, I_list</i>). <i>k_list, I_list</i> contain a list of the principle quantum numbers and prefactors in the sum of the unperturbed states respectively. i.e.
# \begin{equation}
#     |n^{(1)}\rangle \approx |n^{(0)}\rangle + \sum_{k \in k_{list}} I_{list}[k] \ |k^{(0)}\rangle
# \end{equation}
# Use this to show the linear combination of states which creates the perturbed state (increasing the perturbation size makes this clearer). Also plot the perturbed wavefunction for the ground state, what is the interpretation of this?
# 
# 6. With the system in the ground state of the perturbed system ($V'(x) \propto -x$), write down the time dependence once the perturbation is switched off (only consider the above expansion to first order in the sum). 
# 
# 7. (Extra) Plot this evolution.

# ### Solution
# 
# #### Harmonic Potential:

# In[8]:


def perturbation(x):
    return 1000.*x

dE = pm.first_order_energy_sft([0,1,2],perturbation,harmonic_unperturb_wf,h_params,limits = [-np.inf,np.inf])
print(dE)

def perturbation(x):
    return 1000.*x**3

dE = pm.first_order_energy_sft([0,1,2],perturbation,harmonic_unperturb_wf,h_params,limits = [-np.inf,np.inf])
print(dE)


# All odd powers show vanishing first order energy shifts

# In[9]:


def perturbation(x):
    return 1000.*x**2

dE = pm.first_order_energy_sft([0,1,2,3,4],perturbation,harmonic_unperturb_wf,
                               h_params,limits = [-np.inf,np.inf])

fig1 = plt.figure(figsize = (12,12))
ax1 = plt.subplot(221)
ax1.plot([0.5,1.5,2.5,3.5,4.5],dE,'b-')

def perturbation(x):
    return 1000.*x**4

dE = pm.first_order_energy_sft([0,1,2,3,4],perturbation,harmonic_unperturb_wf,
                               h_params,limits = [-np.inf,np.inf])

ax2 = plt.subplot(222)
p = np.polyfit([0.5,1.5,2.5,3.5,4.5],dE,2)
ax2.plot([0.5,1.5,2.5,3.5,4.5],dE,'bx')
n = np.linspace(0.5,4.5,100)
ax2.plot(n,p[0]*n**2+p[2],'r-.')

def perturbation(x):
    return 1000.*x**6

dE = pm.first_order_energy_sft([0,1,2,3,4],perturbation,harmonic_unperturb_wf,
                               h_params,limits = [-np.inf,np.inf])

ax2 = plt.subplot(223)
p = np.polyfit([0.5,1.5,2.5,3.5,4.5],dE,3)
ax2.plot([0.5,1.5,2.5,3.5,4.5],dE,'bx')
n = np.linspace(0.5,4.5,100)
ax2.plot(n,p[0]*n**3+p[2]*n,'r-.')

def perturbation(x):
    return 1000.*x**8

dE = pm.first_order_energy_sft([0,1,2,3,4],perturbation,harmonic_unperturb_wf,
                               h_params,limits = [-np.inf,np.inf])

ax2 = plt.subplot(224)
p = np.polyfit([0.5,1.5,2.5,3.5,4.5],dE,4)
ax2.plot([0.5,1.5,2.5,3.5,4.5],dE,'bx')
n = np.linspace(0.5,4.5,100)
ax2.plot(n,p[0]*n**4+p[2]*n**2+p[4],'r-.')


# All even powers, $x^{2n}$, show polynomial dependence $\Delta E = a_n (n+1/2)^{n} + a_{n-2} (n+1/2)^{n-2} + ...$ down to either constant or linear term
# 
# #### Square Well:
# 
# Only parameters needed in <i>params</i> is well width, $a$, and particle mass, $m$.

# In[10]:


def square_unperturb_wf(params,n):
    a = params[0]
    m = n+1
    def psi_n(x):
        return np.sqrt(2./a)*np.sin((m*np.pi/a)*(x+a/2.))
    return psi_n
    
def square_unperturb_erg(params,n):
    m = n+1
    return (m**2*np.pi**2)/(2*params[1]*params[0]**2)

def square_potential(params):
    def V(x):
        return 0.0*x
    return V


# In[11]:


sq_params = [0.2,2.]
pm.plot_levels(square_unperturb_wf,square_unperturb_erg,square_potential,sq_params,plot_zoom = [25])


# In[12]:


def perturbation(x):
    if np.abs(x) < 0.001:
        return 0.01
    else:
        return 0.

dE = pm.first_order_energy_sft([0,1],perturbation,square_unperturb_wf,sq_params,
                               limits = [-sq_params[0]/2.,sq_params[0]/2.])
print(dE)


# Perturbation theory results suggest $\Delta E^{(1)}_1 \approx \frac{2\epsilon b}{a}$ and $\Delta E^{(1)}_2 \approx 0$. In this case $\frac{2\epsilon b}{a} = 0.0002$ for $\epsilon = 0.001$, $b = 0.002$ and $a = 0.2$

# In[13]:


def perturbation(x):
    return 1000.*x

dE = pm.first_order_energy_sft([0,1],perturbation,square_unperturb_wf,sq_params,
                               limits = [-sq_params[0]/2.,sq_params[0]/2.])
print(dE)


# Both first order energy shifts are, of order, $0$, leading contribution from second order energy shift

# In[14]:


def second_order_energy_sft(n,H,unperturb_wf,unperturb_erg,params,tolerance = 0.01,
                            limits = [-np.inf,np.inf]):
    
    def combine_functions(i,j):
        def integrand(x):
            psi_i = unperturb_wf(params,i)
            psi_j = unperturb_wf(params,j)
            return H(x)*psi_i(x)*psi_j(x)
        return integrand
      
    def E_diff(i,j):
        E_i = unperturb_erg(params,i)
        E_j = unperturb_erg(params,j)
        return E_i-E_j

    def find_k_values(even,k_list,I_list):
        truncate = False
        if even:
            k = 0
        else:
            k = 1
        while not truncate:
            if n != k:
                I, I_err = quad(combine_functions(n,k),limits[0],limits[1])
                E_nk = E_diff(n,k)
        
                if abs(I/E_nk) < tolerance:
                    truncate = True
                else:
                    I_list.append(I/E_nk)
                    k_list.append(k)
                    k += 2
            else:
                k += 2
        return k_list,I_list
    
    k_list = []
    I_list = []
    
    k_list,I_list = find_k_values(True,k_list,I_list)
    k_list,I_list = find_k_values(False,k_list,I_list)
       
    return sum(I_list)


# In[15]:


dE_0 = second_order_energy_sft(0,perturbation,square_unperturb_wf,square_unperturb_erg,sq_params,
                        tolerance = 0.001,limits = [-sq_params[0]/2.,sq_params[0]/2.])
dE_1 = second_order_energy_sft(1,perturbation,square_unperturb_wf,square_unperturb_erg,sq_params,
                        tolerance = 0.001,limits = [-sq_params[0]/2.,sq_params[0]/2.])

print(square_unperturb_erg(sq_params,0),square_unperturb_erg(sq_params,1))
print(dE_0,dE_1)


# At second order, an energy shift is seen for both the ground and first excited states (although with opposite sign).

# In[16]:


n = 0

perturb_n, k_list, I_list = pm.first_order_wf(n,perturbation,square_unperturb_wf,square_unperturb_erg,
                                            sq_params, tolerance = 0.1,
                                            limits = [-sq_params[0]/2.,sq_params[0]/2.],
                                            return_list = True)


k_list.append(n)
I_list.append(1.)

fig2 = plt.figure(figsize = (12,6))
ax = plt.subplot(111)
x = np.linspace(-0.1,0.1,200)

ax.plot(x,perturb_n(x),'b')

V = square_potential(sq_params)
perturb_h = np.array([perturbation(y) for y in x])*0.005

for k in k_list:
    unperturb_k = square_unperturb_wf(sq_params,k)
    if k != n:
        ax.plot(x,I_list[k_list.index(k)]*unperturb_k(x),'r',alpha = 0.5)
    else:
        ax.plot(x,unperturb_k(x),'b',alpha = 0.5)

ax.plot(x,V(x)+perturb_h,'k')
ax.set_xlim((np.min(x),np.max(x)))
ax.set_yticks([])


# The time dependence of the above state is given by:
# \begin{equation}
#     |\psi(t)\rangle = |1^{(0)}\rangle e^{-iE_1 t} + c_{21} |2^{(0)} \rangle e^{-iE_2 t}
# \end{equation}

# In[24]:


import pycav.display as display
import matplotlib.animation as anim

N = 200

def psi_t(i):
    t = i*0.0005
    sum_I = np.sum(I_list)
    for k in k_list:
        unperturb_k = square_unperturb_wf(sq_params,k)
        psi_r[:,i] += (I_list[k_list.index(k)]*unperturb_k(x)*np.cos(square_unperturb_erg(sq_params,k)*t)
                  /(sum_I*np.max(unperturb_k(x))))
        psi_i[:,i] += (I_list[k_list.index(k)]*unperturb_k(x)*np.sin(square_unperturb_erg(sq_params,k)*t)
                  /(sum_I*np.max(unperturb_k(x))))
        psi_a[:,i] += psi_r[:,i]**2+psi_i[:,i]**2


psi_r = np.zeros((x.shape[0],N))
psi_i = np.zeros((x.shape[0],N))
psi_a = np.zeros((x.shape[0],N))

for j in range(N):
    psi_t(j)

def nextframe(arg):
    line_r.set_data(x,psi_r[:,arg])
    line_i.set_data(x,psi_i[:,arg])
    line_a.set_data(x,psi_a[:,arg])
    
fig3 = plt.figure(figsize = (12,9))
ax1 = plt.subplot(211)
ax2 = plt.subplot(212)
ax1.set_xlim([np.min(x),np.max(x)])
ax1.set_ylim([-1.0,1.0])
ax2.set_xlim([np.min(x),np.max(x)])
ax2.set_ylim([0.0,1.])
line_r = ax1.plot(x,np.zeros_like(x),'b')[0]
line_i = ax1.plot(x,np.zeros_like(x),'r')[0]
line_a = ax2.plot(x,np.zeros_like(x),'k')[0]
animate1 = anim.FuncAnimation(fig3, nextframe, frames = N)
animate1 = display.create_animation(animate1, temp = True)
display.display_animation(animate1)


# In[ ]:





# In[ ]: