def distribute_Q_over_procs(self, num_procs): """ num_Q_per_proc is determined as the largest integer dividing the total_Qsteps number. the remainder is placed on rank 0 (if there is a remainder...) """ # set up array of how many Q on each proc nQpp_arr = np.zeros(num_procs).astype(int) # Num_Q_Per_Processor_ARRay proc = 0 for qq in range(self.total_Qsteps): if proc == num_procs: proc = 0 nQpp_arr[proc] = nQpp_arr[proc] + 1 proc = proc + 1 # if any procs have 0 Q, print error message and exit if nQpp_arr.min() == 0: message = 'atleast one processor will do 0 Qpoints.\n' \ ' increase number of procs or decrease number of Q points' raise PSF_exception(message) # print parellelism info message = f'process: 0 Q points: {nQpp_arr[0]:g}\n' for ii in range(1, num_procs): message = message + f' process: {ii:g} Q points: {nQpp_arr[ii]:g}\n' print_stdout(message, msg_type='Q points on each process') # put the Qpoint indicies for each proc in the list self.Q_on_procs = [] shift = 0 for ii in range(num_procs): self.Q_on_procs.append(list(range(shift, shift + nQpp_arr[ii]))) shift = shift + nQpp_arr[ii]
def __init__(self, invars, Qpoints, rank): """ setup freq. grid from MD params and initialize variables to hold SQW i tried padding the FFT w 0's, but spectral leakage was waaayyy worse """ self.rank = rank # the calling mpi rank # effective number of steps, i.e. num in each block that is read and computed self.block_steps = (invars.total_steps // invars.stride) // invars.num_blocks # max freq is actually self.max_freq/2 according to nyquist theorem self.max_freq = 1e-12 / invars.dt / invars.stride * 4.13567 self.df = self.max_freq / self.block_steps # freq. resolution self.num_freq = self.block_steps self.meV = np.linspace(0, self.max_freq, self.num_freq) if self.rank == 0: message = (f'max freq: {self.max_freq/2:2.3f} meV\n' f' number of (positive) freq.: {self.num_freq//2}\n' f' resolution: {self.df:2.3e} meV\n') print_stdout(message, msg_type='frequency Grid') self.sqw = np.zeros((self.num_freq, Qpoints.Qsteps)) # SQW array self.pos = np.zeros( (self.block_steps, invars.num_atoms, 3)) # time-steps, atoms, xyz self.atom_ids = np.zeros( (self.block_steps, invars.num_atoms)).astype(int) # see mod_io self.b_array = np.zeros( (self.block_steps, invars.num_atoms)) # see below self.box_lengths = [0, 0, 0] # read from traj file self.num_blocks = len(invars.blocks) # see mod_invars self.counter = 1
def recompute_lattice(self): """ recompute lattice vectors etc. from data read from traj file. """ self.r_lattice_vectors = np.zeros((3, 3)) self._compute_reciprocal_lattice() # print the lattice/reciprocal lattice message = ( f'real space lattice from trajectory file (Angstrom):\n' f' {self.lattice_vectors[0,0]: 2.3f} {self.lattice_vectors[0,1]: 2.3f}' f' {self.lattice_vectors[0,2]: 2.3f}\n {self.lattice_vectors[1,0]: 2.3f}' f' {self.lattice_vectors[1,1]: 2.3f} {self.lattice_vectors[1,2]: 2.3f}\n' f' {self.lattice_vectors[2,0]: 2.3f} {self.lattice_vectors[2,1]: 2.3f}' f' {self.lattice_vectors[2,2]: 2.3f}\n') print_stdout(message, msg_type='NOTE') message = ( f'reciprocal space lattice from trajectory file (1/Angstrom):\n' f' {self.r_lattice_vectors[0,0]: 2.3f} {self.r_lattice_vectors[0,1]: 2.3f}' f' {self.r_lattice_vectors[0,2]: 2.3f}\n {self.r_lattice_vectors[1,0]: 2.3f}' f' {self.r_lattice_vectors[1,1]: 2.3f} {self.r_lattice_vectors[1,2]: 2.3f}\n' f' {self.r_lattice_vectors[2,0]: 2.3f} {self.r_lattice_vectors[2,1]: 2.3f}' f' {self.r_lattice_vectors[2,2]: 2.3f}') print_stdout(message)
def distribute_Q_over_procs(self,invars,num_ranks): """ num_Q_per_proc is determined as the largest integer dividing the total_Qsteps number. the remainder is placed on rank 0 (if there is a remainder...) should probably use scatter/gather methods from MPI, but don't know how yet """ # set up array of how many Q on each proc nQpp_arr = np.zeros(num_ranks).astype(int) # Num_Q_Per_Processor_ARRay proc = 0 for qq in range(self.total_Qsteps): if proc == num_ranks: proc = 0 nQpp_arr[proc] = nQpp_arr[proc]+1 proc = proc+1 # if any procs have 0 Q, print a warning if nQpp_arr.min() == 0: message = 'atleast one processor will have 0 Qpoints.\n this is probably not effecient.' print_stdout(message,msg_type='WARNING') # print parellelism info message = f'process: 0 Q points: {nQpp_arr[0]:g}\n' for ii in range(1,num_ranks): message = message+f' process: {ii:g} Q points: {nQpp_arr[ii]:g}\n' print_stdout(message,msg_type='Q points on each process') # put the Qpoints for each proc in the list self.Q_on_procs = [] shift = 0 for ii in range(num_ranks): self.Q_on_procs.append(self.total_reduced_Q[shift:shift+nQpp_arr[ii],:]) shift = shift+nQpp_arr[ii]
def parse_trajectory(self,invars,sqw): """ get the trajectories, atom_types, and box sizes from the hdf5 files. can add methods to get the data from differnt file formats, but ultimately the pos, atom_types, and box_lengths arrays that are returend need to be consistent with my code. pos has shape [number of time steps in block, number of atoms, 3-dimensions (i.e. x,y,z)] atom_types has shape [number of time steps in block, number of atoms] note that ids here really means TYPES, but i am too lazy to go change the rest of my code. for now, my code assumes that the ids are sorted at each time step so that they are identical at each step. if they arent sorted, the ids(=TYPES) at each step can be used to assign the correct scattering lenght to the atoms. box_bounds has shape [3], 1 for each cartesian direction. to use the lammps read, lammps should dump using a command like: dump pos all h5md ${dt_dump} pos.h5 position species dump_modify pos sort id """ if sqw.rank == 0: message = 'now reading positions' print_stdout(message,msg_type='NOTE') # the indicies to be sliced. inds = [sqw.block_index*sqw.block_steps,(sqw.block_index+1)*sqw.block_steps] if invars.parse_custom: # get from my old format made using the 'compressor' tool sqw.box_lengths[0] = np.mean(self.handle['box_bounds'][inds[0]:inds[1],1]- self.handle['box_bounds'][inds[0]:inds[1],0]) sqw.box_lengths[1] = np.mean(self.handle['box_bounds'][inds[0]:inds[1],3]- self.handle['box_bounds'][inds[0]:inds[1],2]) sqw.box_lengths[2] = np.mean(self.handle['box_bounds'][inds[0]:inds[1],5]- self.handle['box_bounds'][inds[0]:inds[1],4]) sqw.pos[:,:,:] = self.handle['pos_data'][inds[0]:inds[1],:,:] # get the positins sqw.atom_types[0,:] = self.handle['atom_types'][:] # get the atom TYPES else: # read from lammps output. sqw.box_lengths[0] = np.mean(self.handle['particles']['all']['box']['edges']['value'] [inds[0]:inds[1],0],axis=0) sqw.box_lengths[1] = np.mean(self.handle['particles']['all']['box']['edges']['value'] [inds[0]:inds[1],1],axis=0) sqw.box_lengths[2] = np.mean(self.handle['particles']['all']['box']['edges']['value'] [inds[0]:inds[1],2],axis=0) sqw.pos[:,:,:] = self.handle['particles']['all']['position']['value'][inds[0]:inds[1],:,:] sqw.atom_types[:,:] = self.handle['particles']['all']['species']['value'][inds[0]:inds[1],:] # optionally unimpose minimum image convention if invars.unwrap_pos: if sqw.rank == 0: message = 'unwrapping positions' print_stdout(message,msg_type='NOTE') self._unwrap_positions(invars,sqw)
def __init__(self, invars): """ store lattice and reciprocal lattice vectors """ self.lattice_vectors = invars.lattice_vectors self.r_lattice_vectors = np.zeros((3, 3)) # print cell lengths from INPUT file message = (f'cell lengths from input: {self.lattice_vectors[0,0]} ' f'{self.lattice_vectors[1,1]} ' f'{self.lattice_vectors[2,2]} Angstrom') print_stdout(message, msg_type='NOTE') # print whether or not lattice vectors will be recalculated from traj. file if not invars.recalculate_cell_lengths: message = 'using cell lengths from input\n' print_stdout(message, msg_type='NOTE') else: message = 'using cell lengths from hdf5 trajectory file' print_stdout(message, msg_type='NOTE') # set up reciprocal lattice self._compute_reciprocal_lattice() # print the lattice/reciprocal lattice message = ( f'real space lattice from input file (Angstrom):\n' f' {self.lattice_vectors[0,0]: 2.3f} {self.lattice_vectors[0,1]: 2.3f}' f' {self.lattice_vectors[0,2]: 2.3f}\n {self.lattice_vectors[1,0]: 2.3f}' f' {self.lattice_vectors[1,1]: 2.3f} {self.lattice_vectors[1,2]: 2.3f}\n' f' {self.lattice_vectors[2,0]: 2.3f} {self.lattice_vectors[2,1]: 2.3f}' f' {self.lattice_vectors[2,2]: 2.3f}\n') print_stdout(message, msg_type='NOTE') message = ( f'reciprocal space lattice from input file (1/Angstrom):\n' f' {self.r_lattice_vectors[0,0]: 2.3f} {self.r_lattice_vectors[0,1]: 2.3f}' f' {self.r_lattice_vectors[0,2]: 2.3f}\n {self.r_lattice_vectors[1,0]: 2.3f}' f' {self.r_lattice_vectors[1,1]: 2.3f} {self.r_lattice_vectors[1,2]: 2.3f}\n' f' {self.r_lattice_vectors[2,0]: 2.3f} {self.r_lattice_vectors[2,1]: 2.3f}' f' {self.r_lattice_vectors[2,2]: 2.3f}') print_stdout(message)
def generate_Qpoints(self, invars, lattice): """ generate Q points, either a 2d scan or read list from file. """ # generate list of Q in rlu if invars.Qpoints_file == False: self._Qpoints_from_slice(invars) # 2d scan from input file else: self._Qpoints_from_list(invars) # read Q points from file # print list of Q in rlu message = f'Qpoints: {self.total_Qsteps}' print_stdout(message, msg_type='Brillouin zone path') Q_count = 0 for Q in range(self.total_Qsteps): message = ( f'{Q+1}\t{self.total_reduced_Q[Q,0]: 2.3f} {self.total_reduced_Q[Q,1]: 2.3f} ' f'{self.total_reduced_Q[Q,2]: 2.3f} r.l.u.') print_stdout(message) if Q_count == 49: # break if >= 50 Q points message = ( '...............................\n number of Qpoints is >= 50.' '\n suppressing output.') print_stdout(message) break Q_count = Q_count + 1 self._convert_Q_to_1_over_Angstrom( lattice) # convert the Q from rlu to 1/A
def _loop_over_blocks(self, invars, Qpoints, lattice, traj_file): """ contains outer loop over blocks info about scattering lengths: there should be 1 length per TYPE, in order of types. e.g. for 4 types = 1,2,3,4 there should be for lenghts atom 1 : length 1, atom 2 : lenght 2, etc... i am also assuming that dump_modify sort id was used so that the order of atoms is the same for each step. this can be changed easily if not the case using the atom_ids variable, but that will slow down the calc a little. the b_array variable has shape [num_steps, num_atoms] to vectorize calculating the neutron weighted density-density correlation fn """ for block_index in invars.blocks: # loop over blocks to 'ensemble' average self.block_index = block_index # print progress and start timer if self.rank == 0: start_time = timer() message = f'now on block {self.counter} out of {self.num_blocks}' print_stdout(message, msg_type='NOTE') # get the positions from file traj_file.parse_trajectory(invars, self) # check that the number of b's defined in file are consistent with traj if self.rank == 0: if np.unique(self.atom_ids[0, :]).shape[0] != invars.num_types: message = 'number of types in input file doesnt match simulation' raise PSF_exception(message) # set up array of scattering lengths. for aa in range(invars.num_atoms): self.b_array[0, aa] = invars.b[self.atom_ids[0, aa] - 1] self.b_array = np.tile(self.b_array[0, :].reshape( 1, invars.num_atoms), reps=[self.block_steps, 1]) # box lenghts read from traj file a = self.box_lengths[0] / invars.supercell[0] b = self.box_lengths[1] / invars.supercell[1] c = self.box_lengths[2] / invars.supercell[2] # print box lengths read from traj file to compare to input file if self.rank == 0: message = f'cell lengths from hdf5 file: {a:2.3f} {b:2.3f} {c:2.3f} Angstrom' print_stdout(message, msg_type='NOTE') # recall, only ortho lattice vectors used (for now) if invars.recalculate_cell_lengths: # optionally recalculates from avg in MD file lattice.lattice_vectors = np.array([[a, 0, 0], [0, b, 0], [0, 0, c]]) lattice.recompute_lattice( self.rank) # recompute reciprocal lattice Qpoints.reconvert_Q_points( lattice) # convert Q to 1/A in new basis # do the loop over Q points if self.rank == 0: message = ( 'printing progess for rank 0, which has >= the number of Q on other procs.\n' ' -- now entering loop over Q -- ') print_stdout(message, msg_type='NOTE') for qq in range(Qpoints.Qsteps): if self.rank == 0: message = f' now on Q-point {qq+1} out of {Qpoints.Qsteps}' print_stdout(message) Q = Qpoints.Qpoints[qq, :].reshape( (1, 3)) # these are the ones in 1/Angstrom exp_iQr = np.tile( Q, reps=[self.block_steps, invars.num_atoms, 1]) * self.pos exp_iQr = np.exp(1j * exp_iQr.sum(axis=2)) * self.b_array self.sqw[:, qq] = self.sqw[:, qq] + np.abs( fft(exp_iQr.sum(axis=1)))**2 # print timing to the log file if self.rank == 0: end_time = timer() elapsed_time = (end_time - start_time) / 60 # minutes message = f' elapsed time for this block: {elapsed_time:2.3f} minutes' print_stdout(message, msg_type='TIMING') message = f' time per Q-point: {elapsed_time*60/Qpoints.Qsteps:2.3f} seconds' print_stdout(message) # optionally save progress if invars.save_progress: if self.counter != self.num_blocks: f_name = invars.outfile_prefix + f'_P{self.rank}_B{block_index}.hdf5' mod_io.save_sqw(invars, Qpoints.reduced_Q, self.meV, self.sqw / self.counter, f_name) self.counter = self.counter + 1 # update the counter """
def _check_variables(self): """ where applicable, do some checks on input variables and exit if need be """ # check that the lattice vectors make sense try: self.lattice_vectors = np.array(self.lattice_vectors).reshape( (3, 3)) except: message = 'lattice vectors seem wrong. should be a list of 9 floats with no commas' raise PSF_exception(message) # check that lattice vectors are ortho # the issue is that positions etc. are in cartesian coords with ortho boxes. different lattice # vectors should work, but i haven't tested it yet. it will be necessary to convert Q in 1/A # to cartesian coordinates so that the vectorized multiplication done in mod_sqw._loop_over_blocks # works. if (self.lattice_vectors[0, 1] != 0 or self.lattice_vectors[0, 2] != 0 or self.lattice_vectors[1, 0] != 0 or self.lattice_vectors[1, 2] != 0 or self.lattice_vectors[2, 0] != 0 or self.lattice_vectors[2, 1] != 0): message = 'only ortho. lattice vectors are currently supported. see comments in mod_invars' raise PSF_exception(message) # print the traj file message = f'reading trajectories from file \'{self.traj_file}\'' print_stdout(message, msg_type='NOTE') # check for user defined scattering lenghts (only for ins) if self.ins_xlengths != False: self.num_types = len(self.ins_xlengths) message = 'using user specified scattering lengths (only works for ins, ignored for xray)' print_stdout(message, msg_type='NOTE') else: self.num_types = len(self.types) message = 'using scattering lengths from mod_xlengts' print_stdout(message, msg_type='NOTE') # check experiment type if self.exp_type not in ['xray', 'ins']: message = 'experiment type should be either \'xray\' or \'ins\'' raise PSF_exception(message) else: message = f'the experiment type is \'{self.exp_type}\'' print_stdout(message, msg_type='NOTE') # check that Q paths opts make sense if len(self.Qmin) != 3: message = f'variable Qmin should be a list of 3 floats' raise PSF_exception(message) if len(self.Qmax) != 3: message = f'variable Qmax should be a list of 3 floats' raise PSF_exception(message) # check that the requested blocks make sense if max(self.blocks) >= self.num_blocks or len( self.blocks) > self.num_blocks: message = f'variable blocks should be a list of the blocks to calculate' raise PSF_exception(message) # if the output dir. doesnt exist, create it if not os.path.exists(self.output_dir): message = f'creating directory \'{self.output_dir}\'' print_stdout(message, msg_type='NOTE') os.mkdir(self.output_dir) # check that atleast one of compute_* is not False if not self.compute_sqw and not self.compute_timeavg and not self.compute_bragg: message = ( 'there is nothing to do! set atleast one of compute_sqw, \n compute_timeavg,' ' or compute_bragg to 1 in the input file') raise PSF_exception(message) # check if traj file opens if not os.path.exists(self.traj_file): message = f'file \'{self.traj_file}\' not found' raise PSF_exception(message)
def __init__(self,invars,Qpoints,rank): """ setup freq. grid from MD params and initialize variables to hold SQW i tried padding the FFT w 0's, but spectral leakage was waaayyy worse """ self.rank = rank # the calling mpi rank # effective number of steps, i.e. num in each block that is read and computed self.block_steps = (invars.total_steps//invars.stride)//invars.num_blocks self.pos = np.zeros((self.block_steps,invars.num_atoms,3)) # time-steps, atoms, xyz self.atom_types = np.zeros((self.block_steps,invars.num_atoms)).astype(int) # see mod_io self.box_lengths = [0,0,0] # read from traj file self.num_blocks = len(invars.blocks) # see mod_invars self.counter = 1 # tools to fill xlengths array for ins/xray scattering self.xlengths = np.zeros((self.block_steps,invars.num_atoms)) # this is an array self.xlengths_tools = mod_xlengths.scattering_lengths(invars.num_types) # this is a class object # "rescale" the intensity to make plotting easier. self.common_rescale = 1e6 # only create sqw array if requested. saves time to not do this if only bragg and/or timeavg if invars.compute_sqw: if self.rank == 0: message = 'calculating the dynamical intensity' print_stdout(message,msg_type='NOTE') # setup frequency grid for fft self.num_freq = self.block_steps self.dt_eff = invars.dt*invars.stride*1e-3 # in units of ps, input is fs self.meV = fftfreq(self.num_freq,self.dt_eff)*4.13567 self.df = self.meV[1]-self.meV[0] self.max_freq = self.meV.max() self.sqw_norm = self.num_freq # change this to change how time FT is normalized # print the energy resolution, max, etc if self.rank == 0: message = (f'max freq: {self.max_freq:2.3f} meV\n' f' number of freq.: {self.num_freq}\n' f' resolution: {self.df:2.3e} meV\n') print_stdout(message,msg_type='frequency Grid') message = ('the time Fourier transform is normalized so that the MEAN\n' ' of the dynamical intensity over all energies is equal to the\n' ' total intensity averaged over time') print_stdout(message,msg_type='NORMALIZATION') # the sqw array self.sqw = np.zeros((self.num_freq,Qpoints.Qsteps)) # only create bragg array if requested if invars.compute_bragg: if self.rank == 0: message = 'calculating the Bragg intensity' print_stdout(message,msg_type='NOTE') self.bragg = np.zeros(Qpoints.Qsteps) # only create timeavg array if requested if invars.compute_timeavg: if self.rank == 0: message = 'calculating the time-avgeraged intensity' print_stdout(message,msg_type='NOTE') self.timeavg = np.zeros(Qpoints.Qsteps)
def _loop_over_blocks(self,invars,Qpoints,lattice,traj_file): """ contains outer loop over blocks info about scattering lengths: there should be 1 length per TYPE, in order of types. e.g. for 4 types = 1,2,3,4 there should be for lengths atom 1 : length 1, atom 2 : lenght 2, etc... i am also assuming that dump_modify sort id was used so that the order of atoms is the same for each step. this can be changed easily if not the case using the atom_types variable, but that will slow down the calc a little. the b_array variable has shape [num_steps, num_atoms] to vectorize calculating the neutron weighted density-density correlation fn """ for block_index in invars.blocks: # loop over blocks to 'ensemble' average # used below self.block_index = block_index # print progress and start timer if self.rank == 0: start_time = timer() message = '\n............................................' print_stdout(message) message = f' now on block {self.counter} out of {self.num_blocks}' print_stdout(message,msg_type='NOTE') # get the positions from file traj_file.parse_trajectory(invars,self) # check that the number of b's defined in input file are consistent with traj file if self.rank == 0: if np.unique(self.atom_types[0,:]).shape[0] != invars.num_types: message = 'number of types in input file doesnt match simulation' raise PSF_exception(message) # look up ins scattering lengths OR parameters to compute xray form factors. self.xlengths_tools.map_types_to_data(invars,self) # box lengths read from traj file a = self.box_lengths[0]/invars.supercell[0] b = self.box_lengths[1]/invars.supercell[1] c = self.box_lengths[2]/invars.supercell[2] # print box lengths read from traj file to compare to input file if self.rank == 0: message = f'cell lengths from hdf5 file: {a:2.3f} {b:2.3f} {c:2.3f} Angstrom' print_stdout(message,msg_type='NOTE') # recall, only ortho lattice vectors used (for now) if invars.recalculate_cell_lengths: # optionally recalculates from avg in MD file lattice.lattice_vectors = np.array([[a,0,0],[0,b,0],[0,0,c]]) lattice.recompute_lattice(self.rank) # recompute reciprocal lattice Qpoints.reconvert_Q_points(lattice) # convert Q to 1/A in new basis # --------------------- enter the loop over Qpoints ----------------------------------- if self.rank == 0: Q_start_time = timer() # track time per Q not including the read/write time message = ('printing progess for rank 0, which has >= the number of Q on other procs.\n' ' -- now entering loop over Q -- ') print_stdout(message,msg_type='NOTE') for qq in range(Qpoints.Qsteps): if self.rank == 0: message = f' now on Q-point {qq+1} out of {Qpoints.Qsteps}' print_stdout(message) # the Qpoint to do Q = Qpoints.Qpoints[qq,:].reshape((1,3)) # 1/Angstrom self.Q_norm = np.sqrt(Q[0,0]**2+Q[0,1]**2+Q[0,2]**2) # if xray, need to compute f(|Q|), which are placed in self.xlengths if invars.exp_type == 'xray': self.xlengths_tools.compute_xray_form_fact(self,invars) # space FT by vectorized Q.r dot products and sum over atoms. (tile prepends new axes) exp_iQr = np.tile(Q,reps=[self.block_steps,invars.num_atoms,1])*self.pos # Q.r exp_iQr = np.exp(1j*exp_iQr.sum(axis=2))*self.xlengths # sum over x, y, z exp_iQr = exp_iQr.sum(axis=1) # sum over atoms # compute bragg intensity = |<rho(Q,t)>|**2 if invars.compute_bragg: self.bragg[qq] = self.bragg[qq]+np.abs((exp_iQr).mean())**2/self.common_rescale # compute timeavg intensity = <|rho(Q,t)|**2> if invars.compute_timeavg: self.timeavg[qq] = self.timeavg[qq]+(np.abs(exp_iQr)**2).mean()/self.common_rescale # compute dynamical intensity = |rho(Q,w)|**2 if invars.compute_sqw: self.sqw[:,qq] = (self.sqw[:,qq]+np.abs(fft(exp_iQr))**2/ self.sqw_norm/self.common_rescale) # ------------------------------------------------------------------------------------- # optionally save progress if invars.save_progress: if self.counter != self.num_blocks: if invars.compute_sqw: f_name = invars.outfile_prefix+f'_SQW_P{self.rank}_B{block_index}.hdf5' mod_io.save_sqw(invars,Qpoints.reduced_Q,self.meV,self.sqw/self.counter,f_name) if invars.compute_bragg: f_name = invars.outfile_prefix+f'_BRAGG_P{self.rank}_B{block_index}.hdf5' mod_io.save_bragg(invars,Qpoints.reduced_Q,self.bragg,f_name) if invars.compute_timeavg: f_name = invars.outfile_prefix+f'_TIMEAVG_P{self.rank}_B{block_index}.hdf5' mod_io.save_timeavg(invars,Qpoints.reduced_Q,self.timeavg,f_name) # print timing to the log file if self.rank == 0: end_time = timer() elapsed_time = end_time-start_time Q_time = end_time-Q_start_time io_time = elapsed_time-Q_time Q_time = Q_time/Qpoints.Qsteps # avg over all Q message = f' avg time per Q-point: {Q_time:2.3f} seconds' print_stdout(message,msg_type='TIMING') message = f' total io time: {io_time:2.3f} seconds' print_stdout(message) message = (f' total time for this block: {elapsed_time:2.3f} seconds' f' ({elapsed_time/60:2.3f} minutes)') print_stdout(message) self.counter = self.counter+1 # update the counter
def _loop_over_Q(self,proc,invars,Qpoints): """ this target is passed to each mp Process. it loops over the Q points assigned to each process and puts the results into the Queue. a background process collects all of this and assembles to total S(Q,w), S(Q), etc. arrays. could be easily modified to work in serial. """ # get the inds of which Q to do Q_inds = Qpoints.Q_on_procs[proc] # get how many there are num_Q = len(Q_inds) # return dummy (None) sqw array if not requested if invars.compute_sqw: sqw_pp = np.zeros((self.num_freq,num_Q)) else: sqw_pp = None # return dummy (None) bragg array if not requested if invars.compute_bragg: bragg_pp = np.zeros(num_Q) else: bragg_pp = None # return dummy (None) timeavg array if not requested if invars.compute_timeavg: timeavg_pp = np.zeros(num_Q) else: timeavg_pp = None # loop over the Q points for qq in range(num_Q): # print status if on proc 0 if proc == 0: message = f' now on Q-point {qq+1} out of {num_Q}' print_stdout(message) # the Qpoint to do Q_ind = Q_inds[qq] Q = Qpoints.total_Qpoints[Q_ind,:].reshape((1,3)) # 1/Angstrom self.Q_norm = np.sqrt(Q[0,0]**2+Q[0,1]**2+Q[0,2]**2) # |Q| # if xray, need to compute f(|Q|), which are placed in self.xlengths if invars.exp_type == 'xray': self.xlengths_tools.compute_xray_form_fact(self,invars) # space FT by vectorized Q.r dot products and sum over atoms. (tile prepends new axes) exp_iQr = np.tile(Q,reps=[self.block_steps,invars.num_atoms,1])*self.pos # Q.r exp_iQr = np.exp(1j*exp_iQr.sum(axis=2))*self.xlengths # sum over x, y, z exp_iQr = exp_iQr.sum(axis=1) # sum over atoms # compute bragg intensity = |<rho(Q,t)>|**2 if invars.compute_bragg: bragg_pp[qq] = np.abs((exp_iQr).mean())**2/self.common_rescale # compute timeavg intensity = <|rho(Q,t)|**2> if invars.compute_timeavg: timeavg_pp[qq] = (np.abs(exp_iQr)**2).mean()/self.common_rescale # compute dynamical intensity = |rho(Q,w)|**2 if invars.compute_sqw: sqw_pp[:,qq] = np.abs(fft(exp_iQr))**2/self.sqw_norm/self.common_rescale # put this stuff into the 'Queue' so that some other process can put it into the main arrays self.mp_queue.put([sqw_pp,bragg_pp,timeavg_pp,proc])
def _loop_over_blocks(self,invars,Qpoints,lattice,traj_file): """ contains outer loop over blocks info about scattering lengths: there should be 1 length per TYPE, in order of types. e.g. for 4 types = 1,2,3,4 there should be for lengths atom 1 : length 1, atom 2 : lenght 2, etc... i am also assuming that dump_modify sort id was used so that the order of atoms is the same for each step. this can be changed easily if not the case using the atom_types variable, but that will slow down the calc a little. the b_array variable has shape [num_steps, num_atoms] to vectorize calculating the neutron weighted density-density correlation fn """ for block_index in invars.blocks: # loop over blocks to 'ensemble' average # used below self.block_index = block_index # print progress and start timer start_time = timer() message = '\n............................................' print_stdout(message) message = f' now on block {self.counter} out of {self.num_blocks}' print_stdout(message,msg_type='NOTE') # get the positions from file traj_file.parse_trajectory(invars,self) # check that the number of b's defined in input file are consistent with traj file if np.unique(self.atom_types[0,:]).shape[0] != invars.num_types: message = 'number of types in input file doesnt match simulation' raise PSF_exception(message) # look up ins scattering lengths OR parameters to compute xray form factors. self.xlengths_tools.map_types_to_data(invars,self) # box lengths read from traj file a = self.box_lengths[0]/invars.supercell[0] b = self.box_lengths[1]/invars.supercell[1] c = self.box_lengths[2]/invars.supercell[2] # print box lengths read from traj file to compare to input file message = f'cell lengths from hdf5 file: {a:2.3f} {b:2.3f} {c:2.3f} Angstrom' print_stdout(message,msg_type='NOTE') # recall, only ortho lattice vectors used (for now) if invars.recalculate_cell_lengths: # optionally recalculates from avg in MD file lattice.lattice_vectors = np.array([[a,0,0],[0,b,0],[0,0,c]]) lattice.recompute_lattice() # recompute reciprocal lattice Qpoints.reconvert_Q_points(lattice) # convert Q to 1/A in new basis Q_start_time = timer() # track time per Q not including the read/write time message = ('printing progess for process 0, which has >= the number of Q on other' \ ' processes.\n -- now entering loop over Q -- ') print_stdout(message,msg_type='NOTE') # ----------------------------------------------------------------------------------------- # ------------- multiprocessing part. ------------- # ----------------------------------------------------------------------------------------- # a Queue to hold the retured SQW data self.mp_queue = mp.Queue() # a container to hold the processes procs = [] # loop over processes, setting up the loop over Q on each. for pp in range(invars.num_processes): procs.append(mp.Process(target=self._loop_over_Q, args=(pp,invars,Qpoints))) # now start running the function on each process for proc in procs: proc.start() # note, doing it this way with the queue 'blocks' until the next processes adds to queue if # it is empty. I dont know if this will freeze the whole calculation or just the background # proc that is running the queue. anyway, everything with the queue has to be done before # joining the procs or the data will corrupt/crash the program. # get the stuff calculated on each proc for pp in range(invars.num_processes): sqw_pp, bragg_pp, timeavg_pp, proc = self.mp_queue.get() Q_inds = Qpoints.Q_on_procs[proc] # now put it into main arrays if requested if invars.compute_sqw: self.sqw[:,Q_inds] = self.sqw[:,Q_inds]+sqw_pp if invars.compute_bragg: self.bragg[Q_inds] = self.bragg[Q_inds]+bragg_pp if invars.compute_timeavg: self.timeavg[Q_inds] = self.timeavg[Q_inds]+timeavg_pp # now close the queue and rejoin its proc self.mp_queue.close() self.mp_queue.join_thread() # wait here for all to finish before moving on to the next block for proc in procs: proc.join() # ----------------------------------------------------------------------------------------- # ------------ end of multiprocessing part ---------------- # ----------------------------------------------------------------------------------------- # optionally save progress if invars.save_progress: if self.counter != self.num_blocks: # dont write if this is the last block if invars.compute_sqw: f_name = invars.outfile_prefix+f'_SQW_B{block_index}.hdf5' mod_io.save_sqw(invars,Qpoints.reduced_Q,self.meV,self.sqw/self.counter,f_name) if invars.compute_bragg: f_name = invars.outfile_prefix+f'_BRAGG_B{block_index}.hdf5' mod_io.save_bragg(invars,Qpoints.reduced_Q,self.bragg/self.counter,f_name) if invars.compute_timeavg: f_name = invars.outfile_prefix+f'_TIMEAVG_B{block_index}.hdf5' mod_io.save_timeavg(invars,Qpoints.reduced_Q,self.timeavg/self.counter,f_name) # print timing to screen end_time = timer() elapsed_time = end_time-start_time Q_time = end_time-Q_start_time io_time = elapsed_time-Q_time # time per Qpoint Q_time = Q_time/len(Qpoints.Q_on_procs[0]) # avg over all Q message = f' avg time per Q-point: {Q_time:2.3f} seconds' print_stdout(message,msg_type='TIMING') # time spent in i/o message = f' total io time: {io_time:2.3f} seconds' print_stdout(message) # total time for in this method message = (f' total time for this block: {elapsed_time:2.3f} seconds' f' ({elapsed_time/60:2.3f} minutes)') print_stdout(message) # update the block counter self.counter = self.counter+1
def parse_input(self, input_file): """ read the input file """ # test that the input file exists/isnt broken try: with open(input_file, 'r') as inp: self.input_txt = inp.readlines() except: message = f'input file \'{input_file}\' not found' raise PSF_exception(message) # check the key_words in the file, remove empty lines and comments self._check_file() # get the variables from file self.traj_file = self._parse_str('traj_file', self.traj_file) self.outfile_prefix = self._parse_str('outfile_prefix', self.outfile_prefix) self.output_dir = self._parse_str('output_dir', self.output_dir) self.save_progress = self._parse_bool('save_progress', self.save_progress) self.parse_custom = self._parse_bool('parse_custom', self.parse_custom) self.dt = self._parse_float('dt', self.dt) self.stride = self._parse_int('stride', self.stride) self.total_steps = self._parse_int('total_steps', self.total_steps) self.num_atoms = self._parse_int('num_atoms', self.num_atoms) self.supercell = self._parse_int_list('supercell', self.supercell) self.lattice_vectors = self._parse_float_list('lattice_vectors', self.lattice_vectors) self.unwrap_pos = self._parse_bool('unwrap_pos', self.unwrap_pos) self.recalculate_cell_lengths = self._parse_bool( 'recalculate_cell_lengths', self.recalculate_cell_lengths) self.b = self._parse_float_list('b', self.b) self.num_types = len(self.b) self.Qpoints_file = self._parse_str('Qpoints_file', self.Qpoints_file) self.Qmin = self._parse_float_list('Qmin', self.Qmin) self.Qmax = self._parse_float_list('Qmax', self.Qmax) self.total_Qsteps = self._parse_int('total_Qsteps', self.total_Qsteps) self.num_blocks = self._parse_int('num_blocks', self.num_blocks) self.blocks = list(range(self.num_blocks)) self.blocks = self._parse_int_list('blocks', self.blocks) # check that the lattice vectors make sense try: self.lattice_vectors = np.array(self.lattice_vectors).reshape( (3, 3)) except: message = 'lattice vectors seem wrong. should be a list of 9 floats with no commas' raise PSF_exception(message) # check that lattice vectors are ortho # the issue is that positions etc. are in cartesian coords with ortho boxes. different lattice # vectors should work, but i haven't tested it yet. it will be necessary to convert Q in 1/A # to cartesian coordinates so that the vectorized multiplication done in mod_sqw._loop_over_blocks # works. if (self.lattice_vectors[0, 1] != 0 or self.lattice_vectors[0, 2] != 0 or self.lattice_vectors[1, 0] != 0 or self.lattice_vectors[1, 2] != 0 or self.lattice_vectors[2, 0] != 0 or self.lattice_vectors[2, 1] != 0): message = 'only ortho. lattice vectors are currently supported. see comments in mod_invars' raise PSF_exception(message) # print the traj file message = f'reading trajectories from file \'{self.traj_file}\'' print_stdout(message, msg_type='NOTE') # print the scattering lengths to file message = f' atom-type: 0 b: {self.b[0]: 2.4f}\n' for bb in range(1, self.num_types): message = message + f' atom-type: {bb:2g} b: {self.b[bb]: 2.4f}\n' print_stdout(message, msg_type='scattering lengths (b) in femtometers') # now convert b to Angstrom. 'intensities' are of order 1 this way. with FM, theyre # of order 1e16. this is purely for convenience and doesnt change the physics, since the # intensities are arbitrary units anyway for bb in range(self.num_types): self.b[bb] = self.b[bb] * 1e-5 # check that Q paths opts make sense if len(self.Qmin) != 3: message = f'variable Qmin should be a list of 3 floats' raise PSF_exception(message) if len(self.Qmax) != 3: message = f'variable Qmax should be a list of 3 floats' raise PSF_exception(message) # check that the requested blocks make sense if max(self.blocks) >= self.num_blocks or len( self.blocks) > self.num_blocks: message = f'variable blocks should be a list of the blocks to calculate' raise PSF_exception(message) # if the output dir. doesnt exist, create it if not os.path.exists(self.output_dir): message = f'creating directory \'{self.output_dir}\'' print_stdout(message, msg_type='NOTE') os.mkdir(self.output_dir)