Merge branch 'master' into add_axes_to_rg_fft

d2o/distributor_factory.py

@@ -255,12 +255,11 @@ def _infer_key_type(key):
         return (None, None)
     found_boolean = False
     # Check which case we got:
-    if isinstance(key, tuple) or isinstance(key, slice) or np.isscalar(key):
+    if isinstance(key, slice) or np.isscalar(key):
+        found = 'slicetuple'
+    elif isinstance(key, tuple) or isinstance(key, list):
         # Check if there is something different in the array than
         # scalars and slices
-        if isinstance(key, slice) or np.isscalar(key):
-            key = [key]
-
         scalarQ = np.array(map(np.isscalar, key))
         sliceQ = np.array(map(lambda z: isinstance(z, slice), key))
         if np.all(scalarQ + sliceQ):
@@ -273,8 +272,6 @@ def _infer_key_type(key):
     elif isinstance(key, distributed_data_object):
         found = 'd2o'
         found_boolean = (key.dtype == np.bool_)
-    elif isinstance(key, list):
-        found = 'indexinglist'
     else:
         raise ValueError(about._errors.cstring("ERROR: Unknown keytype!"))
     return (found, found_boolean)

nifty_utilities.py

 # -*- coding: utf-8 -*-
 
 import numpy as np
+from keepers import about
+from itertools import product
+
+
+def get_slice_list(shape, axes):
+    """
+    Helper function which generates slice list(s) to traverse over all
+    combinations of axes, other than the selected axes.
+
+    Parameters
+    ----------
+    shape: tuple
+        Shape of the data array to traverse over.
+    axes: tuple
+        Axes which should not be iterated over.
+
+    Yields
+    -------
+    list
+        The next list of indices and/or slice objects for each dimension.
+
+    Raises
+    ------
+    ValueError
+        If shape is empty.
+    ValueError
+        If axes(axis) does not match shape.
+    """
+
+    if not shape:
+        raise ValueError(about._errors.cstring("ERROR: shape cannot be None."))
+
+    if not all(axis < len(shape) for axis in axes):
+        raise ValueError(
+            about._errors.cstring("ERROR: axes(axis) does not match shape.")
+        )
+
+    axes_select = [0 if x in axes else 1 for x, y in enumerate(shape)]
+
+    axes_iterables = [range(y) for x, y in enumerate(shape) if x not in axes]
+
+    for index in product(*axes_iterables):
+        it_iter = iter(index)
+        slice_list = [
+            next(it_iter)
+            if axis else slice(None, None) for axis in axes_select
+        ]
+        yield slice_list
 
 
 def hermitianize_gaussian(x):
@@ -108,8 +156,8 @@ def convert_nested_list_to_object_array(x):
     dimension_counter = 0
     current_extract = x
     for i in xrange(len(possible_shape)):
-        if isinstance(current_extract, list) == False and\
-            isinstance(current_extract, tuple) == False:
+        if not isinstance(current_extract, list) and \
+                not isinstance(current_extract, tuple):
             break
         current_extract = current_extract[0]
         dimension_counter += 1
@@ -153,8 +201,10 @@ def field_map(ishape, function, *args):
             result = np.empty_like(args[0])
             for i in xrange(np.prod(result.shape)):
                 ii = np.unravel_index(i, result.shape)
-                result[ii] = function(*map(
-                                        lambda z: get_clipped(z, ii), args)
-                                      )
+                result[ii] = function(
+                    *map(
+                        lambda z: get_clipped(z, ii), args
+                    )
+                )
                 # result[ii] = function(*map(lambda z: z[ii], args))
             return result