diff --git a/pmda/parallel.py b/pmda/parallel.py index cce6c0af..fb3409f2 100644 --- a/pmda/parallel.py +++ b/pmda/parallel.py @@ -191,6 +191,7 @@ def __init__(self, universe, atomgroups): self._trajectory = universe.trajectory self._top = universe.filename self._traj = universe.trajectory.filename + self._anchor = universe.anchor_name self._pickles = [pickle.dumps(ag) for ag in atomgroups] @contextmanager @@ -333,7 +334,8 @@ def run(self, step, self._pickles, self._top, - self._traj, ) + self._traj, + self._anchor) blocks.append(task) blocks = delayed(blocks) res = blocks.compute(**scheduler_kwargs) @@ -347,10 +349,11 @@ def run(self, np.array([el.timing_universe for el in res]), time_prepare, conclude.elapsed) return self - def _dask_helper(self, start, stop, step, pickles, top, traj): + def _dask_helper(self, start, stop, step, pickles, top, traj, anchor): """helper function to actually setup dask graph""" with timeit() as b_universe: u = mda.Universe(top, traj) + u.anchor_name = anchor agroups = [pickle.loads(idx) for idx in pickles] res = []