Repository: dask/dask

Issue:
map_overlap should work with smaller-than-depth chunk sizes
Related to #6597, but distinct I think.

If I have an array chunked into individual slices along an axis, but I want to perform an operation across that axis, I would *expect* `map_overlap` to be able to pull as many chunks as it needs along that axis to make the boundary, ie this should work:

```python
from dask import array as da
from scipy import ndimage as ndi

rand = da.random.random((860, 1024, 1024), chunks=(1, 1024, 1024))
filtered = da.map_overlap(
    rand,
    lambda arr: ndi.median_filter(arr, 5),
    depth=(2, 2, 2),
    boundary='reflect',
)
```

This raises:

```pytb
ValueError: The overlapping depth 2 is larger than your
smallest chunk size 1. Rechunk your array
with a larger chunk size or a chunk size that
more evenly divides the shape of your array.
```

But one *should* be able to grab multiple chunks to pad sufficiently?

If auto-rechunking is the preferred solution, then feel free to close this, since I think #6597 covers that adequately.


Discussion:
I think that auto-rechunking isn't necessarily the best solution, but it's
probably sufficient and easy :)

On Tue, Sep 29, 2020 at 9:24 PM Juan Nunez-Iglesias <
notifications@github.com> wrote:

> Related to #6597 <https://github.com/dask/dask/issues/6597>, but distinct
> I think.
>
> If I have an array chunked into individual slices along an axis, but I
> want to perform an operation across that axis, I would *expect*
> map_overlap to be able to pull as many chunks as it needs along that axis
> to make the boundary, ie this should work:
>
> from dask import array as dafrom scipy import ndimage as ndi
> rand = da.random.random((860, 1024, 1024), chunks=(1, 1024, 1024))filtered = da.map_overlap(
>     rand,
>     lambda arr: ndi.median_filter(arr, 5),
>     depth=(2, 2, 2),
>     boundary='reflect',
> )
>
> This raises:
>
> ValueError: The overlapping depth 2 is larger than your
> smallest chunk size 1. Rechunk your array
> with a larger chunk size or a chunk size that
> more evenly divides the shape of your array.
>
> But one *should* be able to grab multiple chunks to pad sufficiently?
>
> If auto-rechunking is the preferred solution, then feel free to close
> this, since I think #6597 <https://github.com/dask/dask/issues/6597>
> covers that adequately.
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/dask/dask/issues/6688>, or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AACKZTH2BGFNKBY4RGQDHE3SIKXHTANCNFSM4R6S6QSQ>
> .
>


Resolving patch:
diff --git a/dask/array/overlap.py b/dask/array/overlap.py
--- a/dask/array/overlap.py
+++ b/dask/array/overlap.py
@@ -422,6 +422,57 @@ def boundaries(x, depth=None, kind=None):
     return x
 
 
+def ensure_minimum_chunksize(size, chunks):
+    """Determine new chunks to ensure that every chunk >= size
+
+    Parameters
+    ----------
+    size: int
+        The maximum size of any chunk.
+    chunks: tuple
+        Chunks along one axis, e.g. ``(3, 3, 2)``
+
+    Examples
+    --------
+    >>> ensure_minimum_chunksize(10, (20, 20, 1))
+    (20, 11, 10)
+    >>> ensure_minimum_chunksize(3, (1, 1, 3))
+    (5,)
+
+    See Also
+    --------
+    overlap
+    """
+    if size <= min(chunks):
+        return chunks
+
+    # add too-small chunks to chunks before them
+    output = []
+    new = 0
+    for c in chunks:
+        if c < size:
+            if new > size + (size - c):
+                output.append(new - (size - c))
+                new = size
+            else:
+                new += c
+        if new >= size:
+            output.append(new)
+            new = 0
+        if c >= size:
+            new += c
+    if new >= size:
+        output.append(new)
+    elif len(output) >= 1:
+        output[-1] += new
+    else:
+        raise ValueError(
+            f"The overlapping depth {size} is larger than your " f"array {sum(chunks)}."
+        )
+
+    return tuple(output)
+
+
 def overlap(x, depth, boundary):
     """Share boundaries between neighboring blocks
 
@@ -441,6 +492,9 @@ def overlap(x, depth, boundary):
     blocks ``{0: 2, 2: 5}`` means share two cells in 0 axis, 5 cells in 2 axis.
     Axes missing from this input will not be overlapped.
 
+    Any axis containing chunks smaller than depth will be rechunked if
+    possible.
+
     Examples
     --------
     >>> import numpy as np
@@ -477,18 +531,14 @@ def overlap(x, depth, boundary):
     depth2 = coerce_depth(x.ndim, depth)
     boundary2 = coerce_boundary(x.ndim, boundary)
 
-    # is depth larger than chunk size?
-    depth_values = [depth2.get(i, 0) for i in range(x.ndim)]
-    for d, c in zip(depth_values, x.chunks):
-        maxd = max(d) if isinstance(d, tuple) else d
-        if maxd > min(c):
-            raise ValueError(
-                "The overlapping depth %d is larger than your\n"
-                "smallest chunk size %d. Rechunk your array\n"
-                "with a larger chunk size or a chunk size that\n"
-                "more evenly divides the shape of your array." % (d, min(c))
-            )
-    x2 = boundaries(x, depth2, boundary2)
+    # rechunk if new chunks are needed to fit depth in every chunk
+    depths = [max(d) if isinstance(d, tuple) else d for d in depth2.values()]
+    new_chunks = tuple(
+        ensure_minimum_chunksize(size, c) for size, c in zip(depths, x.chunks)
+    )
+    x1 = x.rechunk(new_chunks)
+
+    x2 = boundaries(x1, depth2, boundary2)
     x3 = overlap_internal(x2, depth2)
     trim = dict(
         (k, v * 2 if boundary2.get(k, "none") != "none" else 0)
@@ -544,7 +594,9 @@ def map_overlap(
     """Map a function over blocks of arrays with some overlap
 
     We share neighboring zones between blocks of the array, map a
-    function, and then trim away the neighboring strips.
+    function, and then trim away the neighboring strips. If depth is
+    larger than any chunk along a particular axis, then the array is
+    rechunked.
 
     Parameters
     ----------


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
@@ -17,6 +17,7 @@
     reflect,
     periodic,
     overlap,
+    ensure_minimum_chunksize,
 )
 from dask.array.utils import assert_eq, same_keys
 
@@ -574,7 +575,6 @@ def test_depth_equals_boundary_length():
     assert_array_equal(result, expected)
 
 
-@pytest.mark.xfail
 def test_depth_greater_than_boundary_length():
     expected = np.arange(100).reshape(10, 10)
     darr = da.from_array(expected, chunks=(5, 5))
@@ -599,13 +599,31 @@ def test_depth_greater_than_boundary_length():
     assert_array_equal(result, expected)
 
 
-def test_bad_depth_raises():
-    expected = np.arange(144).reshape(12, 12)
-    darr = da.from_array(expected, chunks=(5, 5))
+@pytest.mark.parametrize(
+    "chunks",
+    [
+        ((5, 5, 2), (5, 5, 2)),
+        ((3, 3, 3, 3), (11, 1)),
+    ],
+)
+def test_depth_greater_than_smallest_chunk_combines_chunks(chunks):
+    a = np.arange(144).reshape(12, 12)
+    darr = da.from_array(a, chunks=chunks)
 
     depth = {0: 4, 1: 2}
+    output = overlap(darr, depth=depth, boundary=1)
+
+    assert all(c >= depth[0] * 2 for c in output.chunks[0])
+    assert all(c >= depth[1] * 2 for c in output.chunks[1])
+
 
-    pytest.raises(ValueError, overlap, darr, depth=depth, boundary=1)
+def test_depth_greater_than_dim():
+    a = np.arange(144).reshape(12, 12)
+    darr = da.from_array(a, chunks=(3, 5))
+
+    depth = {0: 13, 1: 4}
+    with pytest.raises(ValueError, match="The overlapping depth"):
+        overlap(darr, depth=depth, boundary=1)
 
 
 def test_none_boundaries():
@@ -702,3 +720,50 @@ def test_trim_boundry(boundary):
     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)
+
+
+def test_map_overlap_rechunks_array_if_needed():
+    # https://github.com/dask/dask/issues/6597
+    expected = np.arange(11)
+    x = da.from_array(expected, chunks=5)
+    y = x.map_overlap(lambda x: x, depth=2, boundary=0)
+    assert all(c >= 2 for c in y.chunks[0])
+    assert_eq(y, expected)
+
+
+def test_map_overlap_rechunks_array_along_multiple_dims_if_needed():
+    # https://github.com/dask/dask/issues/6688
+    rand = da.random.random((860, 1024, 1024), chunks=(1, 1024, 1024))
+    filtered = rand.map_overlap(
+        lambda arr: arr,
+        depth=(2, 2, 2),
+        boundary="reflect",
+    )
+    assert all(all(c >= 2 for c in chunks) for chunks in filtered.chunks)
+
+
+@pytest.mark.parametrize(
+    "chunks,expected",
+    [
+        [(10,), (10,)],
+        [(10, 10), (10, 10)],
+        [
+            (10, 10, 1),
+            (10, 11),
+        ],
+        [(20, 20, 20, 1), (20, 20, 11, 10)],
+        [(20, 20, 10, 1), (20, 20, 11)],
+        [(2, 20, 2, 20), (14, 10, 20)],
+        [(1, 1, 1, 1, 7), (11,)],
+        [(20, 20, 2, 20, 20, 2), (20, 12, 10, 20, 12, 10)],
+    ],
+)
+def test_ensure_minimum_chunksize(chunks, expected):
+    actual = ensure_minimum_chunksize(10, chunks)
+    assert actual == expected
+
+
+def test_ensure_minimum_chunksize_raises_error():
+    chunks = (5, 2, 1, 1)
+    with pytest.raises(ValueError, match="overlapping depth 10 is larger than"):
+        ensure_minimum_chunksize(10, chunks)