#!/usr/bin/python
'''
One-D Hamiltonian Code, Richard P. Muller
Copyright (c) 2000, Materials and Process Simulation Center, Caltech.
All Rights Reserved.

Hamiltonians:
(1) Particle in a box
(2) Harmonic oscillator
(3) One-D Hydrogen atom

Atomic units (hartree, bohr) are used throughout. The particle is
assumed to be an electron.
'''

import os,getopt,sys
from Numeric import *
from LinearAlgebra import *

def ev_sort(eigval,eigvec):
    # Since NumPy returns the eigenvectors/eigenvalues unsorted,
    # perform a sort on them together, based on the eigenvalues.
    # Also, NumPy returns the eigenvectors in the
    # rows, not in the columns (as is normal in diag routines)
    newval = zeros(len(eigval),Float)
    newvec = zeros(eigvec.shape,Float)
    index = argsort(eigval)
    for i in index:
        newval[i] = eigval[index[i]]
        newvec[i,:] = eigvec[index[i],:]

    return newval,newvec

def get_kinetic_energy(N):
    T = zeros((N,N),Float)
    T = T + identity(N) # Diagonal element
    for i in range(N-1): T[i,i+1] = T[i+1,i] = -0.5 # Tri diagonals
    return T

def get_particle_box_potential(N):
    border_width = 5
    V = zeros((N,N),Float)
    for i in range(border_width):
        V[i,i] = 100. 
        V[N-1-i,N-1-i] = 100. 
    return V

def get_harmonic_oscillator_potential(N):
    midpoint = N/2 + 0.5
    V = zeros((N,N),Float)
    K = 0.5/(N*N)
    for i in range(N):
        delx = i - midpoint
        V[i,i] = 0.5*K*delx*delx
    return V

def get_oned_hydrogen_potential(N):
    midpoint = N/2 + 0.5
    V = zeros((N,N),Float)
    Qeff = 3./N
    for i in range(N):
        delx = i - midpoint
        V[i,i] = -Qeff/abs(delx)
    return V

def write_gnuplot_driver(filename):
    file = open(filename,'w')
    file.write('set data style linespoints\n')
    file.write('set title \'Particle in a Box Eigenvectors\'\n')
    file.write('set xlabel \'Point\'\n')
    file.write('set ylabel \'Energy (h) \'\n')
    file.write('plot \'tmp.dat\' using 1, \'tmp.dat\' using 2 \n')
    file.write('pause -1\n')
    file.close()

def write_gnuplot_data_file(filename,vec):
    file = open(filename,'w')
    nvec,nel = vec.shape
    for i in range(nel):
        for j in range(nvec):
            file.write('%f ' % vec[j][i])
        file.write('\n')
    file.close()
    return

def make_even_integer(N):
    half_N = int(N/2)
    return 2*half_N

# Parameters
N = 50
opts,args = getopt.getopt(sys.argv[1:],'n:bsh')

potential = 'hydrogen'
for opt in opts:
    key,value = opt
    if key == '-n':
        N = eval(value)
    elif key == '-b':
        potential = 'box'
    elif key == '-s':
        potential = 'spring'
    elif key == '-h':
        potential = 'hydrogen'

make_even_integer(N) # N has to be even for H-atom
path_to_gnuplot = 'gnuplot'
#path_to_gnuplot = 'C:\Gnuplot3.7\wgnuplot.exe' # Works on Windows
gnu_filename = 'tmp.gnu'
dat_filename = 'tmp.dat'

T = get_kinetic_energy(N)

if potential == 'box':
    V = get_particle_box_potential(N)
elif potential == 'spring':
    V = get_harmonic_oscillator_potential(N)
else:
    V = get_oned_hydrogen_potential(N)
    
H = T + V

val,vec = eigenvectors(H)
val,vec = ev_sort(val,vec) # Eigenvalues aren't sorted in NumPy!

steps = range(N)
write_gnuplot_driver(gnu_filename)
write_gnuplot_data_file(dat_filename,vec[:3])
print val[:6]
os.system('%s tmp.gnu' % path_to_gnuplot)