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
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
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)
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)
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)
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)
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-----------')
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))
def test_read_closed(get_fn): fn_dcd = get_fn('frame0.dcd') with pytest.raises(IOError): f = DCDTrajectoryFile(fn_dcd) f.close() f.read()
def test_write_closed(): f = DCDTrajectoryFile(fn_dcd, "w") f.close() f.write(np.random.randn(10, 3, 3))
def test_read_closed(): f = DCDTrajectoryFile(fn_dcd) f.close() f.read()
def test_write_closed(): f = DCDTrajectoryFile(fn_dcd, 'w') f.close() f.write(np.random.randn(10, 3, 3))
def test_read_closed(): f = DCDTrajectoryFile(fn_dcd) f.close() f.read()
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
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))
def test_read_closed(get_fn): fn_dcd = get_fn('frame0.dcd') with pytest.raises(IOError): f = DCDTrajectoryFile(fn_dcd) f.close() f.read()