omdShrink

#!@Python3_EXECUTABLE@
"""OpenMD File Trimmer
 
Skips every n frames of an OpenMD dump file and loads it into new dump file
 
Usage: omdShrink
 
Options:
  -h, --help              show this help
  -m, --meta-data=...     use specified OpenMD (.dump) file
  -o, --output-file=...   use specified output file
  -s, --split-frame=...   split every n frames
 
 
Example:
   omdShrink -s 10 -m long.dump  -o short.dump
 
"""
 
__copyright__ = "Copyright (c) 2004-present The University of Notre Dame. All Rights Reserved."
__license__ = "OpenMD"
 
import sys
import getopt
import string
import os
 
_haveMDFileName = 0
_haveOutFileName = 0
_haveSplitFrame = 0
 
def usage():
    print(__doc__)
 
def shrinkFile(mdFileName, outputFileName, skipFrameNumber):
    mdFile = open(mdFileName, 'r')
    outFile = open(outputFileName, 'w')
 
    metaDataFrame = []
    haveMetaData = 0
    haveMetaDataEnd = 0
    # Find OpenMD version info first
    line = mdFile.readline()
    while True:
        if '<OOPSE version=' in line or '<OpenMD version=' in line:
            OpenMDversion = line
            break
        line = mdFile.readline()
 
        # Rewind file and find start of MetaData block
 
    mdFile.seek(0)
    breakLoop = False
    line = mdFile.readline()
 
    while True:
        if '<MetaData>' in line:
            while 2:
                metaDataFrame.append(line)
                line = mdFile.readline()
                if '</MetaData>' in line:
                    metaDataFrame.append(line)
                    breakLoop = True
                    break
        if breakLoop:
            break
        line = mdFile.readline()
 
    outFile.write(OpenMDversion)
    for metaline in metaDataFrame:
        if 'sampleTime' in metaline:
            metasplit = metaline.split()
            sampleTime = float(metasplit[2].strip(';'))
            newSampleTime = sampleTime * int(skipFrameNumber)
            outFile.write('sampleTime = %10d;\n' % (newSampleTime))
        else:
            outFile.write(metaline)
 
    mdFile.seek(0)
    framePos = 0
    foundFrames = 0;
    line = mdFile.readline()
    while True:
        if '<Snapshot>' in line:
            division, remainder=divmod(foundFrames, int(skipFrameNumber))
            foundFrames += 1
            if (remainder == 0):
                while not '</Snapshot>' in line:
                    outFile.write(line)
                    line = mdFile.readline()
                outFile.write(line)
        line = mdFile.readline()
        if not line: break
 
    outFile.write("</OpenMD>")
    outFile.close()
    mdFile.close()
 
 
def main(argv):
    try:
        opts, args = getopt.getopt(argv, "hm:o:s:", ["help", "meta-data=", "output-file=", "split-frame", ""])
    except getopt.GetoptError:
        usage()
        sys.exit(2)
    for opt, arg in opts:
        if opt in ("-h", "--help"):
            usage()
            sys.exit()
        elif opt in ("-m", "--meta-data"):
            mdFileName = arg
            global _haveMDFileName
            _haveMDFileName = 1
        elif opt in ("-o", "--output-file"):
            outputFileName = arg
            global _haveOutFileName
            _haveOutFileName = 1
        elif opt in ("-s", "--split-frame"):
            skipFrameNumber = arg
            global _haveSplitFrame
            _haveSplitFrame = 1
    if (_haveMDFileName != 1):
        usage()
        print("No OpenMD input file was specified")
        sys.exit()
    if (_haveOutFileName != 1):
        usage()
        print("No output file was specified")
        sys.exit()
    if (_haveSplitFrame != 1):
        usage()
        print("Frames to split was not specified")
        sys.exit()
    shrinkFile(mdFileName, outputFileName, skipFrameNumber);
 
if __name__ == "__main__":
    if len(sys.argv) == 1:
        usage()
        sys.exit()
    main(sys.argv[1:])