Switch to Grid classes to control how iteration is performed. (#61)

The SimpleGrid class supports non-regular chunks and improves the alignment of
blocks to the true underlying chunks for arrays wrapped by delayed subsetting
operations. Delayed combines are also handled by CompositeGrid class; this
tries to dispatch iteration to each of the component array's grids, reducing
the chance of inefficient iteration across incompatible chunk layouts.

The chunk_shape() generic has been replaced with the chunk_grid() generic,
which allows each array/operation class to determine how to construct an
appropriate grid for its own iteration. We also add a RegularTicks class to
more easily construct regular boundaries for grid construction.

Author: LTLA

README.md

@@ -272,6 +272,7 @@ Any array-like object can be used as a "seed" in a `DelayedArray` provided it ha
 - `dtype` and `shape` properties, like those in NumPy arrays.
 - a method for the `extract_dense_array()` generic.
 - a method for the `is_masked()` generic.
+- a method for the `chunk_grid()` generic.
 
 If the object may contain sparse data, it should also implement:
 
@@ -280,7 +281,6 @@ If the object may contain sparse data, it should also implement:
 
 It may also be desirable to implement:
 
-- a method for the `chunk_shape()` generic.
 - a method for the `create_dask_array()` generic.
 - a method for the `wrap()` generic.
 

src/delayedarray/BinaryIsometricOp.py

@@ -7,7 +7,7 @@
 from .extract_dense_array import extract_dense_array
 from .extract_sparse_array import extract_sparse_array
 from .create_dask_array import create_dask_array
-from .chunk_shape import chunk_shape
+from .chunk_grid import chunk_grid
 from .is_sparse import is_sparse
 from .is_masked import is_masked
 
@@ -130,22 +130,28 @@ def create_dask_array_BinaryIsometricOp(x: BinaryIsometricOp):
     return _execute(ls, rs, x._op)
 
 
-@chunk_shape.register
-def chunk_shape_BinaryIsometricOp(x: BinaryIsometricOp):
-    """See :py:meth:`~delayedarray.chunk_shape.chunk_shape`."""
-    lchunk = chunk_shape(x._left)
-    rchunk = chunk_shape(x._right)
-
-    # Not bothering with taking the lowest common denominator, as that
-    # might be too aggressive and expanding to the entire matrix size.
-    # We instead use the maximum chunk size (which might also expand, e.g.,
-    # if you're combining column-major and row-major matrices; oh well).
-    # Just accept that we'll probably need to break chunks during iteration.
-    output = []
-    for i in range(len(lchunk)):
-        output.append(max(lchunk[i], rchunk[i]))
-
-    return (*output,) 
+@chunk_grid.register
+def chunk_grid_BinaryIsometricOp(x: BinaryIsometricOp):
+    """See :py:meth:`~delayedarray.chunk_grid.chunk_grid`."""
+    lchunk = chunk_grid(x._left)
+    rchunk = chunk_grid(x._right)
+
+    # Favor the chunking for the more expensive grid, to avoid being penalized
+    # heavily from suboptimal chunking for that array. 
+    #
+    # Technically, we could optimize for the case where multiple dimensions
+    # have the same boundaries, in which case we should favor full extraction
+    # of the other dimensions and just iterate over the common dimensions.
+    # This avoids any chunk discrepancies but seems like a pretty unlikely case
+    # - if two arrays of the same shape disagree on the chunk boundaries of one
+    # dimension, they'd probably disagree on the others as well.
+    # 
+    # The other solution is to figure out some high-dimensional caching scheme
+    # for the partially consumed chunks. Sounds like a royal pain.
+    if lchunk.cost > rchunk.cost:
+        return lchunk
+    else:
+        return rchunk
 
 
 @is_sparse.register

src/delayedarray/Cast.py

@@ -5,7 +5,7 @@
 from .extract_dense_array import extract_dense_array
 from .extract_sparse_array import extract_sparse_array
 from .create_dask_array import create_dask_array
-from .chunk_shape import chunk_shape
+from .chunk_grid import chunk_grid
 from .is_sparse import is_sparse
 from .is_masked import is_masked
 
@@ -84,10 +84,10 @@ def create_dask_array_Cast(x: Cast):
     return target.astype(x._dtype)
 
 
-@chunk_shape.register
-def chunk_shape_Cast(x: Cast):
-    """See :py:meth:`~delayedarray.chunk_shape.chunk_shape`."""
-    return chunk_shape(x._seed)
+@chunk_grid.register
+def chunk_grid_Cast(x: Cast):
+    """See :py:meth:`~delayedarray.chunk_grid.chunk_grid`."""
+    return chunk_grid(x._seed)
 
 
 @is_sparse.register

src/delayedarray/Combine.py

@@ -7,7 +7,8 @@
 from .SparseNdarray import _concatenate_SparseNdarrays
 from .extract_sparse_array import extract_sparse_array
 from .create_dask_array import create_dask_array
-from .chunk_shape import chunk_shape
+from .chunk_grid import chunk_grid
+from .Grid import CompositeGrid
 from .is_sparse import is_sparse
 from .is_masked import is_masked
 
@@ -149,23 +150,11 @@ def create_dask_array_Combine(x: Combine):
     return numpy.concatenate((*extracted,), axis=x._along)
 
 
-@chunk_shape.register
-def chunk_shape_Combine(x: Combine):
-    """See :py:meth:`~delayedarray.chunk_shape.chunk_shape`."""
-    chunks = [chunk_shape(s) for s in x._seeds]
-
-    # Not bothering with doing anything too fancy here.  We just use the
-    # maximum chunk size (which might also expand, e.g., if you're
-    # combining column-major and row-major matrices; oh well).  Just accept
-    # that we'll probably need to break chunks during iteration.
-    output = []
-    for i in range(len(x._shape)):
-        dim = []
-        for ch in chunks:
-            dim.append(ch[i])
-        output.append(max(*dim))
-
-    return (*output,) 
+@chunk_grid.register
+def chunk_grid_Combine(x: Combine):
+    """See :py:meth:`~delayedarray.chunk_grid.chunk_grid`."""
+    chunks = [chunk_grid(s) for s in x._seeds]
+    return CompositeGrid(chunks, x._along)
 
 
 @is_sparse.register

src/delayedarray/DelayedArray.py

@@ -18,7 +18,7 @@
 from .extract_sparse_array import extract_sparse_array
 from .apply_over_blocks import apply_over_blocks
 from .create_dask_array import create_dask_array
-from .chunk_shape import chunk_shape
+from .chunk_grid import chunk_grid
 from .is_sparse import is_sparse
 from .is_masked import is_masked
 
@@ -78,6 +78,7 @@ class DelayedArray:
     - A method for the
       :py:meth:`~delayedarray.extract_dense_array.extract_dense_array` generic.
     - A method for the :py:meth:`~delayedarray.is_masked.is_masked` generic.
+    - A method for the :py:meth:`~delayedarray.chunk_grid.chunk_grid` generic.
 
     If the seed contains sparse data, it should also implement:
 
@@ -88,11 +89,11 @@ class DelayedArray:
 
     Optionally, a seed class may have:
 
-    - A method for the :py:meth:`~delayedarray.chunk_shape.chunk_shape` generic,
-      if there is some preferred dimension in which to take chunks of the array.
     - A method for the
       :py:meth:`~delayedarray.create_dask_array.create_dask_array` generic,
       if the seed is not already compatible with the **dask** package.
+    - a method for the `wrap()` generic, to create a ``DelayedArray``
+      subclass that is specific to this seed class.
     """
 
     def __init__(self, seed):
@@ -847,10 +848,10 @@ def create_dask_array_DelayedArray(x: DelayedArray):
     return create_dask_array(x._seed)
 
 
-@chunk_shape.register
-def chunk_shape_DelayedArray(x: DelayedArray):
-    """See :py:meth:`~delayedarray.chunk_shape.chunk_shape`."""
-    return chunk_shape(x._seed)
+@chunk_grid.register
+def chunk_grid_DelayedArray(x: DelayedArray):
+    """See :py:meth:`~delayedarray.chunk_grid.chunk_grid`."""
+    return chunk_grid(x._seed)
 
 
 @is_sparse.register