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_iterload_skip(ref_traj, get_fn):
if ref_traj.fobj is md.formats.PDBTrajectoryFile:
pytest.xfail("PDB Iterloads an extra frame!!")
if ref_traj.fobj is md.formats.GroTrajectoryFile:
pytest.xfail("Not implemented for some reason")
if ref_traj.fext in ('ncrst', 'rst7'):
pytest.skip("Only 1 frame per file format")
top = md.load(get_fn('native.pdb'))
t_ref = md.load(get_fn(ref_traj.fn), top=top)
for cs in [0, 1, 11, 100]:
for skip in [0, 1, 20, 101]:
t = functools.reduce(lambda a, b: a.join(b),
md.iterload(get_fn(ref_traj.fn), skip=skip, top=top, chunk=cs))
eq(t_ref.xyz[skip:], t.xyz)
eq(t_ref.time[skip:], t.time)
eq(t_ref.topology, t.topology)
def iterload(self, i, chunk):
if self.verbose:
print('[MDTraj dataset] iterloading %s' % self.filename(i))
if self._topology is None:
return md.iterload(
self.filename(i), chunk=chunk, stride=self.stride,
atom_indices=self.atom_indices)
else:
return md.iterload(
self.filename(i), chunk=chunk, stride=self.stride,
atom_indices=self.atom_indices, top=self._topology)
def load_data(self):
load_time_start = time.time()
data = []
for tfn in self.filenames:
kwargs = {} if tfn.endswith('h5') else {'top': self.top}
for t in md.iterload(tfn, chunk=self.args.split, **kwargs):
item = np.asarray(md.compute_dihedrals(t, self.indices), np.double)
data.append(item)
print('Loading data into memory + vectorization: %f s' %
(time.time() - load_time_start))
print('''Fitting with %s timeseries from %d trajectories with %d
total observations''' % (len(data), len(self.filenames),
sum(len(e) for e in data)))
return data
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 load_data(self):
load_time_start = time.time()
data = []
for tfn in self.filenames:
kwargs = {} if tfn.endswith('h5') else {'top': self.top}
for t in md.iterload(tfn, chunk=self.args.split, **kwargs):
features = self.featurizer.featurize(t)
data.append(features)
print('Loading data into memory + vectorization: %f s' %
(time.time() - load_time_start))
print('''Fitting with %s timeseries from %d trajectories with %d
total observations''' % (len(data), len(self.filenames),
sum(len(e) for e in data)))
return data
target_filename = chisurf.widgets.save_file('H5-Trajectory file', 'H5-File (*.h5)')
# target_filename = 'clash_dimer.h5'
filename = self.trajectory_filename
stride = self.stride
min_distance = self.min_distance
# Make empty trajectory
frame_0 = md.load_frame(filename, 0)
target_traj = md.Trajectory(xyz=np.empty((0, frame_0.n_atoms, 3)), topology=frame_0.topology)
#atom_indices = np.array(self.atom_list)
atom_selection = self.atom_list
atom_list = target_traj.top.select(atom_selection)
target_traj.save(target_filename)
chunk_size = 1000
for i, chunk in enumerate(md.iterload(filename, chunk=chunk_size, stride=stride)):
xyz = chunk.xyz.copy()
frames_below = below_min_distance(xyz, min_distance, atom_list=atom_list)
selection = np.where(frames_below < 1)[0]
xyz_clash_free = np.take(xyz, selection, axis=0)
with tables.open_file(target_filename, 'a') as table:
table.root.coordinates.append(xyz_clash_free)
times = np.arange(table.root.time.shape[0],
table.root.time.shape[0] + xyz_clash_free.shape[0], dtype=np.float32)
table.root.time.append(times)
traj_2 = md.load_frame(fn2, index=0)
# Create empty trajectory
if self.join_mode == 'time':
traj_join = traj_1.join(traj_2)
axis = 0
elif self.join_mode == 'atoms':
traj_join = traj_1.stack(traj_2)
axis = 1
target_traj = md.Trajectory(xyz=np.empty((0, traj_join.n_atoms, 3)), topology=traj_join.topology)
target_traj.save(target_filename)
chunk_size = self.chunk_size
table = tables.open_file(target_filename, 'a')
for i, (c1, c2) in enumerate(izip(md.iterload(fn1, chunk=chunk_size), md.iterload(fn2, chunk=chunk_size))):
xyz_1 = c1.xyz[::-1] if r1 else c1.xyz
xyz_2 = c2.xyz[::-1] if r2 else c2.xyz
xyz = np.concatenate((xyz_1, xyz_2), axis=axis)
table.root.coordinates.append(xyz)
table.root.time.append(np.arange(i * chunk_size, i * chunk_size + xyz.shape[0], dtype=np.float32))
table.close()
s = 'FrameNbr\t'
for p in self.universe.potentials:
s += '%s\t' % p.name
s += '\n'
chisurf.fio.zipped.open_maybe_zipped(
filename=energy_file,
mode='w'
).write(s)
self.structure = chisurf.structure.TrajectoryFile(
mdtraj.load_frame(
self.trajectory_file, 0
)
)[0]
i = 0
for chunk in mdtraj.iterload(self.trajectory_file):
for frame in chunk:
self.structure.xyz = frame.xyz * 10.0
self.structure.update_dist()
s = '%i\t' % (i * self.stride + 1)
for e in self.universe.getEnergies(self.structure):
s += '%.3f\t' % e
print(s)
s += '\n'
i += 1
open(energy_file, 'a').write(s)
def onSaveTrajectory(self, target_filename=None):
if target_filename is None:
target_filename = str(QtWidgets.QFileDialog.getSaveFileName(None, 'Save H5-Model file', '', 'H5-files (*.h5)'))[0]
filename = self.trajectory_filename
atom_indices = self.atom_list
stride = self.stride
# Make empty trajectory
frame_0 = md.load_frame(filename, 0)
target_traj = md.Trajectory(xyz=np.empty((0, frame_0.n_atoms, 3)), topology=frame_0.topology)
target_traj.save(target_filename)
chunk_size = 1000
table = tables.open_file(target_filename, 'a')
for i, chunk in enumerate(md.iterload(filename, chunk=chunk_size, stride=stride)):
chunk = chunk.superpose(frame_0, frame=0, atom_indices=atom_indices)
xyz = chunk.coordinates.copy()
table.root.coordinates.append(xyz)
table.root.time.append(np.arange(i * chunk_size, i * chunk_size + xyz.shape[0], dtype=np.float32))
table.close()