from numpy import *

# define grid
startx=-5
endx = 5
num_grid_points = 200
mass = 1.  # mass in atomic units
#mass = 39.9*1836
plot_num_states = 10


# generate grid
x = linspace(startx, endx, num_grid_points)
dx = (endx-startx)/float(num_grid_points)


# define potentials 
def HomogeneousPotential(grid,constant=0):
    return [constant for i in grid]

def HarmonicPotential(grid, alpha=1., beta=0., gamma=0.):
    return [ alpha*x**2 + beta*x+gamma for x in grid]

# define Operators
def OpV(potential):
    return diag(potential)

def OpT(M):
    """
    Kinetic energy (1/(2M)*Laplacian) 
        approximated by (-1/2*\delta_{i\pm1}+\delta_ii)/(M*dx**2) 
    """
    T = diag([1./(dx**2*M) for i in range(num_grid_points)])
    for i in range(num_grid_points-1):
        T[i,i+1] = -0.5/M/dx**2
        T[i+1,i] = -0.5/M/dx**2
    return T


# change potentials here
potential = HarmonicPotential(x)
potential = HomogeneousPotential(x)

# generate Hamiltonian and solve it
H = OpT(mass) + OpV(potential)
energies, wf = linalg.eigh(H)


# gnuplot output
print """
#set terminal postscript enh solid color "Helvetica" 22 lw 2
#set output "solved.ps"

set yrange[:%f]
# set xrange[:]

plot """ % ceil(energies[plot_num_states]),
for i in range(plot_num_states):
    print energies[i], " lw 2 lt 2 t '', ",
print "\\"

print """ "-" t "potential" with lines lt 1 lw 2,\\"""
for i in range(plot_num_states-1):
    print """ "" t "" with lines lw 1 lt 3,""", 
print """ "" t "" with lines lw 1 lt 3\n""" 
for (i,v) in enumerate(potential):
    print x[i], v
print "e"
for i in range(plot_num_states):
    for (j,v) in enumerate(wf[:,i]):
        print x[j],v+energies[i]
    print "e"