Find the common block dimensions from the list of block dimensions Currently only implements the simplest possible heuristic: the common block-dimension is the only one that does not span fully span a dimension. This is a conservative choice that allows us to avoid potentially very
(blockdims)
| 4061 | |
| 4062 | |
| 4063 | def common_blockdim(blockdims): |
| 4064 | """Find the common block dimensions from the list of block dimensions |
| 4065 | |
| 4066 | Currently only implements the simplest possible heuristic: the common |
| 4067 | block-dimension is the only one that does not span fully span a dimension. |
| 4068 | This is a conservative choice that allows us to avoid potentially very |
| 4069 | expensive rechunking. |
| 4070 | |
| 4071 | Assumes that each element of the input block dimensions has all the same |
| 4072 | sum (i.e., that they correspond to dimensions of the same size). |
| 4073 | |
| 4074 | Examples |
| 4075 | -------- |
| 4076 | >>> common_blockdim([(3,), (2, 1)]) |
| 4077 | (2, 1) |
| 4078 | >>> common_blockdim([(1, 2), (2, 1)]) |
| 4079 | (1, 1, 1) |
| 4080 | >>> common_blockdim([(2, 2), (3, 1)]) # doctest: +SKIP |
| 4081 | Traceback (most recent call last): |
| 4082 | ... |
| 4083 | ValueError: Chunks do not align |
| 4084 | """ |
| 4085 | if not any(blockdims): |
| 4086 | return () |
| 4087 | non_trivial_dims = {d for d in blockdims if len(d) > 1} |
| 4088 | if len(non_trivial_dims) == 1: |
| 4089 | return first(non_trivial_dims) |
| 4090 | if len(non_trivial_dims) == 0: |
| 4091 | return max(blockdims, key=first) |
| 4092 | |
| 4093 | if np.isnan(sum(map(sum, blockdims))): |
| 4094 | raise ValueError( |
| 4095 | "Arrays' chunk sizes (%s) are unknown.\n\n" |
| 4096 | "A possible solution:\n" |
| 4097 | " x.compute_chunk_sizes()" % blockdims |
| 4098 | ) |
| 4099 | |
| 4100 | if len(set(map(sum, non_trivial_dims))) > 1: |
| 4101 | raise ValueError("Chunks do not add up to same value", blockdims) |
| 4102 | |
| 4103 | # We have multiple non-trivial chunks on this axis |
| 4104 | # e.g. (5, 2) and (4, 3) |
| 4105 | |
| 4106 | # We create a single chunk tuple with the same total length |
| 4107 | # that evenly divides both, e.g. (4, 1, 2) |
| 4108 | |
| 4109 | # To accomplish this we walk down all chunk tuples together, finding the |
| 4110 | # smallest element, adding it to the output, and subtracting it from all |
| 4111 | # other elements and remove the element itself. We stop once we have |
| 4112 | # burned through all of the chunk tuples. |
| 4113 | # For efficiency's sake we reverse the lists so that we can pop off the end |
| 4114 | rchunks = [list(ntd)[::-1] for ntd in non_trivial_dims] |
| 4115 | total = sum(first(non_trivial_dims)) |
| 4116 | i = 0 |
| 4117 | |
| 4118 | out = [] |
| 4119 | while i < total: |
| 4120 | m = min(c[-1] for c in rchunks) |