Repository: dask/dask

Issue:
performance of `map_overlap` with boundary='none' is hindered by `add_dummy_padding`
`map_overlap` provides a few boundary conditions. The performance of the `'none'` boundary condition is seriously affected when the array is trimmed  since `dummy_padding` is added to the array which forces a whole copy of the array.

To illustrate the problem, I define a convenience function that implements the trim logic at the end of `map_overlap`.


Benchmark:

```python
from dask import array as da
from dask.array.overlap import (coerce_boundary, 
                                coerce_depth, 
                                trim_internal, 
                                add_dummy_padding)
shape=(150, 1024, 1024)
chunks=(-1, 512, 512)
dtype='float32'
z = da.zeros(shape=shape, chunks=chunks, dtype=dtype).persist()

# convenience function, inverse of da.overlap.overlap
# logic taken from map_overlap
def trim(x, depth, boundary=None):
    boundary2 = coerce_boundary(x.ndim, boundary)
    depth2 = coerce_depth(x.ndim, depth)
    
    x2 = add_dummy_padding(x, depth2, boundary2)
    return trim_internal(x2, depth2)

# Call once to force python to load all the necessary modules
_ = trim(z, 10, boundary='nearest').persist()
%time _ = trim(z, 10, boundary='none').persist()
%time _ = trim(z, 10, boundary='nearest').persist()
%time _ = trim(z, 10, boundary='none').persist()
%time _ = trim(z, 10, boundary='nearest').persist()
%time _ = trim(z, 10, boundary='none').persist()
%time _ = trim(z, 10, boundary='nearest').persist()
```

Results: 
* 'nearest' finishes in 30 ms
* 'none' finishes in 500 ms

# Computation graph
## Trimming with 'nearest'

```python
trim(z, 10, boundary='nearest').visualize()
```

![image](https://user-images.githubusercontent.com/90008/45063296-44690780-b07b-11e8-9114-6081e2609099.png)

## Trimming with 'none'

```python
trim(z, 10, boundary='none').visualize()
```
![mydask](https://user-images.githubusercontent.com/90008/45063333-6b273e00-b07b-11e8-9b5a-50192e8952b2.png)


This is very likely the issue I was seeing in Issue #3938 

cc: @jni 



Discussion:
> boundary condition is seriously affected when the array is trimmed since dummy_padding is added to the array which forces a whole copy of the array.

It looks like it only affects the outer boundary chunks, but I suspect that in your case that this is the entire array.  

Do you have a recommendation on how to improve the implementation?
~~I agree that in my case it is "all boundary", but you may only want to overlap in a specific dimensions (a use case that was envisioned in `map_overlap`) in which case the `trim` functionality, that you would assume should be a no-op in those dimensions, is actually coming at a big cost.~~ this is probably not true.

I think the logic of `add_dummy_padding` can be moved in `chunk.trim` (or whatever function actually returns the sliced array), to allow the slicing to be aware of the boundary condition. `map_blocks` can pass `block_info` to the trimming function so that it becomes aware what block is being trimmed.

the slicing can then check if 
1. If it is located on boundary.
2. If that boundary was originally passed as 'none'.

Then it can choose to slice the axis as `:`, `overlap:`, `overlap:-overlap` or `:-overlap` accordingly.
That sounds reasonable to me.  Do you have time to help implement this?

On Tue, Sep 4, 2018 at 9:52 PM, Mark Harfouche <notifications@github.com>
wrote:

> I agree that in my case it is "all boundary", but you may only want to
> overlap in a specific dimensions (a use case that was envisioned in
> map_overlap) in which case the trim functionality, that you would assume
> should be a no-op in those dimensions, is actually coming at a big cost.
>
> I think the logic of add_dummy_padding can be moved in chunk.trim (or
> whatever function actually returns the sliced array), to allow the slicing
> to be aware of the boundary condition. map_blocks can pass block_info to
> the trimming function so that it becomes aware what block is being trimmed.
>
> the slicing can then check if
>
>    1. If it is located on boundary.
>    2. If that boundary was originally passed as 'none'.
>
> Then it can choose to slice the axis as :, overlap:, overlap:-overlap or
> :-overlap accordingly.
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/3947#issuecomment-418573385>, or mute
> the thread
> <https://github.com/notifications/unsubscribe-auth/AASszBphJjX2ZjaplIbRUymuPRfytjvcks5uXy5ogaJpZM4WZ9SP>
> .
>

PR #3950 implemented the "solution" where the logic is moved down to `chunk.trim` but I'm not convinced that it is the best solution. I mostly wrote it so I could benchmark my own work without incurring large memory copies and to learn about `block_info`.

My specific issue with PR #3950 is that `chunk.trim` now needs knowledge of options that are coming from higher level code. This makes the api of `chunk.trim` very cluttered.

I think a better solution would be to somehow:
1. Allow `chunk.trim` to handle asymmetric trimming
2. Call `chunk.trim` with the appropriate trimming parameters on each axis.

For now, `chunk.trim` supports different trimmings on a per axis level, but doesn't support asymmetric trimming.

I can probably help implement this, but I would need guidance on the correct choice for the API of `chunk.trim`:

Maybe:
`axes` could take:
1. An iterable of tuples. The inner tuple would be `trim_front`, `trim_back`. `trim_back` would have to be a positive number or `None`
2. If the iterable returns a non-iterable, then it assumes symmetric slicing.
3. If axis is a single value, symmetric trimming on all axes is assumed.
I think that it doesn't matter that much.  We can make some new function like chunk.trim that is only used in this context.  This inner function doesn't have to be public API.



Resolving patch:
diff --git a/dask/array/overlap.py b/dask/array/overlap.py
--- a/dask/array/overlap.py
+++ b/dask/array/overlap.py
@@ -153,7 +153,25 @@ def overlap_internal(x, axes):
     return Array(dsk, name, chunks, dtype=x.dtype)
 
 
-def trim_internal(x, axes):
+def trim_overlap(x, depth, boundary=None):
+    """Trim sides from each block.
+
+    This couples well with the ``map_overlap`` operation which may leave
+    excess data on each block.
+
+    See also
+    --------
+    dask.array.overlap.map_overlap
+
+    """
+
+    # parameter to be passed to trim_internal
+    axes = coerce_depth(x.ndim, depth)
+    boundary2 = coerce_boundary(x.ndim, boundary)
+    return trim_internal(x, axes=axes, boundary=boundary2)
+
+
+def trim_internal(x, axes, boundary=None):
     """ Trim sides from each block
 
     This couples well with the overlap operation, which may leave excess data on
@@ -164,17 +182,52 @@ def trim_internal(x, axes):
     dask.array.chunk.trim
     dask.array.map_blocks
     """
+    boundary = coerce_boundary(x.ndim, boundary)
+
     olist = []
     for i, bd in enumerate(x.chunks):
+        bdy = boundary.get(i, 'none')
         ilist = []
-        for d in bd:
-            ilist.append(d - axes.get(i, 0) * 2)
+        for j, d in enumerate(bd):
+            if bdy != 'none':
+                d = d - axes.get(i, 0) * 2
+            else:
+                d = d - axes.get(i, 0) if j != 0 else d
+                d = d - axes.get(i, 0) if j != len(bd) - 1 else d
+            ilist.append(d)
         olist.append(tuple(ilist))
 
     chunks = tuple(olist)
 
-    return map_blocks(partial(chunk.trim, axes=axes), x, chunks=chunks,
-                      dtype=x.dtype)
+    return map_blocks(partial(_trim, axes=axes, boundary=boundary),
+                      x, chunks=chunks, dtype=x.dtype)
+
+
+def _trim(x, axes, boundary, block_info):
+    """Similar to dask.array.chunk.trim but requires one to specificy the
+    boundary condition.
+
+    ``axes``, and ``boundary`` are assumed to have been coerced.
+
+    """
+    axes = [axes.get(i, 0) for i in range(x.ndim)]
+    axes_back = (-ax if ax else None for ax in axes)
+
+    trim_front = (
+        0 if (chunk_location == 0 and
+              boundary.get(i, 'none') == 'none') else ax
+        for i, (chunk_location, ax) in enumerate(
+            zip(block_info[0]['chunk-location'], axes)))
+    trim_back = (
+        None if (chunk_location == chunks - 1 and
+                 boundary.get(i, 'none') == 'none') else ax
+        for i, (chunks, chunk_location, ax) in enumerate(zip(
+            block_info[0]['num-chunks'],
+            block_info[0]['chunk-location'],
+            axes_back)))
+
+    return x[tuple(slice(front, back)
+             for front, back in zip(trim_front, trim_back))]
 
 
 def periodic(x, axis, depth):
@@ -462,8 +515,7 @@ def map_overlap(x, func, depth, boundary=None, trim=True, **kwargs):
     g = overlap(x, depth=depth2, boundary=boundary2)
     g2 = g.map_blocks(func, **kwargs)
     if trim:
-        g3 = add_dummy_padding(g2, depth2, boundary2)
-        return trim_internal(g3, depth2)
+        return trim_internal(g2, depth2, boundary2)
     else:
         return g2
 


Test patch:
diff --git a/dask/array/tests/test_overlap.py b/dask/array/tests/test_overlap.py
--- a/dask/array/tests/test_overlap.py
+++ b/dask/array/tests/test_overlap.py
@@ -415,3 +415,24 @@ def test_overlap_few_dimensions():
     assert len(a.dask) < len(c.dask)
 
     assert len(c.dask) < 10 * len(a.dask)
+
+
+@pytest.mark.parametrize('boundary',
+                         ['reflect', 'periodic', 'nearest', 'none'])
+def test_trim_boundry(boundary):
+    x = da.from_array(np.arange(24).reshape(4, 6), chunks=(2, 3))
+    x_overlaped = da.overlap.overlap(x, 2,
+                                     boundary={0: 'reflect', 1: boundary})
+    x_trimmed = da.overlap.trim_overlap(x_overlaped, 2,
+                                        boundary={0: 'reflect', 1: boundary})
+    assert np.all(x == x_trimmed)
+
+    x_overlaped = da.overlap.overlap(x, 2,
+                                     boundary={1: boundary})
+    x_trimmed = da.overlap.trim_overlap(x_overlaped, 2,
+                                        boundary={1: boundary})
+    assert np.all(x == x_trimmed)
+
+    x_overlaped = da.overlap.overlap(x, 2, boundary=boundary)
+    x_trimmed = da.overlap.trim_overlap(x_overlaped, 2, boundary=boundary)
+    assert np.all(x == x_trimmed)