slices_from_chunks_overlap

slices_from_chunks_overlap.py

# Proposed slices_from_chunks_overlap function 
# Mofified from slices_from_chunks from dask.array.core
from itertools import product
from dask.array.slicing import cached_cumsum


def slices_from_chunks_overlap(chunks, array_shape, depth=1):
    """Translate chunks tuple to a set of slices in product order
    Parameters
    ----------
    chunks : tuple
        The chunks of the corresponding dask array.
    array_shape : tuple
        Shape of the corresponding dask array.
    depth : int
        The number of pixels to overlap, providing we're not at the array edge.
    Example
    -------
    >>> slices_from_chunks_overlap(((4,), (7, 7)), (4, 14), depth=1)  # doctest: +NORMALIZE_WHITESPACE
     [(slice(0, 5, None), slice(0, 8, None)),
      (slice(0, 5, None), slice(6, 15, None))]
    """
    cumdims = [cached_cumsum(bds, initial_zero=True) for bds in chunks]

    slices = []
    for starts, shapes, maxshape in zip(cumdims, chunks, array_shape):
        inner_slices = []
        for s, dim in zip(starts, shapes):
            slice_start = max(0, s - depth)
            slice_stop = min(s + dim + depth, maxshape)
            inner_slices.append(slice(slice_start, slice_stop))
        slices.append(inner_slices)

    return list(product(*slices))

Fixed bug with Juan, 2021-11-10