def main(fns_in, fn_out, begin, end, stride):
verbosity.level = "low"
print('Multiplexing {:d} beads into one trajectory.'.format(len(fns_in)))
print()
print('input file names:')
for fn in fns_in:
print(fn)
print()
print('output file name:', fn_out)
print()
# Open input trajectory iterators.
trjs_in = [iter_file_name(fn) for fn in fns_in]
# Open output file.
f_out = open_backup(fn_out, 'w')
mode_out = os.path.splitext(fn_out)[1]
# Loop over all frames.
i_frame = 0
i_frame_saved = 0
while True:
# Check the endpoint index, exit if we're done.
if (end > -1) and (i_frame >= end):
break
# Should we save output from this frame?
do_output = (i_frame >= begin) and ((i_frame % stride) == 0)
try:
# Get the frames from all trajectories...
for trj in trjs_in:
frame = trj.next()
# ... and possibly save them in the output trajectory.
if do_output:
print_file(mode_out, frame['atoms'], frame['cell'], f_out)
if do_output:
i_frame_saved += 1
except StopIteration:
# Stop when any of the trajectories runs out of frames.
break
# Count frames and print information on progress.
i_frame += 1
if i_frame % 100 == 0:
print('\rframe {:d}'.format(i_frame), end='')
sys.stdout.flush()
f_out.close()
print()
print()
print('Loaded {:d} frames.'.format(i_frame))
print('Saved {:d} frames.'.format(i_frame_saved))
def contract_trajectory(fns_in, fn_out_template, n_new):
n = len(fns_in)
# Generate output file names.
if n_new == 1:
fns_out = [fn_out_template]
else:
fns_out = [fn_out_template.format(i) for i in range(n_new)]
print "Contracting {:d} beads to {:d} beads.".format(n, n_new)
print
print "input file names:"
for fn in fns_in:
print fn
print
print "output file names:"
for fn in fns_out:
print fn
print
# Open input trajectory iterators.
trjs_in = [iter_file_name(fn) for fn in fns_in]
# Open output files.
fs_out = [open_backup(fn, "w") for fn in fns_out]
mode_out = os.path.splitext(fn_out_template)[1]
# prepare ring polymer rescaler
rescale = nm_rescale(n, n_new)
# Loop over all frames.
i_frame = 0
while True:
try:
# Get the frames for all beads.
frames = [trj.next() for trj in trjs_in]
except StopIteration:
# Stop when any of the trajectories runs out of frames.
break
# Consistency check.
h = frames[0]["cell"].h
natoms = frames[0]["atoms"].natoms
for i in range(n):
# Check that all the cells are the same.
if (frames[i]["cell"].h != h).any():
msg = "Cell for beads {:d} and {:d} differ in frame {:d}."
raise ValueError(msg.format(0, i, i_frame))
# Check that the numbers of atoms are the same.
if frames[i]["atoms"].natoms != natoms:
msg = "Different numbers fo atoms for beads {:d} and {:d} in frame {:d}."
raise ValueError(msg.format(0, i, i_frame))
# Reuse the first frame for output.
cell = frames[0]["cell"]
atoms = frames[0]["atoms"]
# Compose the ring polymer.
q = np.vstack([frame["atoms"].q for frame in frames])
# Contract the coordinates to `n_new` beads.
q_c = rescale.b1tob2(q)
# Save the output data.
for i, f_out in enumerate(fs_out):
atoms.q = q_c[i, :]
print_file(mode_out, atoms, cell, f_out)
# Count frames and print information on progress.
i_frame += 1
if i_frame % 100 == 0:
print "\rframe {:d}".format(i_frame),
sys.stdout.flush()
for f_out in fs_out:
f_out.close()
print
print
print "Processed {:d} frames.".format(i_frame)