def save_eigenfunction_couplings(filename_el, nr_kept, xmin, xmax, xsize, order): """ save_eigenfunction_couplings(filename_el, nr_kept, xmin, xmax, xsize, order) This program sets up the laser interaction hamiltonian for the eigenfunction basis, and stores it in an HDF5 file. This program must be run in parallel. Example ------- To run this program on 5 processors: $ mpirun -n 5 python -c "execfile('vibrational_BO.py');save_eigenstates()" """ #Retrieve the electronic energy curves. f = tables.openFile(filename_el) try: r_grid = f.root.R_grid[:] #Get number of tasks. el_basis_size = f.root.couplings.shape[0] finally: f.close() #Filter function, describing what index pairs should be included in the #calculations. def no_filter(index_pair): """ All couplings included. """ return True def symmetry_filter(index_pair): """ Only include the upper/lower triangular, since the hermeticity means that they are the same. """ i = index_pair[0] j = index_pair[1] if i >= j: return True else: return False #Make a list of the coupling indices that should be included. index_table = create_index_table(el_basis_size, no_filter) nr_tasks = len(index_table) #Initialize the B-spline_basis. spline_basis = vibrational_methods.Bspline_basis(xmin, xmax, xsize, order) vib_basis_size = spline_basis.nr_splines #Generate a filename. filename = name_gen.eigenfunction_couplings(filename_el, spline_basis) #Name of vib states. filename_vib = name_gen.vibrational_eigenstates(filename_el, spline_basis) #Parallel stuff #-------------- #Get processor 'name'. my_id = pypar.rank() #Get total number of processors. nr_procs = pypar.size() #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: f = tables.openFile(filename, 'w') g = tables.openFile(filename_vib) try: f.createArray("/", "electronicFilename", [filename_el]) #Initializing the arrays for the time dependent couplings of H. f.createCArray('/','couplings', tables.atom.FloatAtom(), (nr_kept * el_basis_size, nr_kept * el_basis_size), chunkshape=(nr_kept,nr_kept)) #Energy diagonal. Time independent part of H. energy_diagonal = zeros(nr_kept * el_basis_size) for i in range(el_basis_size): energy_diagonal[nr_kept * i: nr_kept * (i + 1)] = g.root.E[:nr_kept,i] f.createArray("/", "energyDiagonal", energy_diagonal) finally: f.close() g.close() #Save spline info. spline_basis.bsplines.save_spline_info(filename) #---------------------------------- #Setting up the hamiltonian #-------------------------- #Looping over the tasks of this processor. for i in my_tasks: #Retrieve indices. row_index, column_index = index_table[i] #Retrieve electronic couplings. f = tables.openFile(filename_el) try: el_coupling = f.root.couplings[row_index, column_index,:] finally: f.close() # #TODO REMOVE? # #Remove errors from the coupling. (A hack, unfortunately.) # r_grid_2, el_coupling_2 = remove_spikes(r_grid, el_coupling) # # #Setup potential matrix. # couplings = spline_basis.setup_potential_matrix( # r_grid_2, el_coupling_2) # #Setup potential matrix. Aij = <Bi | f(R) | Bj> bfb_matrix = spline_basis.setup_potential_matrix( r_grid, el_coupling) couplings = zeros([nr_kept, nr_kept]) #Retrieve eigensvectors. g = tables.openFile(filename_vib) try: Vr = g.root.V[:,:,row_index] Vc = g.root.V[:,:,column_index] finally: g.close() #Calculate couplings. for r_index in range(nr_kept): for c_index in range(nr_kept): couplings[r_index, c_index] = dot(Vr[:,r_index], dot(Vc[:,c_index])) #First file write. (Send, but not receive baton.) if starter == my_id: #Write to file. spline_basis.save_couplings(filename, couplings, row_index, column_index) #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. spline_basis.save_couplings(filename, couplings, row_index, column_index) #The rest of the file writes. else: #Receiving the baton from the previous writer. pypar.receive(receive_from, buffer = baton) #Write to file. spline_basis.save_couplings(filename, couplings, row_index, column_index) #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("Calculating couplings:", i, len(my_tasks)) #---------------------------- #Letting everyone catch up. pypar.barrier()
def save_all_eigenstates(filename_el, nr_kept, xmin, xmax, xsize, order): """ save_all_eigenstates(filename_el, nr_kept, xmin, xmax, xsize, order) This program solves the vibrational TISE for a set of energy curves, 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 -c "execfile('vibrational_BO.py');save_eigenstates()" """ #Retrieve the electronic energy curves. f = tables.openFile(filename_el) try: r_grid = f.root.R_grid[:] energy_curves = f.root.E[:] finally: f.close() #Initialize the B-spline_basis. spline_basis = vibrational_methods.Bspline_basis(xmin, xmax, xsize, order) spline_basis.setup_kinetic_hamiltonian() spline_basis.setup_overlap_matrix() #Generate a filename. filename = name_gen.vibrational_eigenstates(filename_el, spline_basis) #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(energy_curves) #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: f = tables.openFile(filename, 'w') try: f.createArray("/", "electronicFilename", [filename_el]) f.createArray("/", "R_grid", r_grid) f.createArray("/", "overlap", spline_basis.overlap_matrix) #Initializing the arrays for the eigenvalues and states. f.createCArray('/','E', tables.atom.FloatAtom(), (nr_kept, nr_tasks), chunkshape=(nr_kept, 1)) f.createCArray('/','V', tables.atom.FloatAtom(), (spline_basis.nr_splines, nr_kept, nr_tasks), chunkshape=(spline_basis.nr_splines, nr_kept, 1)) f.createCArray('/','hamiltonian', tables.atom.FloatAtom(), (spline_basis.nr_splines, spline_basis.nr_splines, nr_tasks), chunkshape=(spline_basis.nr_splines, spline_basis.nr_splines, 1)) finally: f.close() #Save spline info. spline_basis.bsplines.save_spline_info(filename) #---------------------------------- #Solving the TISE #---------------- #Looping over the tasks of this processor. for i in my_tasks: #TODO REMOVE? #remove_spikes removes points where the diagonalization has failed. #potential_hamiltonian = spline_basis.setup_potential_matrix( # r_grid, remove_spikes(energy_curves[i,:]) + 1/r_grid) #### #Setup potential matrix. potential_hamiltonian = spline_basis.setup_potential_matrix( r_grid, energy_curves[i,:] + 1/r_grid) #The total hamiltonian. hamiltonian_matrix = (spline_basis.kinetic_hamiltonian + potential_hamiltonian) #Diagonalizing the hamiltonian. E, V = spline_basis.solve(hamiltonian_matrix, nr_kept) #First file write. (Send, but not receive baton.) if starter == my_id: #Write to file. spline_basis.save_eigenstates(filename, E, V, hamiltonian_matrix, 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. spline_basis.save_eigenstates(filename, E, V, hamiltonian_matrix, i) #The rest of the file writes. else: #Receiving the baton from the previous writer. pypar.receive(receive_from, buffer = baton) #Write to file. spline_basis.save_eigenstates(filename, E, V, hamiltonian_matrix, 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("Vibrational BO calculations", i, len(my_tasks)) #---------------------------- #Letting everyone catch up. pypar.barrier()