pickle universe as well

MDAnalysis · Sep 20, 2018 · 61df0c9 · 61df0c9
1 parent ef91190
commit 61df0c9
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/pmda/parallel.py b/pmda/parallel.py
@@ -186,8 +186,7 @@ def __init__(self, universe, atomgroups):
 
         """
         self._trajectory = universe.trajectory
-        self._top = universe.filename
-        self._traj = universe.trajectory.filename
+        self._universe_dump = pickle.dumps(universe)
         self._pickles = [pickle.dumps(ag) for ag in atomgroups]
 
     @contextmanager
@@ -329,8 +328,7 @@ def run(self,
                             min(stop, (b + 1) * bsize * step + start),
                             step,
                             self._pickles,
-                            self._top,
-                            self._traj, )
+                            self._universe_dump)
                     blocks.append(task)
                 blocks = delayed(blocks)
                 res = blocks.compute(**scheduler_kwargs)
@@ -344,10 +342,10 @@ def run(self,
             np.array([el[3] 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, universe_dump):
         """helper function to actually setup dask graph"""
         with timeit() as b_universe:
-            u = mda.Universe(top, traj)
+            u = pickle.loads(universe_dump)
             agroups = [pickle.loads(idx) for idx in pickles]
 
         res = []