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 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 __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 __init__(self, invars, Qpoints):
"""
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
"""
# 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
# "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:
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
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.total_Qsteps))
# only create bragg array if requested
if invars.compute_bragg:
message = 'calculating the Bragg intensity'
print_stdout(message, msg_type='NOTE')
self.bragg = np.zeros(Qpoints.total_Qsteps)
# only create timeavg array if requested
if invars.compute_timeavg:
message = 'calculating the time-avgeraged intensity'
print_stdout(message, msg_type='NOTE')
self.timeavg = np.zeros(Qpoints.total_Qsteps)
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
# TS: 10.13.2021: this line of code doesnt seemed to be used anywhere
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):
# i think this is FIFO
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 _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. \n' \
' contact the author at [email protected] if you need this feature'
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 that the requested number of processes makes sense
if self.num_processes == None:
self.num_processes = os.cpu_count()
if self.num_processes < 1:
message = 'requested number of processes should be 1 or larger'
raise PSF_exception(message)
# 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 = 'each vertex for the Q path should have 3 coords'
raise PSF_exception(message)
self.num_Qpath = len(self.Qmin) // 3
if len(self.Qmin) != len(self.Qmax):
message = f'variable Qmin and Qmax should have same number of vertices'
raise PSF_exception(message)
if len(self.total_Qsteps) != self.num_Qpath:
message = 'number of steps should equal number of paths'
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 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
"""
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:
message = 'unwrapping positions'
print_stdout(message, msg_type='NOTE')
self._unwrap_positions(invars, sqw)