Repository: dask/dask

Issue:
Reblocking along some subset of axes
For the user-facing reblock API in xray, I would like to be able to reblock along only some subset of axes, e.g., `ds.reblock(blockshape={'time': 100})` should only modify blockdims along the `time` dimension, even if there are also `latitude` and `longitude` dimensions.

This looks like a little tricky to do with the current reblock API.

Two options:
1. Expose `blockdims_from_blockshape` in the public API. That way, I can calculate blockdims myself, filling in ones I don't want to change.
2. Support providing only some axes to `blockdims` and `blockshape` in `reblock`. For example, to indicate that only axis 0 should be reblocked but others should be left alone (in a 3D array), we could write `blockshape=(100, None, None)` or `blockshape={0: 100}`.

Thoughts?

CC @PeterDSteinberg



Discussion:
This seems to be the trick for the near term, but I'll have to think about this more.  I agree it is not straighforward as written.

<pre><code>
import dask.array as da
import numpy as np
n = np.random.uniform(0,1,100).reshape((20,5))
x = da.from_array(n, blockdims=((10,),(2,3)))
# use the existing blockdims
changed_column = x.reblock(blockdims=(x.blockdims[0], (1,1,2,1)))
changed_row = x.reblock(blockdims=((5,2,3), x.blockdims[1]))
</code></pre>

The tricky case is when you want to assign a new _blockshape_ only along one dimension. Then you have to use `blockdims_from_blockshape` (or something similar) to turn that blockshape into blockdims.

I see better what you are saying.  I think we should make <pre><code>blockdims_from_blockshape</code></pre> able to be imported, and also we should make more options in the reblock method of dask arrays and the reblock function.  The options should support reblocking a subset of the dimensions, specified by either blockdims or blockshape.  I'll work on that soon.

> I'll work on that soon.

Great, I'm looking forward to it. By the way, `reblock` is awesome!


Resolving patch:
diff --git a/dask/array/reblock.py b/dask/array/reblock.py
--- a/dask/array/reblock.py
+++ b/dask/array/reblock.py
@@ -144,6 +144,33 @@ def intersect_blockdims(old_blockdims=None,
     return cross
 
 
+def blockdims_dict_to_tuple(old, new):
+    """
+
+    >>> blockdims_dict_to_tuple((4, 5, 6), {1: 10})
+    (4, 10, 6)
+    """
+    newlist = list(old)
+    for k, v in new.items():
+        newlist[k] = v
+    return tuple(newlist)
+
+
+def blockshape_dict_to_tuple(old_blockdims, d):
+    """
+
+    >>> blockshape_dict_to_tuple(((4, 4), (5, 5)), {1: 3})
+    ((4, 4), (3, 3, 3, 1))
+    """
+    shape = tuple(map(sum, old_blockdims))
+    new_blockdims = list(old_blockdims)
+    for k, v in d.items():
+        div = shape[k] // v
+        mod = shape[k] % v
+        new_blockdims[k] = (v,) * div + ((mod,) if mod else ())
+    return tuple(new_blockdims)
+
+
 def reblock(x, blockdims=None, blockshape=None):
     """
     Convert blocks in dask array x for new blockdims.
@@ -160,7 +187,12 @@ def reblock(x, blockdims=None, blockshape=None):
     >>> y.blockdims
     ((2, 4, 1), (4, 2, 1), (4, 3), (7,))
 
-    Parameters:
+    blockdims/blockshape also accept dict arguments mapping axis to blockshape
+
+    >>> y = reblock(x, blockshape={1: 2})  # reblock axis 1 with blockshape 2
+
+    Parameters
+    ----------
 
     x:   dask array
     blockdims:  the new block dimensions to create
@@ -168,8 +200,11 @@ def reblock(x, blockdims=None, blockshape=None):
 
     Provide one of blockdims or blockshape.
     """
-
-    if not blockdims:
+    if isinstance(blockdims, dict):
+        blockdims = blockdims_dict_to_tuple(x.blockdims, blockdims)
+    elif isinstance(blockshape, dict):
+        blockdims = blockshape_dict_to_tuple(x.blockdims, blockshape)
+    elif not blockdims:
         blockdims = blockdims_from_blockshape(x.shape, blockshape)
 
     crossed = intersect_blockdims(x.blockdims, blockdims)


Test patch:
diff --git a/dask/array/tests/test_reblock.py b/dask/array/tests/test_reblock.py
--- a/dask/array/tests/test_reblock.py
+++ b/dask/array/tests/test_reblock.py
@@ -148,3 +148,13 @@ def test_reblock_blockshape():
 def test_dtype():
     x = da.ones(5, blockshape=(2,))
     assert x.reblock(blockshape=(1,))._dtype == x._dtype
+
+
+def test_reblock_with_dict():
+    x = da.ones((24, 24), blockshape=(4, 8))
+    y = x.reblock(blockshape={0: 12})
+    assert y.blockdims == ((12, 12), (8, 8, 8))
+
+    x = da.ones((24, 24), blockshape=(4, 8))
+    y = x.reblock(blockdims={0: (12, 12)})
+    assert y.blockdims == ((12, 12), (8, 8, 8))