def add_dummy_padding(x, depth, boundary): """ Pads an array which has 'none' as the boundary type. Used to simplify trimming arrays which use 'none'. >>> import dask.array as da >>> x = da.arange(6, chunks=3) >>> add_dummy_padding(x, {0: 1}, {0: 'none'}).compute() # doctest: +NORMALIZE_WHITESPACE array([..., 0, 1, 2, 3, 4, 5, ...]) """ for k, v in boundary.items(): d = depth.get(k, 0) if v == "none" and d > 0: empty_shape = list(x.shape) empty_shape[k] = d empty_chunks = list(x.chunks) empty_chunks[k] = (d, ) empty = empty_like( getattr(x, "_meta", x), shape=empty_shape, chunks=empty_chunks, dtype=x.dtype, ) out_chunks = list(x.chunks) ax_chunks = list(out_chunks[k]) ax_chunks[0] += d ax_chunks[-1] += d out_chunks[k] = tuple(ax_chunks) x = concatenate([empty, x, empty], axis=k) x = x.rechunk(out_chunks) return x
def nearest(x, axis, depth): """Each reflect each boundary value outwards This mimics what the skimage.filters.gaussian_filter(... mode="nearest") does. """ left = ((slice(None, None, None), ) * axis + (slice(0, 1), ) + (slice(None, None, None), ) * (x.ndim - axis - 1)) right = ((slice(None, None, None), ) * axis + (slice(-1, -2, -1), ) + (slice(None, None, None), ) * (x.ndim - axis - 1)) l = concatenate([x[left]] * depth, axis=axis) r = concatenate([x[right]] * depth, axis=axis) l, r = _remove_overlap_boundaries(l, r, axis, depth) return concatenate([l, x, r], axis=axis)
def repeat(a, repeats, axis=None): if axis is None: if a.ndim == 1: axis = 0 else: raise NotImplementedError("Must supply an integer axis value") if not isinstance(repeats, Integral): raise NotImplementedError("Only integer valued repeats supported") if -a.ndim <= axis < 0: axis += a.ndim elif not 0 <= axis <= a.ndim - 1: raise ValueError("axis(=%d) out of bounds" % axis) if repeats == 0: return a[tuple( slice(None) if d != axis else slice(0) for d in range(a.ndim))] elif repeats == 1: return a cchunks = cached_cumsum(a.chunks[axis], initial_zero=True) slices = [] for c_start, c_stop in sliding_window(2, cchunks): ls = np.linspace(c_start, c_stop, repeats).round(0) for ls_start, ls_stop in sliding_window(2, ls): if ls_start != ls_stop: slices.append(slice(ls_start, ls_stop)) all_slice = slice(None, None, None) slices = [ (all_slice, ) * axis + (s, ) + (all_slice, ) * (a.ndim - axis - 1) for s in slices ] slabs = [a[slc] for slc in slices] out = [] for slab in slabs: chunks = list(slab.chunks) assert len(chunks[axis]) == 1 chunks[axis] = (chunks[axis][0] * repeats, ) chunks = tuple(chunks) result = slab.map_blocks(np.repeat, repeats, axis=axis, chunks=chunks, dtype=slab.dtype) out.append(result) return concatenate(out, axis=axis)
def constant(x, axis, depth, value): """Add constant slice to either side of array""" chunks = list(x.chunks) chunks[axis] = (depth, ) c = full_like( x, value, shape=tuple(map(sum, chunks)), chunks=tuple(chunks), dtype=x.dtype, ) return concatenate([c, x, c], axis=axis)
def periodic(x, axis, depth): """Copy a slice of an array around to its other side Useful to create periodic boundary conditions for overlap """ left = ((slice(None, None, None), ) * axis + (slice(0, depth), ) + (slice(None, None, None), ) * (x.ndim - axis - 1)) right = ((slice(None, None, None), ) * axis + (slice(-depth, None), ) + (slice(None, None, None), ) * (x.ndim - axis - 1)) l = x[left] r = x[right] l, r = _remove_overlap_boundaries(l, r, axis, depth) return concatenate([r, x, l], axis=axis)
def reflect(x, axis, depth): """Reflect boundaries of array on the same side This is the converse of ``periodic`` """ if depth == 1: left = ((slice(None, None, None), ) * axis + (slice(0, 1), ) + (slice(None, None, None), ) * (x.ndim - axis - 1)) else: left = ((slice(None, None, None), ) * axis + (slice(depth - 1, None, -1), ) + (slice(None, None, None), ) * (x.ndim - axis - 1)) right = ((slice(None, None, None), ) * axis + (slice(-1, -depth - 1, -1), ) + (slice(None, None, None), ) * (x.ndim - axis - 1)) l = x[left] r = x[right] l, r = _remove_overlap_boundaries(l, r, axis, depth) return concatenate([l, x, r], axis=axis)
def pad_edge(array, pad_width, mode, **kwargs): """ Helper function for padding edges. Handles the cases where the only the values on the edge are needed. """ kwargs = {k: expand_pad_value(array, v) for k, v in kwargs.items()} result = array for d in range(array.ndim): pad_shapes, pad_chunks = get_pad_shapes_chunks(result, pad_width, (d, )) pad_arrays = [result, result] if mode == "constant": from dask.array.utils import asarray_safe constant_values = kwargs["constant_values"][d] constant_values = [ asarray_safe(c, like=meta_from_array(array), dtype=result.dtype) for c in constant_values ] pad_arrays = [ broadcast_to(v, s, c) for v, s, c in zip(constant_values, pad_shapes, pad_chunks) ] elif mode in ["edge", "linear_ramp"]: pad_slices = [ result.ndim * [slice(None)], result.ndim * [slice(None)] ] pad_slices[0][d] = slice(None, 1, None) pad_slices[1][d] = slice(-1, None, None) pad_slices = [tuple(sl) for sl in pad_slices] pad_arrays = [result[sl] for sl in pad_slices] if mode == "edge": pad_arrays = [ broadcast_to(a, s, c) for a, s, c in zip(pad_arrays, pad_shapes, pad_chunks) ] elif mode == "linear_ramp": end_values = kwargs["end_values"][d] pad_arrays = [ a.map_blocks( linear_ramp_chunk, ev, pw, chunks=c, dtype=result.dtype, dim=d, step=(2 * i - 1), ) for i, (a, ev, pw, c) in enumerate( zip(pad_arrays, end_values, pad_width[d], pad_chunks)) ] elif mode == "empty": pad_arrays = [ empty_like(array, shape=s, dtype=array.dtype, chunks=c) for s, c in zip(pad_shapes, pad_chunks) ] result = concatenate([pad_arrays[0], result, pad_arrays[1]], axis=d) return result