def test_seek(): reference = XTCTrajectoryFile(get_fn('frame0.xtc')).read()[0] with XTCTrajectoryFile(get_fn('frame0.xtc')) 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_seek(): reference = XTCTrajectoryFile(get_fn('frame0.xtc')).read()[0] with XTCTrajectoryFile(get_fn('frame0.xtc')) 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) # offset array is going to be built assert len(f.offsets) == len(reference) 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_xtc_write_wierd_0(): x0 = np.asarray(np.random.randn(100,3,3), dtype=np.float32) x1 = np.asarray(np.random.randn(100,9,3), dtype=np.float32) with XTCTrajectoryFile(temp, 'w') as f: f.write(x0) f.write(x1) xr = XTCTrajectoryFile(temp).read()[0] print(xr.shape)
def test_flush(tmpdir): tmpfn = '{}/traj.xtc'.format(tmpdir) data = np.random.random((5, 100, 3)) with XTCTrajectoryFile(tmpfn, 'w') as f: f.write(data) f.flush() # note that f is still open, so we can now try to read the contents flushed to disk. with XTCTrajectoryFile(tmpfn, 'r') as f2: out = f2.read() np.testing.assert_allclose(out[0], data, atol=1E-3)
def test_xtc_write_wierd_0(tmpdir): x0 = np.asarray(np.random.randn(100, 3, 3), dtype=np.float32) x1 = np.asarray(np.random.randn(100, 9, 3), dtype=np.float32) tmpfn = '{}/traj.xtc'.format(tmpdir) with XTCTrajectoryFile(tmpfn, 'w') as f: f.write(x0) with pytest.raises(ValueError): f.write(x1) xr = XTCTrajectoryFile(tmpfn).read()[0] print(xr.shape)
def test_write_0(): with XTCTrajectoryFile(fn_xtc) as f: xyz = f.read()[0] f = XTCTrajectoryFile(temp, 'w') f.write(xyz) f.close() with XTCTrajectoryFile(temp) as f: xyz2, time2, step2, box2 = f.read() eq(xyz, xyz2)
def test_write_1(): xyz = np.asarray(np.around(np.random.randn(100, 10, 3), 3), dtype=np.float32) time = np.asarray(np.random.randn(100), dtype=np.float32) step = np.arange(100) box = np.asarray(np.random.randn(100,3,3), dtype=np.float32) with XTCTrajectoryFile(temp, 'w') as f: f.write(xyz, time=time, step=step, box=box) with XTCTrajectoryFile(temp) as f: xyz2, time2, step2, box2 = f.read() eq(xyz, xyz2) eq(time, time2) eq(step, step2) eq(box, box2)
def test_write_0(tmpdir, fn_xtc): with XTCTrajectoryFile(fn_xtc) as f: xyz = f.read()[0] tmpfn = '{}/traj.xtc'.format(tmpdir) f = XTCTrajectoryFile(tmpfn, 'w') f.write(xyz) f.close() with XTCTrajectoryFile(tmpfn) as f: xyz2, time2, step2, box2 = f.read() eq(xyz, xyz2)
def test_write_error_0(tmpdir): xyz = np.asarray(np.random.randn(100, 3, 3), dtype=np.float32) tmpfn = '{}/traj.xtc'.format(tmpdir) with XTCTrajectoryFile(tmpfn, 'w') as f: with pytest.raises(ValueError): f.read(xyz)
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_read_atomindices_1(get_fn, fn_xtc): iofile = io.loadh(get_fn('frame0.xtc.h5'), deferred=False) with XTCTrajectoryFile(fn_xtc) as f: xyz, time, step, box = f.read(atom_indices=[0, 1, 2]) assert eq(xyz, iofile['xyz'][:, [0, 1, 2]]) assert eq(step, iofile['step']) assert eq(box, iofile['box']) assert eq(time, iofile['time'])
def test_write_2(tmpdir): xyz = np.asarray(np.around(np.random.randn(100, 10, 3), 3), dtype=np.float32) time = np.asarray(np.random.randn(100), dtype=np.float32) step = np.arange(100) box = np.asarray(np.random.randn(100, 3, 3), dtype=np.float32) tmpfn = '{}/traj.xtc'.format(tmpdir) with XTCTrajectoryFile(tmpfn, 'w') as f: for i in range(len(xyz)): f.write(xyz[i], time=time[i], step=step[i], box=box[i]) with XTCTrajectoryFile(tmpfn) as f: xyz2, time2, step2, box2 = f.read() eq(xyz, xyz2) eq(time, time2) eq(step, step2) eq(box, box2)
def _nc_a_traj(info, gen, gen_fn, has_overlapping_frames): out_fn = "{outdir}/{gen}.{outext}".format(gen=gen, **info['cnv2']) log.debug("Converting to netcdf {} {}".format(gen_fn, out_fn)) with XTCTrajectoryFile(gen_fn, 'r') as xtc: with NetCDFTrajectoryFile(out_fn, 'w') as nc: _nc_a_chunk(xtc, nc, has_overlapping_frames) return out_fn
def test_read_atomindices_2(): iofile = io.loadh(get_fn('frame0.xtc.h5'), deferred=False) with XTCTrajectoryFile(fn_xtc) as f: xyz, time, step, box = f.read(atom_indices=slice(None, None, 2)) yield lambda: eq(xyz, iofile['xyz'][:, ::2]) yield lambda: eq(step, iofile['step']) yield lambda: eq(box, iofile['box']) yield lambda: eq(time, iofile['time'])
def test_read_stride_2(): "read xtc with stride with n_frames" iofile = io.loadh(get_fn('frame0.xtc.h5'), deferred=False) with XTCTrajectoryFile(fn_xtc) as f: xyz, time, step, box = f.read(n_frames=1000, stride=3) yield lambda: eq(xyz, iofile['xyz'][::3]) yield lambda: eq(step, iofile['step'][::3]) yield lambda: eq(box, iofile['box'][::3]) yield lambda: eq(time, iofile['time'][::3])
def test_ragged_1(): # try first writing no box vectors,, and then adding some xyz = np.random.randn(100, 5, 3) time = np.random.randn(100) box = np.random.randn(100, 3, 3) with XTCTrajectoryFile(temp, 'w', force_overwrite=True) as f: f.write(xyz) assert_raises(ValueError, lambda: f.write(xyz, time, box))
def test_tell(): with XTCTrajectoryFile(get_fn('frame0.xtc')) as f: eq(f.tell(), 0) f.read(101) eq(f.tell(), 101) f.read(3) eq(f.tell(), 104)
def test_read_chunk2(): with XTCTrajectoryFile(fn_xtc, 'r', chunk_size_multiplier=1) as f: xyz, time, step, box = f.read() iofile = io.loadh(get_fn('frame0.xtc.h5'), deferred=False) yield lambda: eq(xyz, iofile['xyz']) yield lambda: eq(step, iofile['step']) yield lambda: eq(box, iofile['box']) yield lambda: eq(time, iofile['time'])
def test_read_chunk3(get_fn, fn_xtc): with XTCTrajectoryFile(fn_xtc, chunk_size_multiplier=2) as f: xyz, time, step, box = f.read(n_frames=100) iofile = io.loadh(get_fn('frame0.xtc.h5'), deferred=False) assert eq(xyz, iofile['xyz'][:100]) assert eq(step, iofile['step'][:100]) assert eq(box, iofile['box'][:100]) assert eq(time, iofile['time'][:100])
def test_read_stride(get_fn, fn_xtc): # read xtc with stride iofile = io.loadh(get_fn('frame0.xtc.h5'), deferred=False) with XTCTrajectoryFile(fn_xtc) as f: xyz, time, step, box = f.read(stride=3) assert eq(xyz, iofile['xyz'][::3]) assert eq(step, iofile['step'][::3]) assert eq(box, iofile['box'][::3]) assert eq(time, iofile['time'][::3])
def test_ragged_2(): # try first writing no box vectors, and then adding some xyz = np.random.randn(100, 5, 3) time = np.random.randn(100) box = np.random.randn(100, 3, 3) #from mdtraj.formats import HDF5TrajectoryFile with XTCTrajectoryFile(temp, 'w', force_overwrite=True) as f: f.write(xyz, time=time, box=box) assert_raises(ValueError, lambda: f.write(xyz))
def test_read_stride_n_frames(get_fn, fn_xtc): # read xtc with stride with n_frames iofile = io.loadh(get_fn('frame0.xtc.h5'), deferred=False) for s in strides: with XTCTrajectoryFile(fn_xtc) as f: xyz, time, step, box = f.read(n_frames=1000, stride=s) assert eq(xyz, iofile['xyz'][::s]) assert eq(step, iofile['step'][::s]) assert eq(box, iofile['box'][::s]) assert eq(time, iofile['time'][::s])
def test_read_atomindices_w_stride(get_fn, fn_xtc): # test case for bug: https://github.com/mdtraj/mdtraj/issues/1394 iofile = io.loadh(get_fn('frame0.xtc.h5'), deferred=False) for stride in strides: with XTCTrajectoryFile(fn_xtc) as f: xyz, time, step, box = f.read(atom_indices=[0, 1, 2], stride=stride) assert eq(xyz, iofile['xyz'][:, [0, 1, 2]][::stride]) assert eq(step, iofile['step'][::stride]) assert eq(box, iofile['box'][::stride]) assert eq(time, iofile['time'][::stride])
def test_seek_natoms9(): # create a xtc file with 9 atoms and seek it. with XTCTrajectoryFile(get_fn('frame0.xtc'), 'r') as fh: xyz = fh.read()[0][:, :9, :] with XTCTrajectoryFile(temp, 'w', force_overwrite=True) as f: f.write(xyz) with XTCTrajectoryFile(temp, 'r') as f: eq(f.read(1)[0].shape, (1, 9, 3)) eq(f.tell(), 1) f.seek(99) eq(f.read(1)[0].squeeze(), xyz[99]) # seek relative f.seek(-1, 1) eq(f.read(1)[0].squeeze(), xyz[99]) f.seek(0, 0) eq(f.read(1)[0].squeeze(), xyz[0])
def dump_xtc(self, output_name = Default.TRAJECTORY): """ Extract XYZ data from MONGO database and save as XTC trajectory file. """ #1. Connect to MONGO @ host:port with pymongo.MongoClient(self.db_host(), int(self.db_port())) as server,\ XTCTrajectoryFile(output_name, 'w') as xtc: #2. Define necessary variables temperature = 'structures{t_id}'.format(t_id = self.db_t_id()) data = server[self.db_name()][temperature] #3. Recover necessary data from MONGO self.update_terminal("Recovering data from database '{name}' @ {host}:{port} for temperature ID {t_id} ..."\ .format(name = self.db_name(), host = self.db_host(), port = self.db_port(), t_id = self.db_t_id())) query = {'xyz': 1, '_id': 0} if self.plot_energ(): query['energy'] = 1 query['step'] = 1 self.energy = {} frames = data.find({}, query) #4. Apply FROM and TO (crop data) if int(self.crop.start()) > 0: if int(self.crop.start()) > frames.count(): self.error(Default.INPUT_ERROR, "Minimum crop value is {max}".format(max = frames.count())) frames.skip(int(self.crop.start())) if int(self.crop.end()) > 0: if int(self.crop.end()) > frames.count(): self.error(Default.INPUT_ERROR, "Maximum crop value is {max}".format(max = frames.count())) frames.limit(int(self.crop.end()) - int(self.crop.start())) self.update_progress_bar() #5. Write XTC file, applying PRINT EVERY settings, and WITHOUT X's self.update_terminal("Writting XTC trajectory to file {output_name} ..."\ .format(output_name = output_name)) size = frames.count(with_limit_and_skip = True) for n, frame in enumerate(frames): if self.debug_mode(): print "{:>5d}/{:>5d}".format(n, size) if n % int(self.every()) == 0: #5.1 Extract X's frame_without_Xs = [] for index, atom in enumerate(frame['xyz'], start = 1): if index % Default.ATOMS_PER_RESI == 0 and index > 1: continue frame_without_Xs.append(atom) #5.2 Write XTC file xtc.write(0.1*np.array(frame_without_Xs), step=n, time=n) #5.3 Save energy values if energy plots is requested if self.plot_energ(): self.energy[int(frame['step']) - int(self.crop.start())] = float(frame['energy']) self.update_progress_bar()
def test_read_stride_n_frames_offsets(get_fn, fn_xtc): # read xtc with stride with n_frames and offsets iofile = io.loadh(get_fn('frame0.xtc.h5'), deferred=False) for s in (1, 2, 3, 4, 5): with XTCTrajectoryFile(fn_xtc) as f: f.offsets # pre-compute byte offsets between frames xyz, time, step, box = f.read(n_frames=1000, stride=s) assert eq(xyz, iofile['xyz'][::s]) assert eq(step, iofile['step'][::s]) assert eq(box, iofile['box'][::s]) assert eq(time, iofile['time'][::s])
def regroup_DISK(trajs, topology_file, disctrajs, path, stride=1): """Regroups MD trajectories into clusters according to discretised trajectories. Parameters ---------- trajs : list of strings xtc/dcd/... trajectory file names topology_file : string name of topology file that matches `trajs` disctrajs : list of array-likes discretized trajectories path : string file system path to directory where cluster trajectories are written stride : int stride of disctrajs with respect to the (original) trajs Returns ------- cluster : list of file names or `None`, len(cluster)=np.max(trajs)+1 each element cluster[i] is either `None` if i wasn't found in disctrajs or is a the file name of a new trajectory that holds all frames that were assigned to cluster i. """ # handle single element invocation if not isinstance(trajs, list): trajs = [trajs] if not isinstance(disctrajs, list): disctrajs = [disctrajs] states = np.unique( np.hstack(([np.unique(disctraj) for disctraj in disctrajs]))) states = np.setdiff1d(states, [-1]) # exclude invalid states writer = [None] * (max(states) + 1) cluster = [None] * (max(states) + 1) for i in states: cluster[i] = path + os.sep + ('%d.xtc' % i) writer[i] = XTCTrajectoryFile(cluster[i], 'w', force_overwrite=True) for disctraj, traj in zip(disctrajs, trajs): reader = md.iterload(traj, top=topology_file, stride=stride) start = 0 for chunk in reader: chunk_length = chunk.xyz.shape[0] for i in xrange(chunk_length): cl = disctraj[i + start] if cl != -1: writer[cl].write(chunk.xyz[i, :, :]) # np.newaxis? start += chunk_length # TODO: check that whole disctrajs was used for i in states: writer[i].close() return cluster
def test_ragged_2(tmpdir): # try first writing no box vectors, and then adding some xyz = np.random.randn(100, 5, 3) time = np.random.randn(100) box = np.random.randn(100, 3, 3) tmpfn = '{}/traj.xtc'.format(tmpdir) with XTCTrajectoryFile(tmpfn, 'w', force_overwrite=True) as f: f.write(xyz, time=time, box=box) with pytest.raises(ValueError): f.write(xyz)
def main(argv): print('Importarray: ') print(trajectory.importarray) print('test for loading via gui:') print(trajectory.get_trajectory(argv[0], '')) print(trajectory.importarray) print('test for loading with command + struc&traj:') print(trajectory.get_trajectory(argv[0], argv[1])) print(trajectory.importarray) print('Extratest: Import mdtraj') try: import mdtraj print('success') except ImportError as e: print('error', e) print('Extratest: Import mdanalysis') try: import mdanalysis print('success') except ImportError as e: print('error', e) print('Extratest: mdtraj functionality') try: from mdtraj.formats import XTCTrajectoryFile xtc = XTCTrajectoryFile(argv[0], 'r') xtc_traj = xtc.read(n_frames=1) print('success:', xtc_traj) except Exception as ex: template = "An exception of type {0} occurred. Arguments:\n{1!r}" message = template.format(type(ex).__name__, ex.args) print(message, ex)