def Run(self, output_file=None):
""" Runs cpptraj to actually create the files """
from sys import stdout as sys_stdout
from subprocess import Popen, PIPE
# Accept output_file as both a file object and string object
own_handle = False
if output_file is None:
stdout = sys_stdout
else:
try:
stdout = open(output_file, 'w')
own_handle = True
except TypeError:
stdout = output_file
# Now it's time to run the program
try:
input_string = ''
for action in self.actions:
input_string += action.strip() + '\n'
process = Popen([self.exe, self.prmtop], stdout=stdout, stdin=PIPE)
process.communicate(input_string.encode())
if process.wait():
raise TrajError('Error running %s' % self.program)
finally:
if own_handle: stdout.close()
def Query(self):
""" Finds out how many frames are in the given traj files """
from subprocess import Popen, PIPE
import re
framere = re.compile(r'Frames: (\d+)')
self.traj_sizes = []
# Determine how many frames are in each trajectory of the list
for traj in self.traj_files:
process = Popen(
[self.exe, '-p',
str(self.prmtop), '-y', traj, '-tl'],
stdin=PIPE,
stdout=PIPE)
(output, error) = process.communicate(b'')
if process.wait(): # if it quits with return code != 0
raise TrajError('%s failed when querying %s' %
(self.exe, traj))
output = output.decode()
# Now parse the output to find out how many frames are there. We are
# looking for "Coordinate processing will occur on x frames."
num_frames = framere.findall(output)
if len(num_frames) < 1:
raise TrajError('Could not find number of frames in ' + traj)
elif len(num_frames) > 1:
raise RuntimeError(
'Unexpected output from cpptraj. Has format '
'changed?')
self.traj_sizes.append(int(num_frames[0]))
# end for traj in self.traj_files
self.total_frames = sum(self.traj_sizes)
def Setup(self, startframe=1, endframe=9999999, interval=1):
""" Finds the program and clears the queue and other attributes assigned
during other function calls
"""
self.strip_solvent = False
orig_endframe = endframe
orig_startframe = startframe
# Delete the stripmask, but make its non-existence non-fatal, since it
# won't exist unless Setup is called a second time.
try:
del self.stripmask
except AttributeError:
pass
self.actions = []
# Catch stupid choices
if startframe > endframe:
raise TrajError('Cannot have startframe (%d) > endframe (%d)' %
(startframe, endframe))
if startframe < 0:
raise TrajError('Startframe (%d) < 0' % startframe)
if interval <= 0:
raise TrajError('Interval (%d) <= 0' % interval)
if startframe > self.total_frames:
raise TrajError('start frame (%d) > total frames (%d)' %
(startframe, self.total_frames))
# Bring endframe down if it's > the total number of frames
endframe = min(endframe, self.total_frames)
self.analyzed_frames = int((endframe - startframe) / interval) + 1
# If we have an interval != 1, then it's possible that the endframe that
# the user set is not the actual one that will be used. To make things
# simpler, I will adjust the endframe to what is *actually* used
endframe = startframe + interval * int(
(endframe - startframe) // interval)
# Set up start and end arrays (for each trajectory)
traj_starts = [-1 for i in range(len(self.traj_files))]
traj_ends = [0 for i in range(len(self.traj_files))]
# Now determine where each trajectory starts and ends
for i in range(len(self.traj_files)):
# skip over any trajectories that lie entirely before startframe
if startframe > self.traj_sizes[i]:
traj_starts[i] = -1 # this will tell us to skip this traj file
startframe -= self.traj_sizes[i]
continue
# Now we start at our startframe
traj_starts[i] = startframe
# Now we figure out our last frame, adjusting for interval
last_frame = startframe + interval * int(
(self.traj_sizes[i] - startframe) / interval)
traj_ends[i] = min(endframe, last_frame, self.traj_sizes[i])
# Now determine our new start frame for the next trajectory. First
# we need to know how many frames are left over at the end of the
# last frame, and subtract those from the interval to determine
# where our next trajectory should start from. Also move our endframe
# down, but only by our last_frame! (since the last couple frames in
# our trajectory file could be leftover due to interval > 1)
startframe = interval + last_frame - self.traj_sizes[i]
endframe -= last_frame
if endframe < self.traj_sizes[i]:
break # we're done now
# end for i len(traj_files)
# We now have to trajin each of the analyzed trajectories
for i in range(len(self.traj_files)):
if traj_starts[i] < 0: continue # skip -1's
self.actions.append(
'trajin %s %d %d %d' %
(self.traj_files[i], traj_starts[i], traj_ends[i], interval))
self.actions.append('noprogress') # quash the progress bar
self.processed_frames = (min(orig_endframe, self.total_frames) -
orig_startframe) / interval + 1