Files
python-pygmx/pygmx/__init__.py
2017-08-18 13:22:58 +02:00

38 lines
1.2 KiB
Python

"""
Python wrapper for the gromacs library for their file formats.
Currently this supports xtc and tpr files within the `open` function.
Trajectories in trr format may be read with `.gromacs.reader.TRRReader`, which is experimental.
"""
import os
from .tpxio import TPXReader
from .xtcio import XTCReader, read_xtcframe
from .enxio import EDRFile
from .errors import FileTypeError
from .gromacs.reader import index_filename_for_xtc
from .gromacs.xtcindex import index_xtcfile
FILE_EXTENSIONS = {
'xtc': XTCReader,
'tpr': TPXReader,
'edr': EDRFile,
}
def open(filename, ignore_index_timestamps=False):
"""Open a supported gromacs file. Currently supported file formats: tpr, xtc."""
ext = filename.split('.')[-1]
if ext in FILE_EXTENSIONS:
if ext in ['xtc']:
indexfile = index_filename_for_xtc(filename)
if not os.path.exists(indexfile):
print('Generating Index for xtc file. This may take a while...')
index_xtcfile(filename)
return FILE_EXTENSIONS[ext](filename, indexfile, ignore_timestamps=ignore_index_timestamps)
else:
return FILE_EXTENSIONS[ext](filename)
else:
raise FileTypeError('Filetype {} not supported by pygmx.'.format(ext))