Source code for pyretis.inout.formats.gromacs

# -*- coding: utf-8 -*-
# Copyright (c) 2021, PyRETIS Development Team.
# Distributed under the LGPLv2.1+ License. See LICENSE for more info.
"""Module for handling GROMACS input/output.

Important methods defined here
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

read_gromacs_file (:py:func:`.read_gromacs_file`)
    A method for reading snapshots from a GROMACS GRO file.

read_gromacs_gro_file (:py:func:`.read_gromacs_gro_file`)
    Read a single snapshot from a GROMACS GRO file.

write_gromacs_gro_file (:py:func:`.write_gromacs_gro_file`)
    Write configuration in GROMACS GRO format.

read_gromos96_file (:py:func:`.read_gromos96_file`)
    Read a single configuration GROMACS .g96 file.

write_gromos96_fileĀ (:py:func:`.write_gromos96_file`)
    Write configuration in GROMACS g96 format.

read_xvg_file (:py:func:`.read_xvg_file`)
    For reading .xvg files from GROMACS.

read_trr_header (:py:func:`.read_trr_header`)
    Read a header from an open TRR file.

read_trr_data (:py:func:`.read_trr_data`)
    Read data from an open TRR file.

skip_trr_data (:py:func:`.skip_trr_data`)
    Skip reading data from an open TRR file and move on to the
    next header.

read_trr_file (:py:func:`.read_trr_file`)
    Yield frames from a TRR file.

trr_frame_to_g96 (:py:func:`.trr_frame_to_g96`)
    Dump a specific frame from a TRR file to a .g96 file.

write_trr_frame (:py:func:`.write_trr_frame`)
    A simple method to write to a TRR file.
"""
import logging
import struct
import io
import numpy as np
from pyretis.core.box import box_matrix_to_list
logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
logger.addHandler(logging.NullHandler())


__all__ = [
    'read_gromacs_file',
    'read_gromacs_gro_file',
    'write_gromacs_gro_file',
    'read_gromos96_file',
    'write_gromos96_file',
    'read_xvg_file',
    'read_trr_header',
    'read_trr_data',
    'skip_trr_data',
    'read_trr_file',
    'trr_frame_to_g96',
    'write_trr_frame',
]


# Define formats for the trajectory output:
_GRO_FMT = '{0:5d}{1:5s}{2:5s}{3:5d}{4:8.3f}{5:8.3f}{6:8.3f}'
_GRO_VEL_FMT = _GRO_FMT + '{7:8.4f}{8:8.4f}{9:8.4f}'
_GRO_BOX_FMT = '{:15.9f}'
_G96_FMT = '{0:}{1:15.9f}{2:15.9f}{3:15.9f}\n'
_G96_FMT_FULL = '{0:5d} {1:5s} {2:5s}{3:7d}{4:15.9f}{5:15.9f}{6:15.9f}\n'
_G96_BOX_FMT = '{:15.9f}' * 9 + '\n'
_G96_BOX_FMT_3 = '{:15.9f}' * 3 + '\n'

# Definitions for the TRR reader.
_GROMACS_MAGIC = 1993
_DIM = 3
_TRR_VERSION = 'GMX_trn_file'
_TRR_VERSION_B = b'GMX_trn_file'
_SIZE_FLOAT = struct.calcsize('f')
_SIZE_DOUBLE = struct.calcsize('d')
_HEAD_FMT = '{}13i'
_HEAD_ITEMS = ('ir_size', 'e_size', 'box_size', 'vir_size', 'pres_size',
               'top_size', 'sym_size', 'x_size', 'v_size', 'f_size',
               'natoms', 'step', 'nre', 'time', 'lambda')
TRR_DATA_ITEMS = ('box_size', 'vir_size', 'pres_size',
                  'x_size', 'v_size', 'f_size')


def read_gromacs_lines(lines):
    """Read and parse GROMACS GRO data.

    This method will read a GROMACS file and yield the different
    snapshots found in the file.

    Parameters
    ----------
    lines : iterable
        Some lines of text data representing a GROMACS GRO file.

    Yields
    ------
    out : dict
        This dict contains the snapshot.

    """
    lines_to_read = 0
    snapshot = {}
    read_natoms = False
    gro = (5, 5, 5, 5, 8, 8, 8, 8, 8, 8)
    gro_keys = ('residunr', 'residuname', 'atomname', 'atomnr',
                'x', 'y', 'z', 'vx', 'vy', 'vz')
    gro_type = (int, str, str, int, float, float, float, float, float, float)
    for line in lines:
        if read_natoms:
            read_natoms = False
            lines_to_read = int(line.strip()) + 1
            continue  # just skip to next line
        if lines_to_read == 0:  # new snapshot
            if snapshot:
                _add_matrices_to_snapshot(snapshot)
                yield snapshot
            snapshot = {'header': line.strip()}
            read_natoms = True
        elif lines_to_read == 1:  # read box
            snapshot['box'] = np.array(
                [float(i) for i in line.strip().split()]
            )
            lines_to_read -= 1
        else:  # read atoms
            lines_to_read -= 1
            current = 0
            for i, key, gtype in zip(gro, gro_keys, gro_type):
                val = line[current:current+i].strip()
                if not val:
                    # This typically happens if we try to read velocities
                    # and they are not present in the file.
                    break
                value = gtype(val)
                current += i
                try:
                    snapshot[key].append(value)
                except KeyError:
                    snapshot[key] = [value]
    if snapshot:
        _add_matrices_to_snapshot(snapshot)
        yield snapshot


[docs]def _add_matrices_to_snapshot(snapshot): """Extract positions and velocities as matrices from GROMACS. The extracted positions and velocities will be added to the given snapshot. Parameters ---------- snapshot : dict This dict contains the data read from the GROMACS file. Returns ------- xyz : numpy.array The positions as an array, (N, 3). vel : numpy.array The velocities as an array, (N, 3). """ xyz = np.zeros((len(snapshot['atomnr']), 3)) for i, key in enumerate(('x', 'y', 'z')): if key in snapshot: xyz[:, i] = snapshot[key] vel = np.zeros_like(xyz) for i, key in enumerate(('vx', 'vy', 'vz')): if key in snapshot: vel[:, i] = snapshot[key] snapshot['xyz'] = xyz snapshot['vel'] = vel return xyz, vel
[docs]def read_gromacs_file(filename): """Read GROMACS GRO files. This method will read a GROMACS file and yield the different snapshots found in the file. This file is intended to be used if we want to read all snapshots present in a file. Parameters ---------- filename : string The file to open. Yields ------ out : dict This dict contains the snapshot. Examples -------- >>> from pyretis.inout.formats.gromacs import read_gromacs_file >>> for snapshot in read_gromacs_file('traj.gro'): ... print(snapshot['x'][0]) """ with open(filename, 'r', encoding='utf-8') as fileh: for snapshot in read_gromacs_lines(fileh): yield snapshot
[docs]def read_gromacs_gro_file(filename): """Read a single configuration GROMACS GRO file. This method will read the first configuration from the GROMACS GRO file and return the data as give by :py:func:`.read_gromacs_lines`. It will also explicitly return the matrices with positions, velocities and box size. Parameters ---------- filename : string The file to read. Returns ------- frame : dict This dict contains all the data read from the file. xyz : numpy.array The positions. The array is (N, 3) where N is the number of particles. vel : numpy.array The velocities. The array is (N, 3) where N is the number of particles. box : numpy.array The box dimensions. """ snapshot = None xyz = None vel = None box = None with open(filename, 'r', encoding='utf8') as fileh: snapshot = next(read_gromacs_lines(fileh)) box = snapshot.get('box', None) xyz = snapshot.get('xyz', None) vel = snapshot.get('vel', None) return snapshot, xyz, vel, box
[docs]def write_gromacs_gro_file(outfile, txt, xyz, vel=None, box=None): """Write configuration in GROMACS GRO format. Parameters ---------- outfile : string The name of the file to create. txt : dict of lists of strings This dict contains the information on residue-numbers, names, etc. required to write the GRO file. xyz : numpy.array The positions to write. vel : numpy.array, optional The velocities to write. box: numpy.array, optional The box matrix. """ resnum = txt['residunr'] resname = txt['residuname'] atomname = txt['atomname'] atomnr = txt['atomnr'] npart = len(xyz) with open(outfile, 'w', encoding='utf-8') as output: output.write(f'{txt["header"]}\n') output.write(f'{npart}\n') for i in range(npart): if vel is None: buff = _GRO_FMT.format( resnum[i], resname[i], atomname[i], atomnr[i], xyz[i, 0], xyz[i, 1], xyz[i, 2]) else: buff = _GRO_VEL_FMT.format( resnum[i], resname[i], atomname[i], atomnr[i], xyz[i, 0], xyz[i, 1], xyz[i, 2], vel[i, 0], vel[i, 1], vel[i, 2]) output.write(f'{buff}\n') if box is None: box = ' '.join([_GRO_BOX_FMT.format(i) for i in txt['box']]) else: box = ' '.join([_GRO_BOX_FMT.format(i) for i in box]) output.write(f'{box}\n')
[docs]def read_gromos96_file(filename): """Read a single configuration GROMACS .g96 file. Parameters ---------- filename : string The file to read. Returns ------- rawdata : dict of list of strings This is the raw data read from the file grouped into sections. Note that this does not include the actual positions and velocities as these are returned separately. xyz : numpy.array The positions. vel : numpy.array The velocities. box : numpy.array The simulation box. """ _len = 15 _pos = 24 rawdata = {'TITLE': [], 'POSITION': [], 'VELOCITY': [], 'BOX': [], 'POSITIONRED': [], 'VELOCITYRED': []} section = None with open(filename, 'r', encoding='utf-8', errors='replace') as gromosfile: for lines in gromosfile: new_section = False stripline = lines.strip() if stripline == 'END': continue for key in rawdata: if stripline == key: new_section = True section = key break if new_section: continue rawdata[section].append(lines.rstrip()) txtdata = {} xyzdata = {} for key in ('POSITION', 'VELOCITY'): txtdata[key] = [] xyzdata[key] = [] for line in rawdata[key]: txt = line[:_pos] txtdata[key].append(txt) pos = [float(line[i:i+_len]) for i in range(_pos, 4*_len, _len)] xyzdata[key].append(pos) for line in rawdata[key+'RED']: txt = line[:_pos] txtdata[key].append(txt) pos = [float(line[i:i+_len]) for i in range(0, 3*_len, _len)] xyzdata[key].append(pos) xyzdata[key] = np.array(xyzdata[key]) rawdata['POSITION'] = txtdata['POSITION'] rawdata['VELOCITY'] = txtdata['VELOCITY'] if not rawdata['VELOCITY']: # No velocities were found in the input file. xyzdata['VELOCITY'] = np.zeros_like(xyzdata['POSITION']) logger.info('Input g96 did not contain velocities') if rawdata['BOX']: box = np.array([float(i) for i in rawdata['BOX'][0].split()]) else: box = None logger.info('Input g96 did not contain box vectors.') return rawdata, xyzdata['POSITION'], xyzdata['VELOCITY'], box
[docs]def write_gromos96_file(filename, raw, xyz, vel, box=None): """Write configuration in GROMACS .g96 format. Parameters ---------- filename : string The name of the file to create. raw : dict of lists of strings This contains the raw data read from a .g96 file. xyz : numpy.array The positions to write. vel : numpy.array The velocities to write. box: numpy.array, optional The box matrix. """ _keys = ('TITLE', 'POSITION', 'VELOCITY', 'BOX') with open(filename, 'w', encoding='utf-8') as outfile: for key in _keys: if key not in raw: continue outfile.write(f'{key}\n') for i, line in enumerate(raw[key]): if key == 'POSITION': outfile.write(_G96_FMT.format(line, *xyz[i])) elif key == 'VELOCITY': outfile.write(_G96_FMT.format(line, *vel[i])) elif box is not None and key == 'BOX': if len(box) == 3: outfile.write(_G96_BOX_FMT_3.format(*box)) else: outfile.write(_G96_BOX_FMT.format(*box)) else: outfile.write(f'{line}\n') outfile.write('END\n')
[docs]def read_xvg_file(filename): """Return data in xvg file as numpy array.""" data = [] legends = [] with open(filename, 'r', encoding='utf-8') as fileh: for lines in fileh: if lines.startswith('@ s') and lines.find('legend') != -1: legend = lines.split('legend')[-1].strip() legend = legend.replace('"', '') legends.append(legend.lower()) else: if lines.startswith('#') or lines.startswith('@'): pass else: data.append([float(i) for i in lines.split()]) data = np.array(data) data_dict = {'step': np.arange(tuple(data.shape)[0])} for i, key in enumerate(legends): data_dict[key] = data[:, i+1] return data_dict
def swap_integer(integer): """Convert little/big endian.""" return (((integer << 24) & 0xff000000) | ((integer << 8) & 0x00ff0000) | ((integer >> 8) & 0x0000ff00) | ((integer >> 24) & 0x000000ff)) def swap_endian(endian): """Just swap the string for selecting big/little.""" if endian == '>': return '<' if endian == '<': return '>' raise ValueError('Undefined swap!') def read_struct_buff(fileh, fmt): """Unpack from a file handle with a given format. Parameters ---------- fileh : file object The file handle to unpack from. fmt : string The format to use for unpacking. Returns ------- out : tuple The unpacked elements according to the given format. Raises ------ EOFError We will raise an EOFError if `fileh.read()` attempts to read past the end of the file. """ buff = fileh.read(struct.calcsize(fmt)) if not buff: raise EOFError return struct.unpack(fmt, buff) def read_matrix(fileh, endian, double): """Read a matrix from the TRR file. Here, we assume that the matrix will be of dimensions (_DIM, _DIM). Parameters ---------- fileh : file object The file handle to read from. endian : string Determines the byte order. double : boolean If true, we will assume that the numbers were stored in double precision. Returns ------- mat : numpy.array The matrix as an array. """ if double: fmt = f'{endian}{_DIM**2}d' else: fmt = f'{endian}{_DIM**2}f' read = read_struct_buff(fileh, fmt) mat = np.zeros((_DIM, _DIM)) for i in range(_DIM): for j in range(_DIM): mat[i, j] = read[i * _DIM + j] return mat def read_coord(fileh, endian, double, natoms): """Read a coordinate section from the TRR file. This method will read the full coordinate section from a TRR file. The coordinate section may be positions, velocities or forces. Parameters ---------- fileh : file object The file handle to read from. endian : string Determines the byte order. double : boolean If true, we will assume that the numbers were stored in double precision. natoms : int The number of atoms we have stored coordinates for. Returns ------- mat : numpy.array The coordinates as a numpy array. It will have ``natoms`` rows and ``_DIM`` columns. """ if double: fmt = f'{endian}{natoms * _DIM}d' else: fmt = f'{endian}{natoms * _DIM}f' read = read_struct_buff(fileh, fmt) mat = np.array(read) mat.shape = (natoms, _DIM) return mat def is_double(header): """Determine if we should use double precision. This method determined the precision to use when reading the TRR file. This is based on the header read for a given frame which defines the sizes of certain "fields" like the box or the positions. From this size, the precision can be obtained. Parameters ---------- header : dict The header read from the TRR file. Returns ------- out : boolean True if we should use double precision. """ key_order = ('box_size', 'x_size', 'v_size', 'f_size') size = 0 for key in key_order: if header[key] != 0: if key == 'box_size': size = int(header[key] / _DIM**2) break size = int(header[key] / (header['natoms'] * _DIM)) break if size not in (_SIZE_FLOAT, _SIZE_DOUBLE): raise ValueError('Could not determine size!') return size == _SIZE_DOUBLE
[docs]def read_trr_header(fileh): """Read a header from a TRR file. Parameters ---------- fileh : file object The file handle for the file we are reading. Returns ------- header : dict The header read from the file. """ start = fileh.tell() endian = '>' magic = read_struct_buff(fileh, f'{endian}1i')[0] if magic == _GROMACS_MAGIC: pass else: magic = swap_integer(magic) if not magic == _GROMACS_MAGIC: logger.critical( 'TRR file might be inconsistent! Could find _GROMACS_MAGIC' ) endian = swap_endian(endian) slen = read_struct_buff(fileh, f'{endian}2i') raw = read_struct_buff(fileh, f'{endian}{slen[0] - 1}s') version = raw[0].split(b'\0', 1)[0].decode('utf-8') if not version == _TRR_VERSION: raise ValueError('Unknown format') head_fmt = _HEAD_FMT.format(endian) head_s = read_struct_buff(fileh, head_fmt) header = {} for i, val in enumerate(head_s): key = _HEAD_ITEMS[i] header[key] = val # The next are either floats or double double = is_double(header) if double: fmt = f'{endian}2d' else: fmt = f'{endian}2f' header_r = read_struct_buff(fileh, fmt) header['time'] = header_r[0] header['lambda'] = header_r[1] header['endian'] = endian header['double'] = double return header, fileh.tell() - start
[docs]def skip_trr_data(fileh, header): """Skip coordinates/box data etc. This method is used when we want to skip a data section in the TRR file. Rather than reading the data, it will use the size read in the header to skip ahead to the next frame. Parameters ---------- fileh : file object The file handle for the file we are reading. header : dict The header read from the TRR file. """ offset = sum([header[key] for key in TRR_DATA_ITEMS]) fileh.seek(offset, 1)
[docs]def read_trr_data(fileh, header): """Read box, coordinates etc. from a TRR file. Parameters ---------- fileh : file object The file handle for the file we are reading. header : dict The header read from the file. Returns ------- data : dict The data we read from the file. It may contain the following keys if the data was found in the frame: - ``box`` : the box matrix, - ``vir`` : the virial matrix, - ``pres`` : the pressure matrix, - ``x`` : the coordinates, - ``v`` : the velocities, and - ``f`` : the forces """ data = {} endian = header['endian'] double = header['double'] for key in ('box', 'vir', 'pres'): header_key = f'{key}_size' if header[header_key] != 0: data[key] = read_matrix(fileh, endian, double) for key in ('x', 'v', 'f'): header_key = f'{key}_size' if header[header_key] != 0: data[key] = read_coord(fileh, endian, double, header['natoms']) return data
[docs]def read_trr_file(filename, read_data=True): """Yield frames from a TRR file.""" with open(filename, 'rb') as infile: while True: try: header, _ = read_trr_header(infile) if read_data: data = read_trr_data(infile, header) else: skip_trr_data(infile, header) data = None yield header, data except EOFError: return except struct.error: logger.warning( 'Could not read a frame from the TRR file. Aborting!' ) return
def read_trr_frame(filename, index): """Return a given frame from a TRR file.""" idx = 0 with open(filename, 'rb') as infile: while True: try: header, _ = read_trr_header(infile) if idx == index: data = read_trr_data(infile, header) return header, data skip_trr_data(infile, header) idx += 1 if idx > index: logger.error('Frame %i not found in %s', index, filename) return None, None except EOFError: return None, None
[docs]def trr_frame_to_g96(trr_file, index, outfile): """Dump a TRR frame to a GROMOS96 (.g96) file. Parameters ---------- trr_file : string The TRR file to read from. index : integer The index for the frame to dump from ``trr_file``. outfile : string The g96 file to write. """ header, data = read_trr_frame(trr_file, index) with open(outfile, 'w', encoding='utf-8') as output: output.write('TITLE\n') output.write(f' Dump from TRR, index = {index}, ' f'time = {header["time"]}\n') output.write('END\n') output.write('POSITION\n') for i, posi in enumerate(data['x']): output.write(_G96_FMT_FULL.format(i, 'DUM', 'X', i, *posi)) output.write('END\n') if 'v' in data: output.write('VELOCITY\n') for i, veli in enumerate(data['v']): output.write(_G96_FMT_FULL.format(i, 'DUM', 'X', i, *veli)) output.write('END\n') output.write('BOX\n') box = data['box'] output.write(_G96_BOX_FMT.format(*box_matrix_to_list(box, full=True))) output.write('END\n')
[docs]def _get_chunks(start, end, size): """Yield chunks between start and end of the given size.""" while start < end: new_size = min(size, end - start) start += new_size yield start, new_size
def reverse_trr(filename, outname, print_progress=True): """Reverse a GROMACS TRR file. Parameters ---------- filename : string The file path to the input file. outname : string The file path to the output file. print_progress : boolean, optional This can be used to print out information about the progress to the screen. Might be useful if we are reading a large file. """ # We first loop over the file until we reach the # last header, we skip the data and read all the headers. all_headers = [] buff_size = io.DEFAULT_BUFFER_SIZE with open(filename, 'rb') as infile, open(outname, 'wb') as outfile: while True: try: header, header_size = read_trr_header(infile) header_loc = infile.tell() skip_trr_data(infile, header) all_headers.append((header, header_loc, header_size)) except EOFError: break # Loop through headers in reverse and write data. for header, header_loc, header_size in reversed(all_headers): if print_progress: # pragma: no cover print(f'Processing step {header["step"]} ' f'timeĀ {header["time"]}') data_size = sum([header[key] for key in TRR_DATA_ITEMS]) start = header_loc - header_size infile.seek(start) end = start + header_size + data_size for _, chunk_size in _get_chunks(start, end, buff_size): outfile.write(infile.read(chunk_size))
[docs]def write_trr_frame(filename, data, endian=None, double=False, append=False): """Write data in TRR format to a file. Parameters ---------- filename : string The name/path of the file to write to. data : dict The data we will write to the file. endian : string, optional Select the byte order; big-endian or little-endian. If not specified, the native byte order will be used. double : boolean, optional If True, we will write in double precision. append : boolean, optional If True, we will append to the file, otherwise, the file will be overwritten. """ if append: mode = 'ab' else: mode = 'wb' if double: size = _SIZE_DOUBLE floatfmt = '{}d' else: floatfmt = '{}f' size = _SIZE_FLOAT if endian: floatfmt = endian + floatfmt header = {} for key in _HEAD_ITEMS: header[key] = 0 header['natoms'] = data['natoms'] header['step'] = data['step'] header['box_size'] = size * _DIM * _DIM for i in ('x', 'v', 'f'): if i in data: header[f'{i}_size'] = data['natoms'] * size * _DIM header['endian'] = endian header['double'] = double header['time'] = data['time'] header['lambda'] = data['lambda'] with open(filename, mode) as outfile: write_trr_header(outfile, header, floatfmt, endian=endian) for key in TRR_DATA_ITEMS: if header[key] != 0: matrix = data[key.split('_', 1)[0]] fmt = floatfmt.format(matrix.size) outfile.write(struct.pack(fmt, *matrix.flatten())) return header
def write_trr_header(outfile, header, floatfmt, endian=None): """Write TRR header. Parameters ---------- outfile : filehandle The file we can write to. header : dict The header data for the TRR file. floatfmt : string The string which gives the format for floats. It should indicate if we are writing for double or single precision. endian : string, optional Can be used to force endianess. """ slen = (13, 12) fmt = ['1i', '2i', f'{slen[0] - 1}s', '13i'] if endian: fmt = [endian + i for i in fmt] outfile.write(struct.pack(fmt[0], _GROMACS_MAGIC)) outfile.write(struct.pack(fmt[1], *slen)) outfile.write(struct.pack(fmt[2], _TRR_VERSION_B)) head = [header[key] for key in _HEAD_ITEMS[:13]] outfile.write(struct.pack(fmt[3], *head)) outfile.write(struct.pack(floatfmt.format(1), header['time'])) outfile.write(struct.pack(floatfmt.format(1), header['lambda'])) outfile.flush()