def test_pickle_progress_bar():
traj = pt.datafiles.load_tz2()
with tempfolder():
t2 = ProgressBarTrajectory(traj)
pt.to_pickle(t2, 'test.pk')
t3 = pt.read_pickle('test.pk')
data0 = pt.rmsd(traj)
data1 = pt.rmsd(t2)
data2 = pt.rmsd(t3)
aa_eq(data0, data1)
aa_eq(data0, data2)
def test_progress_log():
"""test_progress_log: simple test, just to make sure it is runnable
"""
traj = pt.datafiles.load_tz2()[:10]
p = ProgressBarTrajectory(traj, style='basic', every=5)
pt.molsurf(p)
p = ProgressBarTrajectory(traj, style='bar', every=5)
pt.molsurf(p)
p = ProgressBarTrajectory(traj, style='circle', every=5)
pt.molsurf(p)
# need to run on Jupyter notebook
if tqdm is not None:
p = ProgressBarTrajectory(traj, style='tqdm')
pt.molsurf(p)
try:
p = ProgressBarTrajectory(traj, style=tqdm_notebook)
pt.molsurf(p)
except traitlets.traitlets.TraitError:
pass
# make sure not loading all coordinates from TrajectoryIterator
traj2 = pt.iterload(fn('tz2.nc'), fn('tz2.parm7'))
traj2._size_limit_in_GB = traj2._estimated_GB - 0.001
with pytest.raises(MemoryError):
traj2.xyz
p2 = ProgressBarTrajectory(traj2)
aa_eq(pt.rmsd(p2), pt.rmsd(traj2))
def worker_by_func(rank,
n_cores=None,
func=None,
traj=None,
args=None,
kwargs=None,
iter_options=None,
apply=None,
progress=False,
progress_params=dict()):
'''worker for pytraj's functions
'''
# need to unpack args and kwargs
if iter_options is None:
iter_options = {}
mask = iter_options.get('mask')
rmsfit = iter_options.get('rmsfit')
autoimage = iter_options.get('autoimage', False)
iter_func = apply
frame_indices = kwargs.pop('frame_indices', None)
if frame_indices is None:
my_iter = traj._split_iterators(n_cores,
rank=rank,
mask=mask,
rmsfit=rmsfit,
autoimage=autoimage)
else:
my_indices = np.array_split(frame_indices, n_cores)[rank]
my_iter = traj.iterframe(frame_indices=my_indices,
mask=mask,
rmsfit=rmsfit,
autoimage=autoimage)
if progress and rank == 0:
from pytraj.utils.progress import ProgressBarTrajectory
my_iter = ProgressBarTrajectory(my_iter,
style='basic',
**progress_params)
n_frames = my_iter.n_frames
kwargs_cp = {}
kwargs_cp.update(kwargs)
if iter_func is not None:
final_iter = WrapBareIterator(iter_func(my_iter), top=my_iter.top)
else:
final_iter = my_iter
data = func(final_iter, *args, **kwargs_cp)
return (data, n_frames)