Turn ParallelAnalysisBase into dask custom collection

.travis.yml

@@ -18,15 +18,15 @@ env:
     - SETUP_CMD="pmda --pep8 --cov pmda"
     # mdanalysis develop from source (see below), which needs
     # minimal CONDA_MDANALYSIS_DEPENDENCIES
-    #- CONDA_DEPENDENCIES="mdanalysis mdanalysistests dask joblib pytest-pep8 mock codecov cython hypothesis sphinx"
-    #- CONDA_MDANALYSIS_DEPENDENCIES="cython mmtf-python six biopython networkx scipy griddataformats gsd hypothesis"
-    - CONDA_MDANALYSIS_DEPENDENCIES="mdanalysis mdanalysistests"
-    - CONDA_DEPENDENCIES="${CONDA_MDANALYSIS_DEPENDENCIES} dask distributed joblib pytest-pep8 mock codecov"
+    - CONDA_DEPENDENCIES="mdanalysis mdanalysistests dask joblib pytest-pep8 mock codecov cython hypothesis sphinx"
+    - CONDA_MDANALYSIS_DEPENDENCIES="cython mmtf-python six biopython networkx scipy griddataformats gsd hypothesis"
+    # - CONDA_MDANALYSIS_DEPENDENCIES="mdanalysis mdanalysistests"
+    # - CONDA_DEPENDENCIES="${CONDA_MDANALYSIS_DEPENDENCIES} dask distributed joblib pytest-pep8 mock codecov"
     - CONDA_CHANNELS='conda-forge'
     - CONDA_CHANNEL_PRIORITY=True
     # install development version of MDAnalysis (needed until the test
     # files for analysis.rdf are available in release 0.19.0)
-    #- PIP_DEPENDENCIES="git+https://github.com/MDAnalysis/mdanalysis#egg=mdanalysis&subdirectory=package git+https://github.com/MDAnalysis/mdanalysis#egg=mdanalysistests&subdirectory=testsuite"
+    - PIP_DEPENDENCIES="git+https://github.com/MDAnalysis/mdanalysis#egg=mdanalysis&subdirectory=package git+https://github.com/MDAnalysis/mdanalysis#egg=mdanalysistests&subdirectory=testsuite"
     - NUMPY_VERSION=stable
     - BUILD_CMD='python setup.py develop'
     - CODECOV=false

pmda/contacts.py

@@ -182,8 +182,6 @@ def is_any_closer(r, r0, dist=2.5):
    :inherited-members:
 
 """
-from __future__ import absolute_import, division
-
 import MDAnalysis as mda
 from MDAnalysis.analysis.contacts import (contact_matrix, hard_cut_q,
                                           radius_cut_q, soft_cut_q)
@@ -237,7 +235,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 +266,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,9 +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
 import numpy as np
@@ -44,7 +41,7 @@ class AnalysisFromFunction(ParallelAnalysisBase):
     >>>     return mda.analysis.align.rotation_matrix(mobile, ref)[0]
     >>> # now run an analysis using the standalone function
     >>> rot = AnalysisFromFunction(rotation_matrix,
-                                   trajectory, mobile, ref).run()
+                                   universe, mobile, ref).run()
     >>> print(rot.results)
 
     Raises
@@ -64,40 +61,24 @@ def __init__(self, function, universe, *args, **kwargs):
             to be 'mobile' Atomgroups if they belong to the same universe. All
             other Atomgroups are assumed to be reference. Here 'mobile' means
             they will be iterated over.
-        Universe : :class:`~MDAnalysis.core.groups.Universe`
+        universe : :class:`~MDAnalysis.core.groups.Universe`
             a :class:`MDAnalysis.core.groups.Universe` (the
-            `atomgroups` must belong to this Universe)
+            `atomgroups` in other args must belong to this Universe)
         *args : list
            arguments for ``function``
         **kwargs : dict
-           keyword arguments for ``function``. keyword arguments with name
-           'universe' or 'atomgroups' will be ignored! Mobile atomgroups to
-           analyze can not be passed as keyword arguments currently.
-
+           keyword arguments for ``function``.
         """
-
         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)
@@ -132,7 +113,7 @@ def analysis_class(function):
     >>> def RotationMatrix(mobile, ref):
     >>>     return mda.analysis.align.rotation_matrix(mobile, ref)[0]
 
-    >>> rot = RotationMatrix(u.trajectory, mobile, ref, step=2).run()
+    >>> rot = RotationMatrix(u, mobile, ref, step=2).run()
     >>> print(rot.results)
 
     See Also
@@ -144,13 +125,12 @@ def analysis_class(function):
     class WrapperClass(AnalysisFromFunction):
         """Custom Analysis Function"""
 
-        def __init__(self, trajectory=None, *args, **kwargs):
-            if not (isinstance(trajectory, ProtoReader) or isinstance(
-                    trajectory, Universe)):
-                print(type(trajectory))
+        def __init__(self, universe=None, *args, **kwargs):
+            if not isinstance(universe, Universe):
+                print(type(universe))
                 raise ValueError(
-                    "First argument needs to be an MDAnalysis reader object.")
-            super(WrapperClass, self).__init__(function, trajectory, *args,
+                    "First argument needs to be an MDAnalysis Universe.")
+            super().__init__(function, universe, *args,
                                                **kwargs)
 
     return WrapperClass

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,14 @@ 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 +295,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 +331,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"""