Refactor pmda after universe can be serialized

pmda/contacts.py

@@ -237,7 +237,8 @@ def __init__(self,
             respective functions for reasonable values.
         """
         universe = mobiles[0].universe
-        super(Contacts, self).__init__(universe, mobiles)
+        super().__init__(universe)
+        self._mobiles = mobiles
 
         if method == 'hard_cut':
             self.fraction_contacts = hard_cut_q
@@ -267,13 +268,13 @@ def __init__(self,
     def _prepare(self):
         self.timeseries = None
 
-    def _single_frame(self, ts, atomgroups):
-        grA, grB = atomgroups
+    def _single_frame(self):
+        grA, grB = self._mobiles
         # compute distance array for a frame
         d = distance_array(grA.positions, grB.positions)
 
         y = np.empty(len(self.r0) + 1)
-        y[0] = ts.frame
+        y[0] = self._ts.frame
         for i, (initial_contacts,
                 r0) in enumerate(zip(self.initial_contacts, self.r0)):
             # select only the contacts that were formed in the reference state

pmda/custom.py

@@ -16,8 +16,6 @@
 same universe and return a value.
 
 """
-from __future__ import absolute_import
-
 from MDAnalysis.core.groups import AtomGroup
 from MDAnalysis.core.universe import Universe
 from MDAnalysis.coordinates.base import ProtoReader
@@ -77,27 +75,15 @@ def __init__(self, function, universe, *args, **kwargs):
         """
 
         self.function = function
-
-        # collect all atomgroups with the same trajectory object as universe
-        trajectory = universe.trajectory
-        arg_ags = []
-        self.other_args = []
-        for arg in args:
-            if isinstance(arg,
-                          AtomGroup) and arg.universe.trajectory == trajectory:
-                arg_ags.append(arg)
-            else:
-                self.other_args.append(arg)
-
-        super(AnalysisFromFunction, self).__init__(universe, arg_ags)
+        super().__init__(universe)
+        self.args = args
         self.kwargs = kwargs
 
     def _prepare(self):
         self.results = []
 
-    def _single_frame(self, ts, atomgroups):
-        args = atomgroups + self.other_args
-        return self.function(*args, **self.kwargs)
+    def _single_frame(self):
+        return self.function(*self.args, **self.kwargs)
 
     def _conclude(self):
         self.results = np.concatenate(self._results)

pmda/density.py

@@ -22,9 +22,6 @@
 MDAnalysis.analysis.density.density_from_Universe
 
 """
-
-from __future__ import absolute_import
-
 import numpy as np
 
 from MDAnalysis.lib.util import fixedwidth_bins
@@ -241,7 +238,7 @@ def __init__(self, atomgroup, delta=1.0, atomselection=None,
                  parameters=None, gridcenter=None, xdim=None, ydim=None,
                  zdim=None):
         u = atomgroup.universe
-        super(DensityAnalysis, self).__init__(u, (atomgroup, ))
+        super().__init__(u)
         self._atomgroup = atomgroup
         self._delta = delta
         self._atomselection = atomselection
@@ -259,10 +256,15 @@ def __init__(self, atomgroup, delta=1.0, atomselection=None,
         elif not updating and atomselection is not None:
             raise ValueError("""With updating=False, the atomselection='{}' is
                         not used and should be None""".format(atomselection))
+        elif updating and atomselection is not None:
+            self._select_atomgroup = atomgroup.select_atoms(atomselection,
+                                                            updating=True)
+        else:
+            self._select_atomgroup = atomgroup
+
 
     def _prepare(self):
-        coord = self.current_coordinates(self._atomgroup, self._atomselection,
-                                         self._updating)
+        coord = self._select_atomgroup.positions
         if self._gridcenter is not None:
             # Generate a copy of smin/smax from coords to later check if the
             # defined box might be too small for the selection
@@ -294,10 +296,10 @@ def _prepare(self):
         self._arange = arange
         self._bins = bins
 
-    def _single_frame(self, ts, atomgroups):
-        coord = self.current_coordinates(atomgroups[0], self._atomselection,
-                                         self._updating)
-        result = np.histogramdd(coord, bins=self._bins, range=self._arange,
+    def _single_frame(self):
+        result = np.histogramdd(self._select_atomgroup.positions,
+                                bins=self._bins,
+                                range=self._arange,
                                 normed=False)
         return result[0]
 
@@ -330,12 +332,3 @@ def _reduce(res, result_single_frame):
         else:
             res += result_single_frame
         return res
-
-    @staticmethod
-    def current_coordinates(atomgroup, atomselection, updating):
-        """Retrieves the current coordinates of all atoms in the chosen atom
-        selection.
-        Note: currently required to allow for updating selections"""
-        ag = atomgroup if not updating else atomgroup.select_atoms(
-                                                            atomselection)
-        return ag.positions

pmda/hbond_analysis.py

@@ -24,8 +24,6 @@
    :members:
 
 """
-from __future__ import absolute_import, division
-
 import numpy as np
 
 from MDAnalysis.lib.distances import capped_distance, calc_angles
@@ -172,7 +170,8 @@ def __init__(self, universe, donors_sel=None, hydrogens_sel=None,
         """
 
         ag = universe.atoms
-        super(HydrogenBondAnalysis, self).__init__(universe, (ag, ))
+        super().__init__(universe)
+        self._atomgroup = ag
         self.donors_sel = donors_sel
         self.hydrogens_sel = hydrogens_sel
         self.acceptors_sel = acceptors_sel
@@ -218,7 +217,7 @@ def guess_hydrogens(self, selection='all', max_mass=1.1, min_charge=0.3):
         :attr:`hydrogens_sel`.
         """
 
-        u = self._universe()
+        u = self._universe
         ag = u.select_atoms(selection)
         hydrogens_ag = ag[
             np.logical_and(
@@ -282,7 +281,7 @@ def guess_donors(self, selection='all', max_charge=-0.5):
             hydrogens_sel = self.guess_hydrogens()
         else:
             hydrogens_sel = self.hydrogens_sel
-        u = self._universe()
+        u = self._universe
         hydrogens_ag = u.select_atoms(hydrogens_sel)
 
         ag = hydrogens_ag.residues.atoms.select_atoms(
@@ -335,7 +334,7 @@ def guess_acceptors(self, selection='all', max_charge=-0.5):
         attribute :attr:`acceptors_sel`.
         """
 
-        u = self._universe()
+        u = self._universe
         ag = u.select_atoms(selection)
         acceptors_ag = ag[ag.charges < max_charge]
         acceptors_list = np.unique(
@@ -393,7 +392,7 @@ def _get_dh_pairs(self, u):
         return donors, hydrogens
 
     def _prepare(self):
-        u = mda.Universe(self._top, self._traj)
+        u = self._universe
         self.hbonds = []
         self.frames = np.arange(self.start, self.stop, self.step)
         self.timesteps = (self.frames*u.trajectory.dt) + u.trajectory[0].time
@@ -410,10 +409,9 @@ def _prepare(self):
         self._donors_ids = donors.ids
         self._hydrogens_ids = hydrogens.ids
 
-    def _single_frame(self, ts, atomgroups):
-        u = atomgroups[0].universe
-
-        box = ts.dimensions
+    def _single_frame(self):
+        u = self._universe
+        box = self._ts.dimensions
 
         # Update donor-hydrogen pairs if necessary
         if self.update_selections:
@@ -459,7 +457,7 @@ def _single_frame(self, ts, atomgroups):
 
         # Store data on hydrogen bonds found at this frame
         hbonds = [[], [], [], [], [], []]
-        hbonds[0].extend(np.full_like(hbond_donors, ts.frame))
+        hbonds[0].extend(np.full_like(hbond_donors, self._ts.frame))
         hbonds[1].extend(hbond_donors.ids)
         hbonds[2].extend(hbond_hydrogens.ids)
         hbonds[3].extend(hbond_acceptors.ids)
@@ -508,7 +506,7 @@ def count_by_type(self):
         resname and atom type of the donor and acceptor atoms in a hydrogen
         bond.
         """
-        u = self._universe()
+        u = self._universe
         d = u.atoms[self.hbonds[:, 1].astype(np.int)]
         a = u.atoms[self.hbonds[:, 3].astype(np.int)]
 
@@ -541,7 +539,7 @@ def count_by_ids(self):
         hydrogen atom id and acceptor atom id in a hydrogen bond.
         """
 
-        u = self._universe()
+        u = self._universe
         d = u.atoms[self.hbonds[:, 1].astype(np.int)]
         h = u.atoms[self.hbonds[:, 2].astype(np.int)]
         a = u.atoms[self.hbonds[:, 3].astype(np.int)]
@@ -558,14 +556,6 @@ def count_by_ids(self):
 
         return unique_hbonds
 
-    def _universe(self):
-        # A Universe containing position information is needed for guessing
-        # donors and acceptors.
-        u = mda.Universe(self._top)
-        if not hasattr(u.atoms, 'positions'):
-            u.load_new(self._positions)
-        return u
-
     @staticmethod
     def _reduce(res, result_single_frame):
         """ Use numpy array append to combine results"""