Beispiel #1
0
def get_frames_num(filepath, file_type, ext=None):
    if file_type == 'coor':
        if ext is None:
            traj = mdtraj_load(filepath)
        else:
            ext2 = ext.lower()
            if ext2 == 'pdb':
                traj = mdtraj_load_pdb(filepath)
            else:
                raise ValueError('Extension "' + ext2 + '" not implemented.')
        numframes = traj.n_frames

    elif file_type == 'traj':

        if ext is None:
            trajfile = mdtraj_open(filepath)
        else:
            ext2 = ext.lower()
            if ext2 == 'dcd':
                trajfile = DCDTrajectoryFile(filepath)
            elif ext2 == 'xtc':
                trajfile = XTCTrajectoryFile(filepath)
            else:
                raise ValueError('Extension "' + ext + '" not implemented.')
        numframes = mdtraj_get_frames_num(trajfile)
        trajfile.close()
    return numframes
Beispiel #2
0
def get_atoms_num(filepath, file_type, ext=None):

    if file_type == 'coor':
        if ext is None:
            traj = mdtraj_load(filepath)
        else:
            ext2 = ext.lower()
            if ext2 == 'pdb':
                traj = mdtraj_load_pdb(filepath)
            else:
                raise ValueError('Extension "' + ext2 + '" not implemented.')
        numatoms = traj.n_atoms
    elif file_type == 'traj':

        if ext is None:
            trajfile = mdtraj_open(filepath)
        else:
            ext2 = ext.lower()
            if ext2 == 'dcd':
                trajfile = DCDTrajectoryFile(filepath)
            elif ext2 == 'xtc':
                trajfile = XTCTrajectoryFile(filepath)
            else:
                raise ValueError('Extension "' + ext + '" not implemented.')
        res = trajfile.read(1)
        xyz = res[0]
        numatoms = xyz.shape[1]
        trajfile.close()
    return numatoms
Beispiel #3
0
def test_write_2():
    xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32)
    box_lengths = 25 * np.ones((500, 3), dtype=np.float32)
    box_angles = 90 * np.ones((500, 3), dtype=np.float32)
    box_lengths[0, 0] = 10.0

    f = DCDTrajectoryFile(temp, "w")
    f.write(xyz, box_lengths, box_angles)
    f.close()

    f = DCDTrajectoryFile(temp)
    xyz2, box_lengths2, box_angles2 = f.read()
    f.close()

    yield lambda: eq(xyz, xyz2)
    yield lambda: eq(box_lengths, box_lengths2)
    yield lambda: eq(box_angles, box_angles2)
Beispiel #4
0
def test_write_2():
    xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32)
    box_lengths = 25 * np.ones((500, 3), dtype=np.float32)
    box_angles = 90 * np.ones((500, 3), dtype=np.float32)
    box_lengths[0, 0] = 10.0

    f = DCDTrajectoryFile(temp, 'w')
    f.write(xyz, box_lengths, box_angles)
    f.close()

    f = DCDTrajectoryFile(temp)
    xyz2, box_lengths2, box_angles2 = f.read()
    f.close()

    yield lambda: eq(xyz, xyz2)
    yield lambda: eq(box_lengths, box_lengths2)
    yield lambda: eq(box_angles, box_angles2)
Beispiel #5
0
def test_write_2(tmpdir):
    fn = '{}/x.dcd'.format(tmpdir)
    xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32)
    box_lengths = 25 * np.ones((500, 3), dtype=np.float32)
    box_angles = 90 * np.ones((500, 3), dtype=np.float32)
    box_lengths[0, 0] = 10.0

    f = DCDTrajectoryFile(fn, 'w')
    f.write(xyz, box_lengths, box_angles)
    f.close()

    f = DCDTrajectoryFile(fn)
    xyz2, box_lengths2, box_angles2 = f.read()
    f.close()

    assert eq(xyz, xyz2)
    assert eq(box_lengths, box_lengths2)
    assert eq(box_angles, box_angles2)
Beispiel #6
0
def test_write_2(tmpdir):
    fn = '{}/x.dcd'.format(tmpdir)
    xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32)
    box_lengths = 25 * np.ones((500, 3), dtype=np.float32)
    box_angles = 90 * np.ones((500, 3), dtype=np.float32)
    box_lengths[0, 0] = 10.0

    f = DCDTrajectoryFile(fn, 'w')
    f.write(xyz, box_lengths, box_angles)
    f.close()

    f = DCDTrajectoryFile(fn)
    xyz2, box_lengths2, box_angles2 = f.read()
    f.close()

    assert eq(xyz, xyz2)
    assert eq(box_lengths, box_lengths2)
    assert eq(box_angles, box_angles2)
Beispiel #7
0
    def read_onefile(self):
        """ Read all snapshots from one dump file 
            
            The keyword filetype is used for different MD engines
            It has four choices:
            'lammps' (default)

            'lammpscenter' (lammps molecular dump with known atom type of each molecule center)
            moltypes is a dict mapping center atomic type to molecular type
            moltypes is also used to select the center atom 
            such as moltypes = {3: 1, 5: 2}

            'gsd' (HOOMD-blue standard output for static properties)
            
            'gsd_dcd' (HOOMD-blue outputs for static and dynamic properties)
        """

        #------------------lammps atomic trajectory---------------------------------
        if self.filetype == 'lammps':
            f = open(self.filename, 'r')
            positions = self.read_lammps(f)
            while positions.any():
                self.Positions.append(positions)
                positions = self.read_lammps(f)

            f.close()
            self.SnapshotNumber = len(self.TimeStep)
            print('--------LAMMPS Atomic Dump Reading Over---------')

#------------------lammps molecular trajectory---------------------------------
        if self.filetype == 'lammpscenter':
            f = open(self.filename, 'r')
            positions = self.read_centertype(f)
            while positions.any():
                self.Positions.append(positions)
                positions = self.read_centertype(f)

            f.close()
            self.SnapshotNumber = len(self.TimeStep)
            print('--------LAMMPS Molecule Center Dump Reading Over---------')

#------------------hoomd-blue gsd trajectory---------------------------------
        if self.filetype == 'gsd':
            import gsd.hoomd
            import gsd

            f = gsd.hoomd.open(self.filename, mode='rb')
            self.read_gsd(f)
            print('---------GSD file reading over-----------')

#------------------hoomd-blue gsd with dcd trajectory---------------------------------
        if self.filetype == 'gsd_dcd':
            import gsd.hoomd
            import gsd, os
            from mdtraj.formats import DCDTrajectoryFile

            gsd_filename = self.filename
            gsd_filepath = os.path.dirname(gsd_filename)
            dcd_filename = gsd_filepath + '/' + os.path.basename(
                gsd_filename)[:-3] + 'dcd'

            f_gsd = gsd.hoomd.open(gsd_filename, mode='rb')
            f_dcd = DCDTrajectoryFile(dcd_filename, 'r')
            self.read_gsd_dcd(f_gsd, f_dcd)
            f_dcd.close()
            print('---------GSD & DCD file reading over-----------')
Beispiel #8
0
def test_write_closed(get_fn):
    fn_dcd = get_fn('frame0.dcd')
    with pytest.raises(IOError):
        f = DCDTrajectoryFile(fn_dcd, 'w')
        f.close()
        f.write(np.random.randn(10, 3, 3))
Beispiel #9
0
def test_read_closed(get_fn):
    fn_dcd = get_fn('frame0.dcd')
    with pytest.raises(IOError):
        f = DCDTrajectoryFile(fn_dcd)
        f.close()
        f.read()
Beispiel #10
0
def test_write_closed():
    f = DCDTrajectoryFile(fn_dcd, "w")
    f.close()
    f.write(np.random.randn(10, 3, 3))
Beispiel #11
0
def test_read_closed():
    f = DCDTrajectoryFile(fn_dcd)
    f.close()
    f.read()
Beispiel #12
0
def test_write_closed():
    f = DCDTrajectoryFile(fn_dcd, 'w')
    f.close()
    f.write(np.random.randn(10, 3, 3))
Beispiel #13
0
def test_read_closed():
    f = DCDTrajectoryFile(fn_dcd)
    f.close()
    f.read()
Beispiel #14
0
class Dcd(TrjHandler):
    '''
    Read and write cell and positions from/to DCD file.

    Velocities, step and time are ignored.

    Currently mstools use mdtraj to support DCD file.
    Therefore, appending is not supported.
    '''

    def __init__(self, file, mode='r'):
        super().__init__()

        try:
            from mdtraj.formats import DCDTrajectoryFile
        except:
            raise ImportError('Currently mstools use mdtraj to parse DCD format. Cannot import mdtraj')

        if mode == 'r':
            self._dcd = DCDTrajectoryFile(file, mode='r')
        elif mode == 'w':
            self._dcd = DCDTrajectoryFile(file, mode='w')
        else:
            raise Exception('Appending not supported for DCD')

    def close(self):
        try:
            self._dcd.close()
        except:
            pass

    def get_info(self):
        self.n_frame = len(self._dcd)
        if self.n_frame == 0:
            raise Exception('Empty DCD file')
        positions, lengths, angles = self._dcd.read(1)
        _, self.n_atom, _ = positions.shape

        return self.n_atom, self.n_frame

    def read_frame(self, i_frame, frame):
        self._dcd.seek(i_frame)
        positions, box_lengths, box_angles = self._dcd.read(1)
        angle = box_angles[0]
        angle[np.abs(angle - 90) < 1E-4] = 90  # in case precision issue
        frame.positions = positions[0] / 10  # convert A to nm
        frame.cell.set_box([box_lengths[0] / 10, angle])  # convert A to nm

    def write_frame(self, frame, subset=None, **kwargs):
        '''
        Write a frame into the opened DCD file

        Parameters
        ----------
        frame : Frame
        subset : list of int, optional
        kwargs : dict
            Ignored
        '''
        if subset is None:
            positions = frame.positions
        else:
            positions = frame.positions[subset]
        self._dcd.write(positions * 10, frame.cell.lengths * 10, frame.cell.angles)  # convert nm to A
Beispiel #15
0
def test_write_closed(get_fn):
    fn_dcd = get_fn('frame0.dcd')
    with pytest.raises(IOError):
        f = DCDTrajectoryFile(fn_dcd, 'w')
        f.close()
        f.write(np.random.randn(10, 3, 3))
Beispiel #16
0
def test_read_closed(get_fn):
    fn_dcd = get_fn('frame0.dcd')
    with pytest.raises(IOError):
        f = DCDTrajectoryFile(fn_dcd)
        f.close()
        f.read()