formattedtext.py

import numpy as np

class FormattedTextReader(object):
    # this is a magic number that we'll put in the coordinates array
    # at the beginning when it's allocated, and then check to make sure
    # at the end that it's been overwritten. this way we're sure that all
    # of our returned data actually was read from the file
    _magic = 123456789.0
    
    def __init__(self, filename, n_atoms):
        self.handle = open(filename, 'r')
        self.n_atoms = n_atoms
        
        # we deposit the numbers after reading them from disk here
        self._buffer = []
        
        # the first line of the trajectory format appears to be a garbage
        # string that we just want to ignore. maybe it's a comment?
        self.handle.readline()
    
    def read(self, n_frames):
        coordinates = self._magic * np.ones((n_frames, self.n_atoms, 3))        

        for i in range(n_frames):
            # fill the buffer with data from this frame
            while len(self._buffer) < 3 * self.n_atoms:
                # add a new line from the file into the buffer
                got = map(float, self.handle.readline().split())

                if len(got) == 0:
                    # this must be the EOF
                    assert len(self._buffer) == 0, 'Read Error'
                    return coordinates[0:i]

                self._buffer.extend(got)
                
            # take the first self.n_atoms*3 entries from the buffer
            frame = np.array(self._buffer[:self.n_atoms*3])
            # and leave the rest
            self._buffer = self._buffer[self.n_atoms*3:]

            coordinates[i] = np.array(frame).reshape(self.n_atoms, 3)

        if self._magic in coordinates:
            raise IOError('Read error. The number of expected coordinates was not parsed')

        return coordinates

test_formattedtext.py

from mdtraj.testing import get_fn, eq
from mdtraj.netcdf import NetCDFFile
 
from formattedtext import FormattedTextReader
 
 
def test_1():
    t = FormattedTextReader('amber-trajectory-examples/implicit-formatted.trj', 22)
    t2 = NetCDFFile('amber-trajectory-examples/implicit-netcdf.trj')
 
    n_coordinates = t2.read()[0]
    f_coordinates = t.read(20)
 
    yield lambda: eq(n_coordinates.shape, (10, 22, 3))
    yield lambda: eq(n_coordinates, f_coordinates, decimal=3)