Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_write_1(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:
f.write(xyz, time=time, step=step, box=box)
with XTCTrajectoryFile(tmpfn) as f:
xyz2, time2, step2, box2 = f.read()
eq(xyz, xyz2)
eq(time, time2)
eq(step, step2)
eq(box, box2)
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_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_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_short_traj(tmpdir):
tmpfn = '{}/traj.xtc'.format(tmpdir)
with XTCTrajectoryFile(tmpfn, 'w') as f:
f.write(np.random.uniform(size=(5, 100000, 3)))
with XTCTrajectoryFile(tmpfn, 'r') as f:
assert len(f) == 5, len(f)
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_read_chunk1(get_fn, fn_xtc):
with XTCTrajectoryFile(fn_xtc, 'r', chunk_size_multiplier=0.5) as f:
xyz, time, step, box = f.read()
iofile = io.loadh(get_fn('frame0.xtc.h5'), deferred=False)
assert eq(xyz, iofile['xyz'])
assert eq(step, iofile['step'])
assert eq(box, iofile['box'])
assert eq(time, iofile['time'])
def _regroup_DISK_subset(states, trajs, topology_file, disctrajs, path, stride):
writer = [None] * (max(states) + 1)
out_fnames = []
for i in states:
out_fname = path + os.sep + ('%d.xtc' % i)
out_fnames.append(out_fname)
writer[i] = XTCTrajectoryFile(out_fname, '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 in states:
writer[cl].write(chunk.xyz[i, :, :])
start += chunk_length
for i in states:
writer[i].close()
return out_fnames
def save_xtc(self, filename, force_overwrite=True):
"""Save trajectory to Gromacs XTC format
Parameters
----------
filename : str
filesystem path in which to save the trajectory
force_overwrite : bool, default=True
Overwrite anything that exists at filename, if its already there
"""
with XTCTrajectoryFile(filename, 'w', force_overwrite=force_overwrite) as f:
f.write(xyz=in_units_of(self.xyz, Trajectory._distance_unit, f.distance_unit),
time=self.time,
box=in_units_of(self.unitcell_vectors, Trajectory._distance_unit, f.distance_unit))
def save_xtc(self, filename, force_overwrite=True):
"""Save trajectory to Gromacs XTC format
Parameters
----------
filename : str
filesystem path in which to save the trajectory
force_overwrite : bool, default=True
Overwrite anything that exists at filename, if its already there
"""
with XTCTrajectoryFile(filename, 'w', force_overwrite=force_overwrite) as f:
f.write(xyz=self.xyz, time=self.time, box=self.unitcell_vectors)