filterTrj.py

#!/usr/bin/python

## BY: MOHAMED HESHAM MOSELHY
## OUTPUT A LIST OF TRAJECTORY FILES THAT START AT starttime AND ENDS WITH endtime
## Assumes that all the trajectory files correspond to checkpoint files with the same prefix
## For example: 1.trr or 2.###trr (where ### is anything or nothing) correspond to 1.cpt and 2.cpt, respectively
## Usage: filterTrj.py starttime endtime [outputfile - optional]
## Time units for start and end times are in picoseconds
## outputfile will be overridden if it exists


import sys
from glob import glob
import subprocess
import os
from progress.bar import Bar
import fnmatch

def get_time(prefix):
	output = subprocess.check_output(['{ gmx dump -cp ' + str(prefix) + '.cpt | head -n19 | tail -n1; } 2>/dev/null'], shell=True)
	t = float(output.split('= ')[1].rstrip('\n'))
	return t

def get_trj(prefix):
	pattern = '{}.*trr'.format(prefix)
	return fnmatch.filter(dirlist, pattern)[0]

if len(sys.argv) < 3 or len(sys.argv) > 4:
	sys.stderr.write('Usage: {0} starttime endtime (ps) [outputfile - optional]\n'.format(sys.argv[0]))
	sys.exit(1)

outputfile = '/dev/null'
start_time = float(sys.argv[1])
end_time = float(sys.argv[2])
if len(sys.argv) == 4:
	outputfile = sys.argv[3]

fp = open(outputfile, 'w')

cpt_files = glob('?.cpt') + glob('[0-9]?.cpt')
cpt_files = [int(x.rstrip('*.cpt')) for x in cpt_files]
cpt_files.sort()
times = [None] * (len(cpt_files) + 1)
out = []
dirlist = os.listdir('.')

bar = Bar("Reading: ", max=cpt_files[-1])
# Iterate through each cpt file, (including the last one, hence the +1)
for i in range(1, cpt_files[-1]+1):
	t = get_time(i)
	curr = times[i] = t
	# print('\n' + str(curr))
	if curr >= end_time:
		out += [get_trj(i)]
		print("\nDone at {}.cpt".format(i))
		break
	elif start_time < curr:
		out += [get_trj(i)]
	bar.next()
bar.finish()

out = [os.path.realpath(p) for p in out]
output = ' '.join(out)

print(output)
fp.write(output + '\n')