Helper function for padding edges. Handles the cases where the only the values on the edge are needed.
(array, pad_width, mode, **kwargs)
| 1018 | |
| 1019 | |
| 1020 | def pad_edge(array, pad_width, mode, **kwargs): |
| 1021 | """ |
| 1022 | Helper function for padding edges. |
| 1023 | |
| 1024 | Handles the cases where the only the values on the edge are needed. |
| 1025 | """ |
| 1026 | |
| 1027 | kwargs = {k: expand_pad_value(array, v) for k, v in kwargs.items()} |
| 1028 | |
| 1029 | result = array |
| 1030 | for d in range(array.ndim): |
| 1031 | pad_shapes, pad_chunks = get_pad_shapes_chunks( |
| 1032 | result, pad_width, (d,), mode=mode |
| 1033 | ) |
| 1034 | pad_arrays = [result, result] |
| 1035 | |
| 1036 | if mode == "constant": |
| 1037 | from dask.array.utils import asarray_safe |
| 1038 | |
| 1039 | constant_values = kwargs["constant_values"][d] |
| 1040 | constant_values = [ |
| 1041 | asarray_safe(c, like=meta_from_array(array), dtype=result.dtype) |
| 1042 | for c in constant_values |
| 1043 | ] |
| 1044 | |
| 1045 | pad_arrays = [ |
| 1046 | broadcast_to(v, s, c) |
| 1047 | for v, s, c in zip(constant_values, pad_shapes, pad_chunks) |
| 1048 | ] |
| 1049 | elif mode in ["edge", "linear_ramp"]: |
| 1050 | pad_slices = [result.ndim * [slice(None)], result.ndim * [slice(None)]] |
| 1051 | pad_slices[0][d] = slice(None, 1, None) |
| 1052 | pad_slices[1][d] = slice(-1, None, None) |
| 1053 | pad_slices = [tuple(sl) for sl in pad_slices] |
| 1054 | |
| 1055 | pad_arrays = [result[sl] for sl in pad_slices] |
| 1056 | |
| 1057 | if mode == "edge": |
| 1058 | pad_arrays = [ |
| 1059 | broadcast_to(a, s, c) |
| 1060 | for a, s, c in zip(pad_arrays, pad_shapes, pad_chunks) |
| 1061 | ] |
| 1062 | elif mode == "linear_ramp": |
| 1063 | end_values = kwargs["end_values"][d] |
| 1064 | |
| 1065 | pad_arrays = [ |
| 1066 | a.map_blocks( |
| 1067 | linear_ramp_chunk, |
| 1068 | ev, |
| 1069 | pw, |
| 1070 | chunks=c, |
| 1071 | dtype=result.dtype, |
| 1072 | dim=d, |
| 1073 | step=(2 * i - 1), |
| 1074 | ) |
| 1075 | for i, (a, ev, pw, c) in enumerate( |
| 1076 | zip(pad_arrays, end_values, pad_width[d], pad_chunks) |
| 1077 | ) |
no test coverage detected
searching dependent graphs…