Repository: dask/dask

Issue:
Issue with trim_excess in dask.array.coarsen
I am trying to use ``dask.array.coarsen`` on 3-d arrays and am finding that the result depends on the chunking of the data, which doesn't seem right. The following example illustrates the issue:

```python
In [1]: import numpy as np                                                      

In [2]: import dask.array as da                                                 

In [3]: array1 = da.random.random((623, 768, 768)).rechunk((322, 256, 256))     

In [4]: array2 = da.random.random((623, 768, 768)).rechunk((623, 768, 768))     

In [5]: import numpy as np                                                      

In [6]: da.coarsen(np.nanmean, array1, {1: 3}, trim_excess=True).shape          
Out[6]: (623, 255, 768)

In [7]: da.coarsen(np.nanmean, array2, {1: 3}, trim_excess=True).shape          
Out[7]: (623, 256, 768)
```

Since 768 is divisible by 3, I would expect the final array to have shape ``(623, 256, 768)`` as for ``array2``. The result for ``array1`` seems incorrect, so I think this is a bug?



Discussion:
I think the issue is that the coarsen is happening on each chunk so you are getting a different result than expected because 256 is not divisible by 3. That is why you get an error with array1 if you take out `trim_excess=True`. So I am not entirely sure if this is worth changing or not. 

The relevant code is here:

https://github.com/dask/dask/blob/366c7998398bc778c4aa5f4b6bb22c25b584fbc1/dask/array/routines.py#L1303:L1305

I feel like a top-level function such as this should probably return the same result independent of how the data is chunked (unlike a function like say ``map_blocks`` which explicitly operates on chunks)
`da.coarsen` doesn't apply any overlaps between neighboring blocks, so the output is always dependent on the chunking of the input. There may be a bug here in terms of the output shape; I'm not familiar enough with `coarsen` to know, perhaps @mrocklin remembers - this method is one of the first things added to dask in 2015 and hasn't been touched much since. However, the output *values* will still differ depending on chunking, since no windowing is used.

For windowed computations that don't depend on how the data is chunked I recommend looking at `da.map_overlap`.
> The result for array1 seems incorrect, so I think this is a bug?

Looks like it to me

> I feel like a top-level function such as this should probably return the same result independent of how the data is chunked 

+1

I think that the solution here is to verify that all chunks of the coarsened dimensions align with the coarsening, and if not, to call a minimally invasive rechunk automatically within this function.

This shouldn't add too much extra cost, and should make sure that we get correct and consistent results.
If someone wanted to solve this I would recommend the following approach:

Make a function that takes in the chunks for a single dimension and an integer coarsenting value and produces a new chunking that is aligned with that value

```python
def aligned_coarse_chunks(chunks: List[int], multiple: int) -> List[int]:
    """ Returns a new chunking aligned with the coarsening multiple"""

>>> aligned_coarse_chunks((20, 10, 15, 23, 24), 10)
(20, 10, 20, 20, 20, 2)  
```
Note that the sum of the two lists should be the same, all numbers should be multiples of the `multiple`, and any excess should be at the end.

Then I would apply that function to all of the dimensions that are being chunked with the right multiple.  If the resulting chunking is different from the inputs `.chunks` value then I would call rechunk with that combined tuple.
Here's my attempt at this function

``` python
from typing import List


def aligned_coarse_chunks(chunks: List[int], multiple: int) -> List[int]:
    """ Returns a new chunking aligned with the coarsening multiple"""

    def choose_new_size(multiple, q, left):
        """ 
        See if multiple * q is a good choice when 'left' elements are remaining.
        Else return multiple * (q-1)
        """
        possible = multiple * q
        if (left - possible) > 0:
            return possible
        else:
            return multiple * (q - 1)

    # print(chunks)
    # print(sum(chunks))
    newchunks = []
    left = sum(chunks) - sum(newchunks)
    chunkgen = (c for c in chunks)
    while left > 0:
        if left < multiple:
            newchunks.append(left)
            break

        chunk_size = next(chunkgen, 0)
        if chunk_size == 0:
            chunk_size = multiple

        q, r = divmod(chunk_size, multiple)
        # print(c0, left, q, r)
        if q == 0:
            continue
        elif r == 0:
            newchunks.append(chunk_size)
        elif r >= 5:
            newchunks.append(choose_new_size(multiple, q + 1, left))
        else:
            newchunks.append(choose_new_size(multiple, q, left))

        left = sum(chunks) - sum(newchunks)
        # print(newchunks, left)

    print(f"{chunks} → {newchunks}")

    # checks
    assert sum(chunks) == sum(newchunks)
    if sum(chunks) % multiple == 0:
        lastind = None
    else:
        lastind = -1
    assert all(c % multiple == 0 for c in newchunks[slice(lastind)])

    return tuple(newchunks)

```

For these tests:
``` python
aligned_coarse_chunks((20, 10, 15, 23, 24), 10)
aligned_coarse_chunks((20, 10, 15, 42, 23, 24), 10)
aligned_coarse_chunks((20, 10, 15, 47, 23, 24), 10)
aligned_coarse_chunks((2, 10, 15, 47, 23, 24), 10)
aligned_coarse_chunks((10, 20, 30, 40, 2), 10)
```

I get
```
(20, 10, 15, 23, 24) → [20, 10, 20, 20, 20, 2]
(20, 10, 15, 23, 24) → [20, 10, 20, 20, 20, 2]
(20, 10, 15, 42, 23, 24) → [20, 10, 20, 40, 20, 20, 4]
(20, 10, 15, 47, 23, 24) → [20, 10, 20, 50, 20, 10, 9]
(2, 10, 15, 47, 23, 24) → [10, 20, 50, 20, 20, 1]
(10, 20, 30, 40, 2) → [10, 20, 30, 40, 2]
```

Does this look sensible?

Resolving patch:
diff --git a/dask/array/routines.py b/dask/array/routines.py
--- a/dask/array/routines.py
+++ b/dask/array/routines.py
@@ -5,6 +5,7 @@
 from functools import wraps, partial
 from numbers import Real, Integral
 from distutils.version import LooseVersion
+from typing import List
 
 import numpy as np
 from tlz import concat, sliding_window, interleave
@@ -1403,17 +1404,76 @@ def piecewise(x, condlist, funclist, *args, **kw):
     )
 
 
+def aligned_coarsen_chunks(chunks: List[int], multiple: int) -> List[int]:
+    """
+    Returns a new chunking aligned with the coarsening multiple.
+    Any excess is at the end of the array.
+
+    Examples
+    --------
+    >>> aligned_coarsen_chunks(chunks=(1, 2, 3), multiple=4)
+    (4, 2)
+    >>> aligned_coarsen_chunks(chunks=(1, 20, 3, 4), multiple=4)
+    (20, 4, 4)
+    >>> aligned_coarsen_chunks(chunks=(20, 10, 15, 23, 24), multiple=10)
+    (20, 10, 20, 20, 20, 2)
+    """
+
+    def choose_new_size(multiple, q, left):
+        """
+        See if multiple * q is a good choice when 'left' elements are remaining.
+        Else return multiple * (q-1)
+        """
+        possible = multiple * q
+        if (left - possible) > 0:
+            return possible
+        else:
+            return multiple * (q - 1)
+
+    newchunks = []
+    left = sum(chunks) - sum(newchunks)
+    chunkgen = (c for c in chunks)
+    while left > 0:
+        if left < multiple:
+            newchunks.append(left)
+            break
+
+        chunk_size = next(chunkgen, 0)
+        if chunk_size == 0:
+            chunk_size = multiple
+
+        q, r = divmod(chunk_size, multiple)
+        if q == 0:
+            continue
+        elif r == 0:
+            newchunks.append(chunk_size)
+        elif r >= 5:
+            newchunks.append(choose_new_size(multiple, q + 1, left))
+        else:
+            newchunks.append(choose_new_size(multiple, q, left))
+
+        left = sum(chunks) - sum(newchunks)
+
+    return tuple(newchunks)
+
+
 @wraps(chunk.coarsen)
 def coarsen(reduction, x, axes, trim_excess=False, **kwargs):
-    if not trim_excess and not all(
-        bd % div == 0 for i, div in axes.items() for bd in x.chunks[i]
-    ):
+    if not trim_excess and not all(x.shape[i] % div == 0 for i, div in axes.items()):
         msg = "Coarsening factor does not align with block dimensions"
         raise ValueError(msg)
 
     if "dask" in inspect.getfile(reduction):
         reduction = getattr(np, reduction.__name__)
 
+    new_chunks = {}
+    for i, div in axes.items():
+        aligned = aligned_coarsen_chunks(x.chunks[i], div)
+        if aligned != x.chunks[i]:
+            new_chunks[i] = aligned
+    if new_chunks:
+        x = x.rechunk(new_chunks)
+
     name = "coarsen-" + tokenize(reduction, x, axes, trim_excess)
     dsk = {
         (name,)


Test patch:
diff --git a/dask/array/tests/test_routines.py b/dask/array/tests/test_routines.py
--- a/dask/array/tests/test_routines.py
+++ b/dask/array/tests/test_routines.py
@@ -1521,13 +1521,31 @@ def test_coarsen():
 
 def test_coarsen_with_excess():
     x = da.arange(10, chunks=5)
-    assert_eq(da.coarsen(np.min, x, {0: 3}, trim_excess=True), np.array([0, 5]))
+    assert_eq(da.coarsen(np.min, x, {0: 5}, trim_excess=True), np.array([0, 5]))
     assert_eq(
         da.coarsen(np.sum, x, {0: 3}, trim_excess=True),
-        np.array([0 + 1 + 2, 5 + 6 + 7]),
+        np.array([0 + 1 + 2, 3 + 4 + 5, 6 + 7 + 8]),
     )
 
 
+def test_coarsen_bad_chunks():
+
+    x1 = da.arange(10, chunks=5)
+    x2 = x1.rechunk((1, 2, 3, 4))
+    assert_eq(da.coarsen(np.sum, x1, {0: 5}), da.coarsen(np.sum, x2, {0: 5}))
+
+
+def test_aligned_coarsen_chunks():
+
+    from ..routines import aligned_coarsen_chunks as acc
+
+    assert acc((20, 10, 15, 23, 24), 10) == (20, 10, 20, 20, 20, 2)
+    assert acc((20, 10, 15, 42, 23, 24), 10) == (20, 10, 20, 40, 20, 20, 4)
+    assert acc((20, 10, 15, 47, 23, 24), 10) == (20, 10, 20, 50, 20, 10, 9)
+    assert acc((2, 10, 15, 47, 23, 24), 10) == (10, 20, 50, 20, 20, 1)
+    assert acc((10, 20, 30, 40, 2), 10) == (10, 20, 30, 40, 2)
+
+
 def test_insert():
     x = np.random.randint(10, size=(10, 10))
     a = da.from_array(x, chunks=(5, 5))