def calculate(self,invars,Qpoints,lattice,traj_file):
"""
calculate SQW from MD data
"""
# this call does all of the 'work'
self._loop_over_blocks(invars,Qpoints,lattice,traj_file)
# average over the blocks
if invars.compute_sqw:
self.sqw = self.sqw/self.num_blocks
if invars.compute_bragg:
self.bragg = self.bragg/self.num_blocks
if invars.compute_timeavg:
self.timeavg = self.timeavg/self.num_blocks
# optionally save progress.
if invars.save_progress:
if invars.compute_sqw:
f_name = invars.outfile_prefix+f'_SQW_P{self.rank}_BF.hdf5' # final file
mod_io.save_sqw(invars,Qpoints.reduced_Q,self.meV,self.sqw,f_name)
if invars.compute_bragg:
f_name = invars.outfile_prefix+f'_BRAGG_P{self.rank}_BF.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}_BF.hdf5'
mod_io.save_timeavg(invars,Qpoints.reduced_Q,self.timeavg,f_name)
# free up memory (probably is already done by garbage collection)
del self.pos, self.atom_types, self.xlengths
# gather and save bragg results
if invars.compute_bragg:
# get from all ranks
bragg_from_ranks = [sqw.bragg]
for ii in range(1, num_ranks):
bragg_ii = comm.recv(source=ii, tag=12)
bragg_from_ranks.append(bragg_ii)
# assemble into a single array
bragg_total = np.zeros(Qpoints.total_Qsteps)
mod_utils.assemble_timeavg(bragg_from_ranks, bragg_total)
# write to final file
f_name = invars.outfile_prefix + f'_BRAGG_FINAL.hdf5'
mod_io.save_bragg(invars, Qpoints.total_reduced_Q, bragg_total,
f_name)
# ---------------------------------------------------------------------
# gather and save timeavg results
if invars.compute_timeavg:
# get from all ranks
timeavg_from_ranks = [sqw.timeavg]
for ii in range(1, num_ranks):
timeavg_ii = comm.recv(source=ii, tag=13)
timeavg_from_ranks.append(timeavg_ii)
# assemble into a single array
timeavg_total = np.zeros(Qpoints.total_Qsteps)
mod_utils.assemble_timeavg(timeavg_from_ranks, timeavg_total)
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_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
# save the SQW results
if invars.compute_sqw:
# save total SQW to final hdf5 file
f_name = invars.outfile_prefix + f'_SQW_FINAL.hdf5'
mod_io.save_sqw(invars, Qpoints.total_reduced_Q, sqw.meV, sqw.sqw,
f_name)
# ----------------------------------------------------------------------------
# save bragg results
if invars.compute_bragg:
# write to final file
f_name = invars.outfile_prefix + f'_BRAGG_FINAL.hdf5'
mod_io.save_bragg(invars, Qpoints.total_reduced_Q, sqw.bragg, f_name)
# ----------------------------------------------------------------------------
# save timeavg results
if invars.compute_timeavg:
# write to final file
f_name = invars.outfile_prefix + f'_TIMEAVG_FINAL.hdf5'
mod_io.save_timeavg(invars, Qpoints.total_reduced_Q, sqw.timeavg,
f_name)
# ----------------------------------------------------------------------------
# calculate and print elapsed time to do this file
inner_end = timer()