def align_all_phases(self):
"""
align_all_phases()
Makes sure the eigenfunctions saved to file does not change sign from
one R to another.
"""
#File where the eigenstates are stored.
filename = name_gen.electronic_eigenstates_R(self.basis.config)
m_max = self.basis.config.m_max
#q_max is the same as mu_max.
q_max = self.basis.config.mu_max
f = tables.openFile(filename, "r+")
try:
for m in range(-m_max, m_max + 1):
m_group = name_gen.m_name(m)
for q in range(q_max + 1):
q_group = name_gen.q_name(q)
V = eval("f.root.%s.%s.V"%(m_group, q_group))
#Flipping the appropriate wavefunctions.
self.align_phases(V)
finally:
f.close()
def calculate_dipole_eig_R(self, index_array, R):
"""
dipole_matrix_eig, E = calculate_dipole_eig_R(index_array, R)
Calculates the dipole matrix in the eigenstate basis for a given
internuclear distance <R>.
Parameters
----------
index_array : (3 x N) int array, where row i contains
[m value, q value, index] of the i-th eigenstate basis function
in the eigenstate HDF5 file.
R : float, the internuclear distance for which the dipole matrix
should be calculated.
Returns
-------
dipole_matrix_eig : (N x N) complex array, containing the dipole
matrix in the eigenstate basis.
E : 1D float array, containing the energies of the eigenstate basis
function.
"""
#Create the correct basis.
temp_config = config.Config(self.config)
temp_config.R = R
basis = psc_basis.PSC_basis(temp_config)
#Computes the dipole matrix in the psc basis.
dipole_matrix_psc = self.dipole_psc(basis)
#Initialize eigenstate/value arrays.
V = zeros([basis.basis_size, index_array.shape[0]], dtype = complex)
E = zeros([index_array.shape[0]])
f = tables.openFile(self.eigenstate_file)
try:
#Find index of current R in the R_grid.
R_index = find(f.root.R_grid[:] == R)[0]
for i, [m, q, index] in enumerate(index_array):
m_group = name_gen.m_name(m)
q_group = name_gen.q_name(q)
V[:,i] = eval("f.root.%s.%s.V[:, %i, %i]"%(
m_group, q_group, index, R_index))
E[i] = eval("f.root.%s.%s.E[%i, %i]"%(
m_group, q_group, index, R_index))
finally:
f.close()
#Transform the dipole matrix to the eigenstate basis.
dipole_matrix_eig = dot(conj(transpose(V)), dot(dipole_matrix_psc, V))
return dipole_matrix_eig, E
def sort_in_m_and_q(self, E, V):
"""
E_and_V_dictionary = sort_in_m_and_q(E, V)
Sorts the eigenfunctions according to the m and q quantum number,
(i.e. number of nodes in the eta part of the function
(polar angle function, roughly)). Returns a dictionary with m and q
values as keys.
Parameters
----------
E : 1D float array, containing the sorted eigenvalues.
V : 2D complex array, containing the normalized eigenvectors.
Returns
-------
E_and_V_dictionary : nested dictionary instance, with m and q numbers as keys,
and a tuple of corresponding E and V constituents as values.
Notes
-----
The q quantum number seems to correspond to the l quantum number in
atoms. It is explained further in
Barmaki et al., Phys. Rev. A 69, 043403 (2004).
"""
#Initialize dictionaries.
index_dictionary = {}
E_and_V_dictionary = {}
#Getting the q quantum numbers from the eigenfunction.
q = self.extract_q(V)
#Loop over eigenstates.
for i in range(V.shape[1]):
#Singling out a wavefunction.
v = V[:,i]
#Getting the m quantum number from the eigenfunction.
m = self.extract_m(v)
m_key = name_gen.m_name(m)
q_key = name_gen.q_name(q[i])
#Add the index to the index dictionary.
if m_key in index_dictionary:
if q_key in index_dictionary[m_key]:
index_dictionary[m_key][q_key].append(i)
else:
index_dictionary[m_key][q_key] = [i]
else:
index_dictionary[m_key] = {q_key:[i]}
#Putting the actual eigenstates/energies in the dictionary.
for m_key, m_value in index_dictionary.iteritems():
for q_key, indices in m_value.iteritems():
if m_key in E_and_V_dictionary:
E_and_V_dictionary[m_key][q_key] = (E[indices],
V[:, indices])
else:
E_and_V_dictionary[m_key] = {q_key:(E[indices],
V[:, indices])}
return E_and_V_dictionary
def BO_dipole_couplings(self, m_list, q_list, E_lim):
"""
BO_dipole_couplings(m_list, q_list, E_lim)
Parallel program that calculates the dipole couplings for a
z-polarized laser in lenght gauge. An eigenstate basis is used, of
states whose quantum numbers are in <m_list> and <q_list>, that have
energies below <E_lim>. The couplings are stored to an HDF5 file.
Parameters
----------
m_list : list of integers, containing the m values wanted in
the basis.
q_list : list of integers, containing the q values wanted in
the basis.
E_lim : float, the upper limit of the energies wanted in
the basis, for R ~ 2.0.
Notes
-----
I sometimes observe unnatural spikes in the couplings
(as a function of R), which should be removed before the couplings
are used. I don't know why they are there.
Example
-------
>>> filename = "el_states_m_0_nu_70_mu_25_beta_1_00_theta_0_00.h5"
>>> tdse = tdse_electron.TDSE_length_z(filename = filename)
>>> m = [0]
>>> q = [0,1,2,3]
>>> E_lim = 5.0
>>> tdse.BO_dipole_couplings(m, q, E_lim)
"""
#Name of the HDF5 file where the couplings will be saved.
self.coupling_file = name_gen.electronic_eig_couplings_R(self,
m_list, q_list, E_lim)
#Parallel stuff
#--------------
#Get processor 'name'.
my_id = pypar.rank()
#Get total number of processors.
nr_procs = pypar.size()
#Size of eigenstate basis. (Buffer for broadcast.)
basis_size_buffer = r_[0]
#Get number of tasks.
f = tables.openFile(self.eigenstate_file)
try:
R_grid = f.root.R_grid[:]
finally:
f.close()
nr_tasks = len(R_grid)
#Get a list of the indices of this processors share of R_grid.
my_tasks = nice_stuff.distribute_work(nr_procs, nr_tasks, my_id)
#The processors will be writing to the same file.
#In order to avoid problems, the procs will do a relay race of writing to
#file. This is handeled by blocking send() and receive().
#Hopefully there will not be to much waiting.
#ID of the processor that will start writing.
starter = 0
#ID of the processor that will be the last to write.
ender = (nr_tasks - 1) % nr_procs
#Buffer for the baton, i.e. the permission slip for file writing.
baton = r_[0]
#The processor one is to receive the baton from.
receive_from = (my_id - 1) % nr_procs
#The processor one is to send the baton to.
send_to = (my_id + 1) % nr_procs
#-------------------------------
#Initializing the HDF5 file
#--------------------------
if my_id == 0:
#Initialize index list.
index_array = []
#Find the index of the R closest to 2.0.
R_index = argmin(abs(R_grid - 2.0))
#Choose basis functions.
f = tables.openFile(self.eigenstate_file)
try:
for m in m_list:
m_group = name_gen.m_name(m)
for q in q_list:
q_group = name_gen.q_name(q)
for i in range(self.config.nu_max + 1):
if eval("f.root.%s.%s.E[%i,%i]"%(m_group, q_group,
i, R_index)) > E_lim:
break
else:
#Collect indices of the basis functions.
index_array.append(r_[m, q, i])
finally:
f.close()
#Cast index list as an array.
index_array = array(index_array)
#Number of eigenstates in the basis.
basis_size = len(index_array)
print basis_size, "is the basis size"
basis_size_buffer[0] = basis_size
f = tables.openFile(self.coupling_file, 'w')
try:
f.createArray("/", "R_grid", R_grid)
#Saving the index array.
f.createArray("/", "index_array", index_array)
#Initializing the arrays for the couplings and energies.
f.createCArray('/', 'E',
tables.atom.FloatAtom(),
(basis_size, nr_tasks),
chunkshape=(basis_size, 1))
f.createCArray('/', 'couplings',
tables.atom.ComplexAtom(16),
(basis_size, basis_size, nr_tasks),
chunkshape=(basis_size, basis_size, 1))
finally:
f.close()
#Save config instance.
self.config.save_config(self.coupling_file)
#----------------------------------
#Calculating the dipole couplings
#--------------------------------
#Broadcasting the basis size from processor 0.
pypar.broadcast(basis_size_buffer, 0)
#Initializing the index array.
if my_id != 0:
index_array = zeros([basis_size_buffer[0], 3], dtype=int)
#Broadcasting the index array from proc. 0.
pypar.broadcast(index_array, 0)
#Looping over the tasks of this processor.
for i in my_tasks:
#Calculate the dipole couplings for one value of R.
couplings, E = self.calculate_dipole_eig_R(index_array, R_grid[i])
#First file write. (Send, but not receive baton.)
if starter == my_id:
#Write to file.
self.save_dipole_eig_R(couplings, E, R_grid[i])
#Avoiding this statement 2nd time around.
starter = -1
#Sending the baton to the next writer.
pypar.send(baton, send_to, use_buffer = True)
#Last file write. (Receive, but not send baton.)
elif i == my_tasks[-1] and ender == my_id :
#Receiving the baton from the previous writer.
pypar.receive(receive_from, buffer = baton)
#Write to file.
self.save_dipole_eig_R(couplings, E, R_grid[i])
#The rest of the file writes.
else:
#Receiving the baton from the previous writer.
pypar.receive(receive_from, buffer = baton)
#Write to file.
self.save_dipole_eig_R(couplings, E, R_grid[i])
#Sending the baton to the next writer.
pypar.send(baton, send_to, use_buffer = True)
#Showing the progress of the work.
if my_id == 0:
nice_stuff.status_bar("Electronic dipole couplings:",
i, len(my_tasks))
#----------------------------
#Letting everyone catch up.
pypar.barrier()
def save_electronic_eigenstates(m_max, nu_max, mu_max, R_grid, beta, theta):
"""
save_electronic_eigenstates(m_max, nu_max, mu_max, R_grid, beta, theta)
This program solves the electronic TISE for a range of internuclear
distances, given in <R_grid>, and stores them in an HDF5 file.
This program must be run in parallel.
Example
-------
To run this program on 5 processors:
$ mpirun -n 5 python electronic_BO.py
"""
#Parallel stuff
#--------------
#Get processor 'name'.
my_id = pypar.rank()
#Get total number of processors.
nr_procs = pypar.size()
#Get number of tasks.
nr_tasks = len(R_grid)
#Get a list of the indices of this processors share of R_grid.
my_tasks = nice_stuff.distribute_work(nr_procs, nr_tasks, my_id)
#The processors will be writing to the same file.
#In order to avoid problems, the procs will do a relay race of writing to
#file. This is handeled by blocking send() and receive().
#Hopefully there will not be to much waiting.
#ID of the processor that will start writing.
starter = 0
#ID of the processor that will be the last to write.
ender = (nr_tasks - 1) % nr_procs
#Buffer for the baton, i.e. the permission slip for file writing.
baton = r_[0]
#The processor one is to receive the baton from.
receive_from = (my_id - 1) % nr_procs
#The processor one is to send the baton to.
send_to = (my_id + 1) % nr_procs
#-------------------------------
#Initializing the HDF5 file
#--------------------------
if my_id == 0:
#Creates a config instance.
my_config = config.Config(m = m_max, nu = nu_max, mu = mu_max,
R = R_grid[0], beta = beta, theta = theta)
#Number of basis functions.
basis_size = (2 * m_max + 1) * (nu_max + 1) * (mu_max + 1)
#Generate a filename.
filename = name_gen.electronic_eigenstates_R(my_config)
f = tables.openFile(filename, 'w')
try:
f.createArray("/", "R_grid", R_grid)
#Looping over the m values.
for m in range(-1 * m_max, m_max + 1):
#Creating an m group in the file.
m_group = name_gen.m_name(m)
f.createGroup("/", m_group)
#Looping over th q values.
for q in range(mu_max + 1):
#Creating a q group in the m group in the file.
q_group = name_gen.q_name(q)
f.createGroup("/%s/"%m_group, q_group)
#Initializing the arrays for the eigenvalues and states.
f.createCArray('/%s/%s/'%(m_group, q_group),'E',
tables.atom.FloatAtom(),
(basis_size/(mu_max + 1), nr_tasks),
chunkshape=(basis_size/(mu_max + 1), 1))
f.createCArray('/%s/%s/'%(m_group, q_group),'V',
tables.atom.ComplexAtom(16),
(basis_size, basis_size/(mu_max + 1), nr_tasks),
chunkshape=(basis_size, basis_size/(mu_max + 1), 1))
finally:
f.close()
#Save config instance.
my_config.save_config(filename)
#----------------------------------
#Solving the TISE
#----------------
#Looping over the tasks of this processor.
for i in my_tasks:
#Creating TISE instance.
tise = tise_electron.TISE_electron(m = m_max, nu = nu_max,
mu = mu_max, R = R_grid[i], beta = beta, theta = theta)
#Diagonalizing the hamiltonian.
E,V = tise.solve()
#First file write. (Send, but not receive baton.)
if starter == my_id:
#Write to file.
tise.save_eigenfunctions_R(E, V, R_grid[i])
#Avoiding this statement 2nd time around.
starter = -1
#Sending the baton to the next writer.
pypar.send(baton, send_to, use_buffer = True)
#Last file write. (Receive, but not send baton.)
elif i == my_tasks[-1] and ender == my_id :
#Receiving the baton from the previous writer.
pypar.receive(receive_from, buffer = baton)
#Write to file.
tise.save_eigenfunctions_R(E, V, R_grid[i])
#The rest of the file writes.
else:
#Receiving the baton from the previous writer.
pypar.receive(receive_from, buffer = baton)
#Write to file.
tise.save_eigenfunctions_R(E, V, R_grid[i])
#Sending the baton to the next writer.
pypar.send(baton, send_to, use_buffer = True)
#Showing the progress of the work.
if my_id == 0:
nice_stuff.status_bar("Electronic BO calculations",
i, len(my_tasks))
#----------------------------
#Letting everyone catch up.
pypar.barrier()
#Since the sign of the eigenfunctions are completely arbitrary, one must
#make sure they do not change sign from one R to another.
if my_id == 0:
tise.align_all_phases()
#Letting 0 catch up.
pypar.barrier()
