Trajectory slicing made completely pythonic MDAnalysis#918

package/MDAnalysis/analysis/base.py

@@ -27,16 +27,15 @@
 classes.
 
 """
-from six.moves import range, zip
-
 import inspect
 import logging
+
 import numpy as np
 import six
-
 from MDAnalysis import coordinates
 from MDAnalysis.core.groups import AtomGroup
 from MDAnalysis.lib.log import ProgressMeter, _set_verbose
+from six.moves import range, zip
 
 logger = logging.getLogger(__name__)
 
@@ -223,7 +222,7 @@ def __init__(self, function, trajectory=None, *args, **kwargs):
         """
         if (trajectory is not None) and (not isinstance(
                 trajectory, coordinates.base.Reader)):
-            args = args + (trajectory, )
+            args = args + (trajectory,)
             trajectory = None
 
         if trajectory is None:

package/MDAnalysis/coordinates/base.py

@@ -124,28 +124,27 @@
 """
 from __future__ import absolute_import
 
-from six.moves import range
-import six
-
+import copy
 import warnings
+import weakref
 
 import numpy as np
-import copy
-import weakref
+import six
+from six.moves import range
 
+from . import core
+from .. import NoDataError
 from .. import (
     _READERS,
     _SINGLEFRAME_WRITERS,
     _MULTIFRAME_WRITERS,
 )
-from ..core import flags
 from .. import units
-from ..lib.util import asiterable, Namespace
-from . import core
-from .. import NoDataError
-
 from ..auxiliary.base import AuxReader
 from ..auxiliary.core import auxreader
+from ..core import flags
+from ..lib.util import asiterable, Namespace
+
 
 class Timestep(object):
     """Timestep data for one frame
@@ -1191,6 +1190,7 @@ def __getitem__(self, frame):
         ----
         *frame* is a 0-based frame index.
         """
+
         def apply_limits(frame):
             if frame < 0:
                 frame += len(self)
@@ -1214,6 +1214,7 @@ def listiter(frames):
                     if not isinstance(f, (int, np.integer)):
                         raise TypeError("Frames indices must be integers")
                     yield self._read_frame_with_aux(apply_limits(f))
+
             return listiter(frame)
         elif isinstance(frame, slice):
             start, stop, step = self.check_slice_indices(
@@ -1306,22 +1307,23 @@ def check_slice_indices(self, start, stop, step):
 
         if step > 0 and stop < start:
             raise IndexError("Stop frame is lower than start frame")
+        elif step > 0 and stop >= nframes:
+            stop = nframes
         elif step < 0 and start < stop:
             raise IndexError("Start frame is lower than stop frame")
-        if not (0 <= start < nframes) or stop > nframes:
-            raise IndexError(
-                "Frame start/stop outside of the range of the trajectory.")
+        elif step < 0 and start >= nframes:
+            start = nframes - 1
 
         return start, stop, step
 
     def __repr__(self):
         return ("<{cls} {fname} with {nframes} frames of {natoms} atoms>"
                 "".format(
-                    cls=self.__class__.__name__,
-                    fname=self.filename,
-                    nframes=self.n_frames,
-                    natoms=self.n_atoms
-                ))
+            cls=self.__class__.__name__,
+            fname=self.filename,
+            nframes=self.n_frames,
+            natoms=self.n_atoms
+        ))
 
     def add_auxiliary(self, auxname, auxdata, format=None, **kwargs):
         """Add auxiliary data to be read alongside trajectory.
@@ -1421,7 +1423,7 @@ def next_as_aux(self, auxname):
         aux = self._check_for_aux(auxname)
         ts = self.ts
         # catch up auxiliary if it starts earlier than trajectory
-        while aux.step_to_frame(aux.step+1, ts) < 0:
+        while aux.step_to_frame(aux.step + 1, ts) < 0:
             next(aux)
         # find the next frame that'll have a representative value
         next_frame = aux.next_nonempty_frame(ts)
@@ -1555,7 +1557,6 @@ def rename_aux(self, auxname, new):
         setattr(self.ts.aux, new, self.ts.aux[auxname])
         delattr(self.ts.aux, auxname)
 
-
     def get_aux_descriptions(self, auxnames=None):
         """Get descriptions to allow reloading the specified auxiliaries.
 
@@ -1586,7 +1587,6 @@ def get_aux_descriptions(self, auxnames=None):
         return descriptions
 
 
-
 class Reader(ProtoReader):
     """Base class for trajectory readers that extends :class:`ProtoReader` with a
     :meth:`__del__` method.
@@ -1611,6 +1611,7 @@ class Reader(ProtoReader):
        Provides kwargs to be passed to :class:`Timestep`
 
     """
+
     def __init__(self, filename, convert_units=None, **kwargs):
         super(Reader, self).__init__()
 
@@ -1664,6 +1665,7 @@ class Writer(six.with_metaclass(_Writermeta, IObase)):
     See Trajectory API definition in :mod:`MDAnalysis.coordinates.__init__` for
     the required attributes and methods.
     """
+
     def convert_dimensions_to_unitcell(self, ts, inplace=True):
         """Read dimensions from timestep *ts* and return appropriate unitcell.
 
@@ -1730,7 +1732,7 @@ def has_valid_coordinates(self, criteria, x):
         x = np.ravel(x)
         return np.all(criteria["min"] < x) and np.all(x <= criteria["max"])
 
-    # def write_next_timestep(self, ts=None)
+        # def write_next_timestep(self, ts=None)
 
 
 class SingleFrameReader(ProtoReader):

package/MDAnalysis/core/universe.py

@@ -52,28 +52,27 @@
 .. autofunction:: Merge
 
 """
-import six
-
+import copy
 import errno
-import numpy as np
 import logging
-import copy
+import sys
 import uuid
 
 import MDAnalysis
-import sys
+import numpy as np
+import six
 
-from .. import _ANCHOR_UNIVERSES
-from ..lib import util
-from ..lib.util import cached
-from ..lib.log import ProgressMeter, _set_verbose
-from ..exceptions import NoDataError
 from . import groups
+from ._get_readers import get_reader_for, get_parser_for
 from .groups import (GroupBase, Atom, Residue, Segment,
                      AtomGroup, ResidueGroup, SegmentGroup)
-from ._get_readers import get_reader_for, get_parser_for
 from .topology import Topology
 from .topologyattrs import AtomAttr, ResidueAttr, SegmentAttr
+from .. import _ANCHOR_UNIVERSES
+from ..exceptions import NoDataError
+from ..lib import util
+from ..lib.log import ProgressMeter, _set_verbose
+from ..lib.util import cached
 
 logger = logging.getLogger("MDAnalysis.core.universe")
 
@@ -273,10 +272,10 @@ def _generate_from_topology(self):
         self.atoms = AtomGroup(np.arange(self._topology.n_atoms), self)
 
         self.residues = ResidueGroup(
-                np.arange(self._topology.n_residues), self)
+            np.arange(self._topology.n_residues), self)
 
         self.segments = SegmentGroup(
-                np.arange(self._topology.n_segments), self)
+            np.arange(self._topology.n_segments), self)
 
         # Update Universe namespace with segids
         for seg in self.segments:
@@ -381,10 +380,10 @@ def load_new(self, filename, format=None, in_memory=False, **kwargs):
                              " have the same number of atoms!\n"
                              "Topology number of atoms {top_n_atoms}\n"
                              "Trajectory: {fname} Number of atoms {trj_n_atoms}".format(
-                                 form=self.trajectory.format,
-                                 top_n_atoms=len(self.atoms),
-                                 fname=filename,
-                                 trj_n_atoms=self.trajectory.n_atoms))
+                form=self.trajectory.format,
+                top_n_atoms=len(self.atoms),
+                fname=filename,
+                trj_n_atoms=self.trajectory.n_atoms))
 
         if in_memory:
             self.transfer_to_memory(kwargs.get("in_memory_frame_interval", 1))
@@ -602,10 +601,10 @@ def _process_attr(self, attr):
             raise ValueError('Length of {attr} does not'
                              ' match number of {obj}s.\n'
                              'Expect: {n:d} Have: {m:d}'.format(
-                                 attr=attr.attrname,
-                                 obj=attr.per_object,
-                                 n=n_dict[attr.per_object],
-                                 m=len(attr)))
+                attr=attr.attrname,
+                obj=attr.per_object,
+                n=n_dict[attr.per_object],
+                m=len(attr)))
 
         self._class_bases[GroupBase]._add_prop(attr)
 
@@ -935,13 +934,13 @@ def Merge(*args):
             tg = tg.atomgroup_intersection(ag, strict=True)
 
             # Map them so they refer to our new indices
-            new_idx = [tuple(map(lambda x:mapping[x], entry))
+            new_idx = [tuple(map(lambda x: mapping[x], entry))
                        for entry in tg.indices]
             bondidx.extend(new_idx)
             if hasattr(tg, '_bondtypes'):
                 types.extend(tg._bondtypes)
             else:
-                types.extend([None]*len(tg))
+                types.extend([None] * len(tg))
         if any(t is None for t in types):
             attrs.append(bonds_class(bondidx))
         else: