added Product class (from larray_editor) to make product for Sequences

similar to itertools.product but can be indexed and has a length

Author: gdementen

larray/util/misc.py

@@ -833,3 +833,70 @@ def __iter__(self):
 
     def __repr__(self):
         return 'SequenceZip({})'.format(self.sequences)
+
+
+# TODO: remove Product from larray_editor.utils (it is almost identical)
+class Product(object):
+    """
+    Represents the `cartesian product` of several sequences.
+
+    This is very similar to itertools.product but only accepts sequences and acts as a sequence (it can be
+    indexed and has a len).
+
+    Parameters
+    ----------
+    sequences : Iterable of Sequence
+        Sequences on which to apply the cartesian product.
+
+    Examples
+    --------
+    >>> p = Product([['a', 'b', 'c'], [1, 2]])
+    >>> for i in range(len(p)):
+    ...     print(p[i])
+    ('a', 1)
+    ('a', 2)
+    ('b', 1)
+    ('b', 2)
+    ('c', 1)
+    ('c', 2)
+    >>> p[1:4]
+    [('a', 2), ('b', 1), ('b', 2)]
+    >>> p[-3:]
+    [('b', 2), ('c', 1), ('c', 2)]
+    >>> list(p)
+    [('a', 1), ('a', 2), ('b', 1), ('b', 2), ('c', 1), ('c', 2)]
+    """
+    def __init__(self, sequences):
+        self.sequences = sequences
+        assert len(sequences)
+        shape = [len(a) for a in self.sequences]
+        self._div_mod = [(int(np.prod(shape[i + 1:])), shape[i])
+                         for i in range(len(shape))]
+        self._length = np.prod(shape)
+
+    def __len__(self):
+        return self._length
+
+    def __getitem__(self, key):
+        if isinstance(key, (int, np.integer)):
+            if key >= self._length:
+                raise IndexError("index %d out of range for Product of length %d" % (key, self._length))
+            # this is similar to np.unravel_index but a tad faster for scalars
+            return tuple(array[key // div % mod]
+                         for array, (div, mod) in zip(self.sequences, self._div_mod))
+        else:
+            assert isinstance(key, slice), "key (%s) has invalid type (%s)" % (key, type(key))
+            start, stop, step = key.indices(self._length)
+            div_mod = self._div_mod
+            arrays = self.sequences
+            # XXX: we probably want to return another Product object with an updated start/stop to stay
+            #      lazy in that case too.
+            return [tuple(array[idx // div % mod]
+                          for array, (div, mod) in zip(arrays, div_mod))
+                    for idx in range(start, stop, step)]
+
+    def __iter__(self):
+        return product(*self.sequences)
+
+    def __repr__(self):
+        return 'Product({})'.format(self.sequences)