1"""Index construction for blocks and nested sub-blocks."""2from __future__ import annotations3import numpy as np456def equal_blocks(N: int, K: int):7 """Contiguous equal blocks; requires K | N. Returns array (K, m) of indices."""8 if N % K:9 raise ValueError(f"K={K} does not divide N={N}")10 m = N // K11 return np.arange(N).reshape(K, m)121314def random_blocks(N: int, K: int, rng: np.random.Generator):15 """Random equal blocks (uniform permutation), (K, m)."""16 return rng.permutation(N).reshape(K, N // K)171819def unequal_blocks(sizes):20 """Contiguous blocks with the given sizes -> list of index arrays."""21 sizes = [int(s) for s in sizes]22 out, start = [], 023 for s in sizes:24 out.append(np.arange(start, start + s)); start += s25 return out262728def nested_sub_blocks(block: np.ndarray, s: int):29 """Split one block (length m) into s equal contiguous sub-blocks: (s, m/s)."""30 m = block.shape[-1]31 if m % s:32 raise ValueError(f"split multiplier {s} does not divide block size {m}")33 return block.reshape(block.shape[:-1] + (s, m // s))343536def check_splits(m: int, splits):37 for s in splits:38 if m % s:39 raise ValueError(f"block size m={m} not divisible by split multiplier s={s}")40 return True414243def skewed_sizes(N: int, K: int, ratio: float):44 """Unequal block sizes: sizes grow geometrically so that max/min ~= ratio; sum = N (integer-rounded)."""45 g = np.geomspace(1.0, ratio, K)46 sizes = np.floor(g / g.sum() * N).astype(int)47 sizes[-1] += N - sizes.sum()48 return sizes49