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)
| 4266 | |
| 4267 | |
| 4268 | def common_blockdim(blockdims): |
| 4269 | """Find the common block dimensions from the list of block dimensions |
| 4270 | |
| 4271 | Currently only implements the simplest possible heuristic: the common |
| 4272 | block-dimension is the only one that does not span fully span a dimension. |
| 4273 | This is a conservative choice that allows us to avoid potentially very |
| 4274 | expensive rechunking. |
| 4275 | |
| 4276 | Assumes that each element of the input block dimensions has all the same |
| 4277 | sum (i.e., that they correspond to dimensions of the same size). |
| 4278 | |
| 4279 | Examples |
| 4280 | -------- |
| 4281 | >>> common_blockdim([(3,), (2, 1)]) |
| 4282 | (2, 1) |
| 4283 | >>> common_blockdim([(1, 2), (2, 1)]) |
| 4284 | (1, 1, 1) |
| 4285 | >>> common_blockdim([(2, 2), (3, 1)]) # doctest: +SKIP |
| 4286 | Traceback (most recent call last): |
| 4287 | ... |
| 4288 | ValueError: Chunks do not align |
| 4289 | """ |
| 4290 | if not any(blockdims): |
| 4291 | return () |
| 4292 | non_trivial_dims = {d for d in blockdims if len(d) > 1} |
| 4293 | if len(non_trivial_dims) == 1: |
| 4294 | return first(non_trivial_dims) |
| 4295 | if len(non_trivial_dims) == 0: |
| 4296 | return max(blockdims, key=first) |
| 4297 | |
| 4298 | if np.isnan(sum(map(sum, blockdims))): |
| 4299 | raise ValueError( |
| 4300 | f"Arrays' chunk sizes ({blockdims}) are unknown.\n\n" |
| 4301 | "A possible solution:\n" |
| 4302 | " x.compute_chunk_sizes()" |
| 4303 | ) |
| 4304 | |
| 4305 | if len(set(map(sum, non_trivial_dims))) > 1: |
| 4306 | raise ValueError("Chunks do not add up to same value", blockdims) |
| 4307 | |
| 4308 | # We have multiple non-trivial chunks on this axis |
| 4309 | # e.g. (5, 2) and (4, 3) |
| 4310 | |
| 4311 | # We create a single chunk tuple with the same total length |
| 4312 | # that evenly divides both, e.g. (4, 1, 2) |
| 4313 | |
| 4314 | # To accomplish this we walk down all chunk tuples together, finding the |
| 4315 | # smallest element, adding it to the output, and subtracting it from all |
| 4316 | # other elements and remove the element itself. We stop once we have |
| 4317 | # burned through all of the chunk tuples. |
| 4318 | # For efficiency's sake we reverse the lists so that we can pop off the end |
| 4319 | rchunks = [list(ntd)[::-1] for ntd in non_trivial_dims] |
| 4320 | total = sum(first(non_trivial_dims)) |
| 4321 | i = 0 |
| 4322 | |
| 4323 | out = [] |
| 4324 | while i < total: |
| 4325 | m = min(c[-1] for c in rchunks) |