def test_read_5(): with DCDTrajectoryFile(fn_dcd) as f: xyz_ref, box_lengths_ref, box_angles_ref = f.read() with DCDTrajectoryFile(fn_dcd) as f: xyz, box_lengths, box_angles = f.read(atom_indices=[1, 2, 5]) yield lambda: eq(xyz_ref[:, [1, 2, 5], :], xyz)
def test_seek(): reference = DCDTrajectoryFile(get_fn('frame0.dcd')).read()[0] with DCDTrajectoryFile(get_fn('frame0.dcd')) as f: eq(f.tell(), 0) eq(f.read(1)[0][0], reference[0]) eq(f.tell(), 1) xyz = f.read(1)[0][0] eq(xyz, reference[1]) eq(f.tell(), 2) f.seek(0) eq(f.tell(), 0) xyz = f.read(1)[0][0] eq(f.tell(), 1) eq(xyz, reference[0]) f.seek(5) eq(f.read(1)[0][0], reference[5]) eq(f.tell(), 6) f.seek(-5, 1) eq(f.tell(), 1) eq(f.read(1)[0][0], reference[1])
def test_read_6(): with DCDTrajectoryFile(fn_dcd) as f: xyz_ref, box_lengths_ref, box_angles_ref = f.read() with DCDTrajectoryFile(fn_dcd) as f: xyz, box_lengths, box_angles = f.read( atom_indices=slice(None, None, 2)) yield lambda: eq(xyz_ref[:, ::2, :], xyz)
def test_read_5(get_fn): fn_dcd = get_fn('frame0.dcd') with DCDTrajectoryFile(fn_dcd) as f: xyz_ref, box_lengths_ref, box_angles_ref = f.read() with DCDTrajectoryFile(fn_dcd) as f: xyz, box_lengths, box_angles = f.read(atom_indices=[1, 2, 5]) assert eq(xyz_ref[:, [1, 2, 5], :], xyz)
def test_read_2(): "DCDReader: check nframes" xyz1, box_lengths1, box_angles1 = DCDTrajectoryFile(fn_dcd).read() xyz2, box_lengths2, box_angles2 = DCDTrajectoryFile(fn_dcd).read(10000) yield lambda: eq(xyz1, xyz2) yield lambda: eq(box_lengths1, box_lengths2) yield lambda: eq(box_angles1, box_angles2)
def test_write_0(): with DCDTrajectoryFile(fn_dcd) as f: xyz = f.read()[0] with DCDTrajectoryFile(temp, 'w') as f: f.write(xyz) with DCDTrajectoryFile(temp) as f: xyz2 = f.read()[0] eq(xyz, xyz2)
def test_write_1(): xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32) with DCDTrajectoryFile(temp, 'w') as f: f.write(xyz) with DCDTrajectoryFile(temp) as f: xyz2 = f.read()[0] eq(xyz, xyz2)
def test_read_6(get_fn): fn_dcd = get_fn('frame0.dcd') with DCDTrajectoryFile(fn_dcd) as f: xyz_ref, box_lengths_ref, box_angles_ref = f.read() with DCDTrajectoryFile(fn_dcd) as f: xyz, box_lengths, box_angles = f.read( atom_indices=slice(None, None, 2)) assert eq(xyz_ref[:, ::2, :], xyz)
def test_read_2(get_fn): # check nframes fn_dcd = get_fn('frame0.dcd') xyz1, box_lengths1, box_angles1 = DCDTrajectoryFile(fn_dcd).read() xyz2, box_lengths2, box_angles2 = DCDTrajectoryFile(fn_dcd).read(10000) assert eq(xyz1, xyz2) assert eq(box_lengths1, box_lengths2) assert eq(box_angles1, box_angles2)
def test_read_3(): "DCDReader: check streaming read of frames 1 at a time" xyz_ref, box_lengths_ref, box_angles_ref = DCDTrajectoryFile(fn_dcd).read() reader = DCDTrajectoryFile(fn_dcd) for i in range(len(xyz_ref)): xyz, box_lenths, box_angles = reader.read(1) eq(xyz_ref[np.newaxis, i], xyz) eq(box_lengths_ref[np.newaxis, i], box_lenths) eq(box_angles_ref[np.newaxis, i], box_angles)
def test_write_1(tmpdir): fn = '{}/x.dcd'.format(tmpdir) xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32) with DCDTrajectoryFile(fn, 'w') as f: f.write(xyz) with DCDTrajectoryFile(fn) as f: xyz2 = f.read()[0] eq(xyz, xyz2)
def test_read_stride_2(): "Read dcd with stride when n_frames is supplied (different code path)" with DCDTrajectoryFile(fn_dcd) as f: xyz1, box_lengths1, box_angles1 = f.read() with DCDTrajectoryFile(fn_dcd) as f: xyz2, box_lengths2, box_angles2 = f.read(n_frames=1000, stride=2) yield lambda: eq(xyz1[::2], xyz2) yield lambda: eq(box_lengths1[::2], box_lengths2) yield lambda: eq(box_angles1[::2], box_angles2)
def test_read_stride(): "Read dcd with stride" with DCDTrajectoryFile(fn_dcd) as f: xyz1, box_lengths1, box_angles1 = f.read() with DCDTrajectoryFile(fn_dcd) as f: xyz2, box_lengths2, box_angles2 = f.read(stride=2) yield lambda: eq(xyz1[::2], xyz2) yield lambda: eq(box_lengths1[::2], box_lengths2) yield lambda: eq(box_angles1[::2], box_angles2)
def test_write_0(tmpdir, get_fn): fn_dcd = get_fn('frame0.dcd') fn = '{}/x.dcd'.format(tmpdir) with DCDTrajectoryFile(fn_dcd) as f: xyz = f.read()[0] with DCDTrajectoryFile(fn, 'w') as f: f.write(xyz) with DCDTrajectoryFile(fn) as f: xyz2 = f.read()[0] eq(xyz, xyz2)
def test_read_3(get_fn): # Check streaming read of frames 1 at a time fn_dcd = get_fn('frame0.dcd') xyz_ref, box_lengths_ref, box_angles_ref = DCDTrajectoryFile(fn_dcd).read() reader = DCDTrajectoryFile(fn_dcd) for i in range(len(xyz_ref)): xyz, box_lenths, box_angles = reader.read(1) eq(xyz_ref[np.newaxis, i], xyz) eq(box_lengths_ref[np.newaxis, i], box_lenths) eq(box_angles_ref[np.newaxis, i], box_angles)
def test_read_stride_2(get_fn): # Read dcd with stride when n_frames is supplied (different code path) fn_dcd = get_fn('frame0.dcd') with DCDTrajectoryFile(fn_dcd) as f: xyz1, box_lengths1, box_angles1 = f.read() with DCDTrajectoryFile(fn_dcd) as f: xyz2, box_lengths2, box_angles2 = f.read(n_frames=1000, stride=2) assert eq(xyz1[::2], xyz2) assert eq(box_lengths1[::2], box_lengths2) assert eq(box_angles1[::2], box_angles2)
def test_read_stride(get_fn): # Read dcd with stride fn_dcd = get_fn('frame0.dcd') with DCDTrajectoryFile(fn_dcd) as f: xyz1, box_lengths1, box_angles1 = f.read() with DCDTrajectoryFile(fn_dcd) as f: xyz2, box_lengths2, box_angles2 = f.read(stride=2) assert eq(xyz1[::2], xyz2) assert eq(box_lengths1[::2], box_lengths2) assert eq(box_angles1[::2], box_angles2)
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 test_write_3(): "checking" xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32) box_lengths = 25 * np.ones((600, 3), dtype=np.float32) with DCDTrajectoryFile(temp, 'w') as f: f.write(xyz, box_lengths)
def write_trajectory_dcd(netcdf_filename, topology, pdb_trajectory_filename, dcd_trajectory_filename): """ Write trajectory. Parameters ---------- netcdf_filename : str NetCDF filename. topology : Topology Topology object pdb_trajectory_filename : str PDB trajectory output filename dcd_trajectory_filename : str Output trajectory filename. """ ncfile = netCDF4.Dataset(netcdf_filename, 'r') [nsamples, nstates] = ncfile.variables['logZ'].shape # Write reference.pdb file from simtk.openmm.app import PDBFile outfile = open(pdb_trajectory_filename, 'w') positions = unit.Quantity(ncfile.variables['positions'][0,:,:], unit.nanometers) PDBFile.writeFile(topology, positions, file=outfile) outfile.close() # TODO: Export as DCD trajectory with MDTraj from mdtraj.formats import DCDTrajectoryFile with DCDTrajectoryFile(dcd_trajectory_filename, 'w') as f: f.write(ncfile.variables['positions'][:,:,:] * 10.0) # angstroms
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 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 test_do_overwrite(tmpdir): fn = '{}/x.dcd'.format(tmpdir) with open(fn, 'w') as f: f.write('a') with DCDTrajectoryFile(fn, 'w', force_overwrite=True) as f: f.write(np.random.randn(10, 5, 3))
def test_read_4(): "DCDReader: check streaming read followed by reading the 'rest'" xyz_ref, box_lengths_ref, box_angles_ref = DCDTrajectoryFile(fn_dcd).read() reader = DCDTrajectoryFile(fn_dcd) for i in range(int(len(xyz_ref) / 2)): xyz, box_lenths, box_angles = reader.read(1) eq(xyz_ref[np.newaxis, i], xyz) eq(box_lengths_ref[np.newaxis, i], box_lenths) eq(box_angles_ref[np.newaxis, i], box_angles) xyz_rest, box_rest, angles_rest = reader.read() yield lambda: eq(xyz_ref[i + 1:], xyz_rest) yield lambda: eq(box_lengths_ref[i + 1:], box_rest) yield lambda: eq(box_angles_ref[i + 1:], angles_rest) yield lambda: len(xyz_ref) == i + len(xyz_rest)
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_3(tmpdir): fn = '{}/x.dcd'.format(tmpdir) xyz = np.array(np.random.randn(500, 10, 3), dtype=np.float32) box_lengths = 25 * np.ones((600, 3), dtype=np.float32) with DCDTrajectoryFile(fn, 'w') as f: with pytest.raises(ValueError): f.write(xyz, box_lengths)
def test_dont_overwrite(tmpdir): fn = '{}/x.dcd'.format(tmpdir) with open(fn, 'w') as f: f.write('a') with pytest.raises(IOError): with DCDTrajectoryFile(fn, 'w', force_overwrite=False) as f: f.write(np.random.randn(10, 5, 3))
def test_tell(): with DCDTrajectoryFile(get_fn('frame0.dcd')) as f: eq(f.tell(), 0) f.read(101) eq(f.tell(), 101) f.read(3) eq(f.tell(), 104)
def test_read_1(): with BINPOSTrajectoryFile(fn_binpos, chunk_size_multiplier=0.5) as f: xyz = f.read() with DCDTrajectoryFile(fn_dcd) as f: xyz2 = f.read()[0] xyz3 = io.loadh(get_fn('frame0.binpos.h5'), 'xyz') yield lambda: eq(xyz[1:], xyz2) yield lambda: eq(xyz, xyz3)
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)