Repository: dask/dask

Issue:
Reduce I/O by fusing repeated getitems on arrays
Consider the following term

``` python
    # x[1000:2000][15:20]
    (getitem,
      (getitem, x, (slice(1000, 2000),)),
      slice(15, 20)
```

This is equivalent to the following much faster term

``` python
    # x[1015: 1020]
    (getitem, x, (slice(1015, 1020),))
```

This happens a lot when we use `from_array` and then slice

``` python
>>> d = da.from_array(x, blockshape=(1000,))  # First round of getitem
>>> d[1015:1020]  # Second round of getitem
```

If we can optimize then we cut down _a lot_ on I/O.  This is also pretty typical as people block in a bulk way but pull out thin slices.  This ends up being very important when slicing into large volumetric arrays (e.g. @shoyer 's work with [xray](xray.readthedocs.org)) and gives us a lot of leeway with choosing bad chunk sizes.

I would break this problem into two sub-problems
1.  A function to operate on terms, enacting the transformation shown in the above example, although being aware of multiple slices and possibly integer values as well.  This requires attention to detail but is not very complex from a dask perspective.  If we can get either of them to dive into these problems again @nevermindewe and @PeterDSteinberg have shown talent here.
2.  A function to apply such a transformation within nested terms (e.g. what happens if this subterm is within yet another function call).  This is more in line with previous work done by @eriknw or @jcrist.  In particular in relation to recent work on Rewrite Rules #82 how do we cleanly extend traversing terms and applying general callables?  What's the best way to enable these transformations?



Discussion:
This is actually useful enough that I wrote a whole special class to encapsulate this logic for xray, [`LazilyIndexedArray`](https://github.com/xray/xray/blob/0d164d848401209971ded33aea2880c1fdc892cb/xray/core/indexing.py#L226). There may be some relevant logic there for task 1 (e.g., how to slice a slice object).

Well, matching can easily find these terms, `RewriteRule` could be generalized to take in an arbitrary right hand side. If it's a function, it accepts a dict representing the match. If it's a term, a specific subs function is created and linked in, so that rewriting is always match -> validate (for non linear patterns) -> apply func. I have something like this already written up for SymPy, I'll port it over the weekend/early next week.

Here is a first pass on problem 1

``` python
from operator import getitem

def test_fuse_getitem():
    pairs = [((getitem, (getitem, 'x', slice(1000, 2000)), slice(15, 20)),
              (getitem, 'x', slice(1015, 1020))),

             ((getitem, (getitem, 'x', (slice(1000, 2000), slice(100, 200))),
                        (slice(15, 20), slice(50, 60))),
              (getitem, 'x', (slice(1015, 1020), slice(150, 160)))),

             ((getitem, (getitem, 'x', slice(1000, 2000)), 10),
              (getitem, 'x', 1010)),

             ((getitem, (getitem, 'x', (slice(1000, 2000), 10)),
                        (slice(15, 20),)),
              (getitem, 'x', (slice(1015, 1020), 10))),

             ((getitem, (getitem, 'x', (10, slice(1000, 2000))),
                        (slice(15, 20),)),
              (getitem, 'x', (10, slice(1015, 1020))))

        ]

    for inp, expected in pairs:
        result = fuse_getitem(inp)
        assert result == expected


def fuse_getitem(term):
    (_, (_, x, a), b) = term

    c = fuse_slice(a, b)

    return (getitem, x, c)


def normalize_slice(s):
    start, stop, step = s.start, s.stop, s.step
    if start is None:
        start = 0
    if step is None:
        step = 1
    return slice(start, stop, step)

def fuse_slice(a, b):
    """

    >>> fuse_slice(slice(1000, 2000), slice(10, 15))
    slice(1010, 1015, None)

    >>> fuse_slice(slice(1000, 2000), 10)
    1010

    >>> fuse_slice(slice(1000, 2000, 5), 10)
    1050

    >>> fuse_slice(slice(1000, 2000, 5), slice(10, 20, 2))
    slice(1050, 1100, 10)
    """

    if isinstance(a, slice):
        a = normalize_slice(a)
    if isinstance(b, slice):
        b = normalize_slice(b)

    if isinstance(a, slice) and isinstance(b, int):
        return a.start + b*(a.step if a.step is not None else 1)
    if isinstance(a, slice) and isinstance(b, slice):
        start = a.start + a.step * b.start
        if b.stop is not None:
            stop = a.start + a.step * b.stop
        else:
            stop = None
        if a.stop is not None:
            if stop is not None:
                stop = min(a.stop, stop)
            else:
                stop = a.stop
            stop = stop - ((stop - start) % b.step)
        step = a.step * b.step
        if step == 1:
            step = None
        return slice(start, stop, step)
    if isinstance(a, tuple) and not isinstance(b, tuple):
        b = (b,)
    if isinstance(a, tuple) and isinstance(b, tuple):
        j = 0
        result = list()
        for i in range(len(a)):
            if isinstance(a[i], int):
                result.append(a[i])
                continue
            result.append(fuse_slice(a[i], b[j]))
            j += 1
        return tuple(result)
    raise NotImplementedError()
```

Doesn't handle the stranger slicing conditions (e.g. negative steps) nor is it exhaustively tested.

Addressed problem 2 with #106. Just make a function that accepts a dict of matches and returns a term, and use it as the rhs of the rule.


Resolving patch:
diff --git a/dask/array/core.py b/dask/array/core.py
--- a/dask/array/core.py
+++ b/dask/array/core.py
@@ -14,10 +14,11 @@
         merge, curry, compose, reduce)
 import numpy as np
 from . import chunk
-from .slicing import slice_array, insert_many, remove_full_slices
+from .slicing import slice_array, insert_many
 from ..utils import deepmap, ignoring
 from ..async import inline_functions
 from ..optimize import cull, inline
+from .optimization import optimize
 from ..compatibility import unicode
 from .. import threaded, core
 from ..context import _globals
@@ -817,12 +818,8 @@ def get(dsk, keys, get=None, **kwargs):
     2. Use custom score function
     """
     get = get or _globals['get'] or threaded.get
-    fast_functions=kwargs.get('fast_functions',
-                             set([getitem, np.transpose]))
-    dsk2 = cull(dsk, list(core.flatten(keys)))
-    dsk3 = remove_full_slices(dsk2)
-    dsk4 = inline_functions(dsk3, fast_functions=fast_functions)
-    return get(dsk4, keys, **kwargs)
+    dsk2 = optimize(dsk, keys, **kwargs)
+    return get(dsk2, keys, **kwargs)
 
 
 def unpack_singleton(x):
diff --git a/dask/array/optimization.py b/dask/array/optimization.py
new file mode 100644
--- /dev/null
+++ b/dask/array/optimization.py
@@ -0,0 +1,195 @@
+from ..optimize import cull, fuse
+from ..core import flatten
+from ..async import inline_functions
+from ..optimize import dealias
+from operator import getitem
+from dask.rewrite import RuleSet, RewriteRule
+from toolz import valmap
+import numpy as np
+
+
+def optimize(dsk, keys, **kwargs):
+    """ Optimize dask for array computation
+
+    1.  Cull tasks not necessary to evaluate keys
+    2.  Remove full slicing, e.g. x[:]
+    3.  Inline fast functions like getitem and np.transpose
+    """
+    fast_functions=kwargs.get('fast_functions',
+                             set([getitem, np.transpose]))
+    dsk2 = cull(dsk, list(flatten(keys)))
+    dsk3 = remove_full_slices(dsk2)
+    dsk4 = fuse(dsk3)
+    dsk5 = valmap(rewrite_rules.rewrite, dsk4)
+    dsk6 = inline_functions(dsk5, fast_functions=fast_functions)
+    return dsk6
+
+
+def is_full_slice(task):
+    """
+
+    >>> is_full_slice((getitem, 'x',
+    ...                 (slice(None, None, None), slice(None, None, None))))
+    True
+    >>> is_full_slice((getitem, 'x',
+    ...                 (slice(5, 20, 1), slice(None, None, None))))
+    False
+    """
+    return (isinstance(task, tuple) and
+            task[0] == getitem and
+            (task[2] == slice(None, None, None) or
+             isinstance(task[2], tuple) and
+             all(ind == slice(None, None, None) for ind in task[2])))
+
+
+def remove_full_slices(dsk):
+    """ Remove full slices from dask
+
+    See Also:
+        dask.optimize.inline
+
+    Example
+    -------
+
+
+    >>> dsk = {'a': (range, 5),
+    ...        'b': (getitem, 'a', (slice(None, None, None),)),
+    ...        'c': (getitem, 'b', (slice(None, None, None),)),
+    ...        'd': (getitem, 'c', (slice(None, 5, None),)),
+    ...        'e': (getitem, 'd', (slice(None, None, None),))}
+
+    >>> remove_full_slices(dsk)  # doctest: +SKIP
+    {'a': (range, 5),
+     'e': (getitem, 'a', (slice(None, 5, None),))}
+    """
+    full_slice_keys = set(k for k, task in dsk.items() if is_full_slice(task))
+    dsk2 = dict((k, task[1] if k in full_slice_keys else task)
+                 for k, task in dsk.items())
+    dsk3 = dealias(dsk2)
+    return dsk3
+
+
+a, b, x = '~a', '~b', '~x'
+
+
+def fuse_slice_dict(match):
+    # Making new dimensions?  Need two getitems
+    # TODO: well, we could still optimize the non-None axes
+
+    try:
+        c = fuse_slice(match[a], match[b])
+        return (getitem, match[x], c)
+    except NotImplementedError:
+        return (getitem, (getitem, match[x], match[a]),
+                         match[b])
+
+
+rewrite_rules = RuleSet(RewriteRule((getitem, (getitem, x, a), b),
+                                    fuse_slice_dict,
+                                    (a, b, x)))
+
+
+def normalize_slice(s):
+    """ Replace Nones in slices with integers
+
+    >>> normalize_slice(slice(None, None, None))
+    slice(0, None, 1)
+    """
+    start, stop, step = s.start, s.stop, s.step
+    if start is None:
+        start = 0
+    if step is None:
+        step = 1
+    if start < 0 or step < 0 or stop is not None and stop < 0:
+        raise NotImplementedError()
+    return slice(start, stop, step)
+
+
+def fuse_slice(a, b):
+    """ Fuse stacked slices together
+
+    Fuse a pair of repeated slices into a single slice:
+
+    >>> fuse_slice(slice(1000, 2000), slice(10, 15))
+    slice(1010, 1015, None)
+
+    This also works for tuples of slices
+
+    >>> fuse_slice((slice(100, 200), slice(100, 200, 10)),
+    ...            (slice(10, 15), [5, 2]))
+    (slice(110, 115, None), [150, 120])
+
+    And a variety of other interesting cases
+
+    >>> fuse_slice(slice(1000, 2000), 10)  # integers
+    1010
+
+    >>> fuse_slice(slice(1000, 2000, 5), slice(10, 20, 2))
+    slice(1050, 1100, 10)
+
+    >>> fuse_slice(slice(1000, 2000, 5), [1, 2, 3])  # lists
+    [1005, 1010, 1015]
+
+    >>> fuse_slice(None, slice(None, None))  # doctest: +SKIP
+    None
+    """
+    # None only works if the second side is a full slice
+    if a is None and b == slice(None, None):
+        return None
+
+    # Replace None with 0 and one in start and step
+    if isinstance(a, slice):
+        a = normalize_slice(a)
+    if isinstance(b, slice):
+        b = normalize_slice(b)
+
+    if isinstance(a, slice) and isinstance(b, int):
+        if b < 0:
+            raise NotImplementedError()
+        return a.start + b*a.step
+
+    if isinstance(a, slice) and isinstance(b, slice):
+        start = a.start + a.step * b.start
+        if b.stop is not None:
+            stop = a.start + a.step * b.stop
+        else:
+            stop = None
+        if a.stop is not None:
+            if stop is not None:
+                stop = min(a.stop, stop)
+            else:
+                stop = a.stop
+            stop = stop
+        step = a.step * b.step
+        if step == 1:
+            step = None
+        return slice(start, stop, step)
+
+    if isinstance(b, list):
+        return [fuse_slice(a, bb) for bb in b]
+    if isinstance(a, list) and isinstance(b, (int, slice)):
+        return a[b]
+
+    if isinstance(a, tuple) and not isinstance(b, tuple):
+        b = (b,)
+
+    # If given two tuples walk through both, being mindful of uneven sizes
+    # and newaxes
+    if isinstance(a, tuple) and isinstance(b, tuple):
+        j = 0
+        result = list()
+        for i in range(len(a)):
+            #  axis ceased to exist  or we're out of b
+            if isinstance(a[i], int) or j == len(b):
+                result.append(a[i])
+                continue
+            while b[j] is None:  # insert any Nones on the rhs
+                result.append(None)
+                j += 1
+            result.append(fuse_slice(a[i], b[j]))  # Common case
+            j += 1
+        while j < len(b):  # anything leftover on the right?
+            result.append(b[j])
+            j += 1
+        return tuple(result)
+    raise NotImplementedError()
diff --git a/dask/array/slicing.py b/dask/array/slicing.py
--- a/dask/array/slicing.py
+++ b/dask/array/slicing.py
@@ -2,7 +2,6 @@
 from toolz import merge, first, accumulate
 from operator import getitem, add
 from ..compatibility import long
-from ..optimize import dealias
 import numpy as np
 
 
@@ -480,50 +479,6 @@ def new_blockdim(dim_shape, lengths, index):
     return [(slc.stop - slc.start) // slc.step for slc in slices]
 
 
-def is_full_slice(task):
-    """
-
-    >>> is_full_slice((getitem, 'x',
-    ...                 (slice(None, None, None), slice(None, None, None))))
-    True
-    >>> is_full_slice((getitem, 'x',
-    ...                 (slice(5, 20, 1), slice(None, None, None))))
-    False
-    """
-    return (isinstance(task, tuple) and
-            task[0] == getitem and
-            (task[2] == slice(None, None, None) or
-             isinstance(task[2], tuple) and
-             all(ind == slice(None, None, None) for ind in task[2])))
-
-
-def remove_full_slices(dsk):
-    """ Remove full slices from dask
-
-    See Also:
-        dask.optimize.inline
-
-    Example
-    -------
-
-
-    >>> dsk = {'a': (range, 5),
-    ...        'b': (getitem, 'a', (slice(None, None, None),)),
-    ...        'c': (getitem, 'b', (slice(None, None, None),)),
-    ...        'd': (getitem, 'c', (slice(None, 5, None),)),
-    ...        'e': (getitem, 'd', (slice(None, None, None),))}
-
-    >>> remove_full_slices(dsk)  # doctest: +SKIP
-    {'a': (range, 5),
-     'e': (getitem, 'a', (slice(None, 5, None),))}
-    """
-    full_slice_keys = set(k for k, task in dsk.items() if is_full_slice(task))
-    dsk2 = dict((k, task[1] if k in full_slice_keys else task)
-                 for k, task in dsk.items())
-    dsk3 = dealias(dsk2)
-    return dsk3
-
-
 def replace_ellipsis(n, index):
     """ Replace ... with slices, :, : ,:
 


Test patch:
diff --git a/dask/array/tests/test_array_core.py b/dask/array/tests/test_array_core.py
--- a/dask/array/tests/test_array_core.py
+++ b/dask/array/tests/test_array_core.py
@@ -722,3 +722,12 @@ def test_reductions():
     assert eq(da.nansum(a), np.nansum(x))
     assert eq(da.nanvar(a), np.nanvar(x))
     assert eq(da.nanstd(a), np.nanstd(x))
+
+
+def test_optimize():
+    x = np.arange(5).astype('f4')
+    a = da.from_array(x, blockshape=(2,))
+    expr = a[1:4] + 1
+    result = optimize(expr.dask, expr._keys())
+    assert isinstance(result, dict)
+    assert all(key in result for key in expr._keys())
diff --git a/dask/array/tests/test_optimization.py b/dask/array/tests/test_optimization.py
new file mode 100644
--- /dev/null
+++ b/dask/array/tests/test_optimization.py
@@ -0,0 +1,99 @@
+from dask.array.optimization import (getitem, rewrite_rules, optimize,
+        remove_full_slices, fuse_slice)
+from dask.utils import raises
+
+
+def test_fuse_getitem():
+    pairs = [((getitem, (getitem, 'x', slice(1000, 2000)), slice(15, 20)),
+              (getitem, 'x', slice(1015, 1020))),
+
+             ((getitem, (getitem, 'x', (slice(1000, 2000), slice(100, 200))),
+                        (slice(15, 20), slice(50, 60))),
+              (getitem, 'x', (slice(1015, 1020), slice(150, 160)))),
+
+             ((getitem, (getitem, 'x', slice(1000, 2000)), 10),
+              (getitem, 'x', 1010)),
+
+             ((getitem, (getitem, 'x', (slice(1000, 2000), 10)),
+                        (slice(15, 20),)),
+              (getitem, 'x', (slice(1015, 1020), 10))),
+
+             ((getitem, (getitem, 'x', (10, slice(1000, 2000))),
+                        (slice(15, 20),)),
+              (getitem, 'x', (10, slice(1015, 1020)))),
+
+             ((getitem, (getitem, 'x', (slice(1000, 2000), slice(100, 200))),
+                        (slice(None, None), slice(50, 60))),
+              (getitem, 'x', (slice(1000, 2000), slice(150, 160)))),
+
+             ((getitem, (getitem, 'x', (None, slice(None, None))),
+                        (slice(None, None), 5)),
+              (getitem, 'x', (None, 5))),
+
+             ((getitem, (getitem, 'x', (slice(1000, 2000), slice(10, 20))),
+                        (slice(5, 10),)),
+              (getitem, 'x', (slice(1005, 1010), slice(10, 20)))),
+
+             ((getitem, (getitem, 'x', (slice(1000, 2000),)),
+                        (slice(5, 10), slice(10, 20))),
+              (getitem, 'x', (slice(1005, 1010), slice(10, 20))))
+
+        ]
+
+    for inp, expected in pairs:
+        result = rewrite_rules.rewrite(inp)
+        assert result == expected
+
+
+def test_optimize_with_getitem_fusion():
+    dsk = {'a': 'some-array',
+           'b': (getitem, 'a', (slice(10, 20), slice(100, 200))),
+           'c': (getitem, 'b', (5, slice(50, 60)))}
+
+    result = optimize(dsk, ['c'])
+    expected = {'c': (getitem, 'some-array', (15, slice(150, 160)))}
+    assert result == expected
+
+
+def test_optimize_slicing():
+    dsk = {'a': (range, 10),
+           'b': (getitem, 'a', (slice(None, None, None),)),
+           'c': (getitem, 'b', (slice(None, None, None),)),
+           'd': (getitem, 'c', (slice(None, 5, None),)),
+           'e': (getitem, 'd', (slice(None, None, None),))}
+
+    expected = {'a': (range, 10),
+                'e': (getitem, 'a', (slice(None, 5, None),))}
+
+    assert remove_full_slices(dsk) == expected
+
+
+def test_fuse_slice():
+    assert fuse_slice(slice(10, 15), slice(0, 5, 2)) == slice(10, 15, 2)
+
+    assert fuse_slice((slice(100, 200),), (None, slice(10, 20))) == \
+            (None, slice(110, 120))
+    assert fuse_slice((slice(100, 200),), (slice(10, 20), None)) == \
+            (slice(110, 120), None)
+    assert fuse_slice((1,), (None,)) == \
+            (1, None)
+    assert fuse_slice((1, slice(10, 20)), (None, None, 3, None)) == \
+            (1, None, None, 13, None)
+
+    assert (raises(NotImplementedError,
+                  lambda: fuse_slice(slice(10, 15, 2), -1)) or
+            fuse_slice(slice(10, 15, 2), -1) == 14)
+
+
+def test_fuse_slice_with_lists():
+    assert fuse_slice(slice(10, 20, 2), [1, 2, 3]) == [12, 14, 16]
+    assert fuse_slice([10, 20, 30, 40, 50], [3, 1, 2]) == [40, 20, 30]
+    assert fuse_slice([10, 20, 30, 40, 50], 3) == 40
+    assert fuse_slice([10, 20, 30, 40, 50], -1) == 50
+    assert fuse_slice([10, 20, 30, 40, 50], slice(1, None, 2)) == [20, 40]
+
+
+def test_hard_fuse_slice_cases():
+    term = (getitem, (getitem, 'x', (None, slice(None, None))),
+                     (slice(None, None), 5))
+    assert rewrite_rules.rewrite(term) == (getitem, 'x', (None, 5))
diff --git a/dask/array/tests/test_slicing.py b/dask/array/tests/test_slicing.py
--- a/dask/array/tests/test_slicing.py
+++ b/dask/array/tests/test_slicing.py
@@ -1,6 +1,6 @@
 import dask
 import dask.array as da
-from dask.array.slicing import slice_array, _slice_1d, take, remove_full_slices
+from dask.array.slicing import slice_array, _slice_1d, take
 from operator import getitem
 import numpy as np
 from pprint import pprint
@@ -324,16 +324,3 @@ def test_slicing_and_blockdims():
     o = da.ones((24, 16), blockdims=((4, 8, 8, 4), (2, 6, 6, 2)))
     t = o[4:-4, 2:-2]
     assert t.blockdims == ((8, 8), (6, 6))
-
-
-def test_optimize_slicing():
-    dsk = {'a': (range, 10),
-           'b': (getitem, 'a', (slice(None, None, None),)),
-           'c': (getitem, 'b', (slice(None, None, None),)),
-           'd': (getitem, 'c', (slice(None, 5, None),)),
-           'e': (getitem, 'd', (slice(None, None, None),))}
-
-    expected = {'a': (range, 10),
-                'e': (getitem, 'a', (slice(None, 5, None),))}
-
-    assert remove_full_slices(dsk) == expected