示例#1
0
def vib_problem_parameters():
    """
    filename_el, nr_kept, xmin, xmax, xsize, order = vib_problem_parameters()

    This is the place to choose the problem parameters of the vibrational
    problem.
    """
    
    #Number of vibrational states to be kept when diagonalizing.
    nr_kept = 200

    #Start of the R grid.
    xmin = 0.5
    
    #End of the R grid.
    xmax = 18.5

    #Approximat number of B-splines.
    xsize = 399

    #Order of the B-splines.
    order = 5
    
    #Name of the file where the electronic couplings are stored.
    #--------------------
    #Independent problem:
    #filename_el = ("el_couplings_lenght_z_m_0_nu_25_mu_15_beta_1_00_" + 
    #	"theta_0_00_m_0_q_3_E_5_00.h5")
    #-------------------- 
    #Generate name from info in el_tise_problem_parameters().
    m_max, nu_max, mu_max, R_grid, beta, theta = el_tise_problem_parameters()
    
    conf = config.Config(m = m_max, nu = nu_max, mu = mu_max, 
	R = R_grid[0], beta = beta, theta = theta)
    
    tdse_instance = tdse_electron.TDSE_length_z(conf = conf)

    filename, m_list, q_list, e_lim = el_tdse_problem_parameters()
    
    filename_el = name_gen.electronic_eig_couplings_R(tdse_instance, 
	m_list, q_list, e_lim)
    #--------------------

    return  filename_el, nr_kept, xmin, xmax, xsize, order
    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()