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
def _broadcast_any(ar, shape, chunks):
if isinstance(ar, Array):
return broadcast_to(ar, shape).rechunk(chunks)
if isinstance(ar, np.ndarray):
return np.ascontiguousarray(np.broadcast_to(ar, shape))
def pad_stats(array, pad_width, mode, stat_length):
"""
Helper function for padding boundaries with statistics from the array.
In cases where the padding requires computations of statistics from part
or all of the array, this function helps compute those statistics as
requested and then adds those statistics onto the boundaries of the array.
"""
if mode == "median":
raise NotImplementedError("`pad` does not support `mode` of `median`.")
stat_length = expand_pad_value(array, stat_length)
result = np.empty(array.ndim * (3, ), dtype=object)
for idx in np.ndindex(result.shape):
axes = []
select = []
pad_shape = []
pad_chunks = []
for d, (i, s, c, w, l) in enumerate(
zip(idx, array.shape, array.chunks, pad_width, stat_length)):
if i < 1:
axes.append(d)
select.append(slice(None, l[0], None))
pad_shape.append(w[0])
pad_chunks.append(w[0])
elif i > 1:
axes.append(d)
select.append(slice(s - l[1], None, None))
pad_shape.append(w[1])
pad_chunks.append(w[1])
else:
select.append(slice(None))
pad_shape.append(s)
pad_chunks.append(c)
axes = tuple(axes)
select = tuple(select)
pad_shape = tuple(pad_shape)
pad_chunks = tuple(pad_chunks)
result_idx = array[select]
if axes:
if mode == "maximum":
result_idx = result_idx.max(axis=axes, keepdims=True)
elif mode == "mean":
result_idx = result_idx.mean(axis=axes, keepdims=True)
elif mode == "minimum":
result_idx = result_idx.min(axis=axes, keepdims=True)
result_idx = broadcast_to(result_idx, pad_shape, chunks=pad_chunks)
if mode == "mean":
if np.issubdtype(array.dtype, np.integer):
result_idx = rint(result_idx)
result_idx = result_idx.astype(array.dtype)
result[idx] = result_idx
result = block(result.tolist())
return result