2015-10-06-213210.ipynb

Kernel: Python 2

import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

x = np.linspace(0,2*np.pi,1000)
dx = x[1]-x[0]

s = np.sin(x)
ds = np.gradient(s,2)/dx
dds = np.gradient(ds,2)/dx

plt.plot(x,s)
plt.plot(x,ds)
plt.plot(x,dds)
plt.title("Finite Differences of sin(x)")
plt.xlabel("x")
plt.ylabel("$\sin$, $d\sin/dx$, $d^2\!\sin/dx^2$")

<matplotlib.text.Text at 0x7f3733782b50>

class Schrodinger(object):
    
    def __init__(self, V, start=-1, end=1, npts=100, mass=1.0):
        self.m = mass
        self.n = npts
        self.x = np.linspace(start,end,npts)
        self.dx = abs(self.x[1]-self.x[0])
        self.__init_laplacian__()
        
        self.Vx = V(self.x)        
        self.H = (-0.5/self.m)*self.laplace + np.diag(self.Vx)
        self.E = np.array([])
        self.psi = []
    
    def __init_laplacian__(self):
        M = -2*np.identity(self.n,'d')
        for i in xrange(1,self.n):
            M[i,i-1] = M[i-1,i] = 1
        self.laplace = M / self.dx**2
        
    def diagonalize(self):
        E,U = np.linalg.eigh(self.H)
        psi = []
        for i, e in enumerate(E):
            psi.append(U[:,i]/np.sqrt(self.dx) + e)
        self.E = E
        self.psi = psi
    
    def plot_spectra(self,levels=3):
        if not self.psi: return
        plt.plot(self.x, self.Vx ,color='k')
        for i in xrange(0,levels):
            plt.axhline(y=self.E[i],color='k',ls=":")
            plt.plot(self.x,self.psi[i])
        plt.title("Spectra and Eigenfunctions")
        plt.xlabel("Position")
        plt.ylabel("Energy")

s = Schrodinger(lambda x : 160*x**2)
s.diagonalize()
s.plot_spectra(10)

s = Schrodinger(lambda x : -40*np.cos(5*x))
s.diagonalize()
s.plot_spectra(5)