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)
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)
