import random
def plan(samples, cold_every, seed, family_size):
for name, value in (("samples", samples), ("cold_every", cold_every),
("family_size", family_size)):
if type(value) is not int or value < 1:
raise ValueError(f"randomization: invalid {name}")
if type(seed) is not int:
raise ValueError("randomization: invalid seed")
blocks = (samples + cold_every - 1) // cold_every
if blocks > 20:
raise ValueError("randomization: at most 20 seed blocks supported")
rng = random.Random(seed)
if blocks % 2 == 0:
signs = [1] * (blocks // 2) + [-1] * (blocks // 2)
rng.shuffle(signs)
kind = "balanced"
else:
signs = [rng.choice((-1, 1)) for _ in range(blocks)]
kind = "independent"
return {
"version": 1,
"unit": "cold_seed",
"kind": kind,
"seed": seed,
"samples": samples,
"cold_every": cold_every,
"family_size": family_size,
"blocks": [
{"samples": list(range(start, min(start + cold_every, samples))),
"order": ["base", "head"] if signs[i] == 1 else ["head", "base"]}
for i, start in enumerate(range(0, samples, cold_every))
],
}
def checked(design):
if not isinstance(design, dict):
raise ValueError("randomization: invalid design")
try:
expected = plan(design["samples"], design["cold_every"], design["seed"], design["family_size"])
except KeyError as error:
raise ValueError("randomization: missing design metadata") from error
if design != expected:
raise ValueError("randomization: schedule does not match its declared design")
return design
def ordered(sample, arms, design):
names = {arm[0] for arm in arms}
if design is None or not {"base", "head"} <= names:
return arms
order = design["blocks"][sample // design["cold_every"]]["order"]
by_name = {arm[0]: arm for arm in arms}
return [by_name[name] for name in order] + [
arm for arm in arms if arm[0] not in order
]
def verify_pairs(base, head, design, phase, allow_incomplete=False):
if design is None:
return None
checked(design)
expected = set(range(design["samples"]))
if phase == "cold":
expected = {block["samples"][0] for block in design["blocks"]}
if not set(base) <= expected or not set(head) <= expected:
raise ValueError("randomization: unexpected sample")
for sample in base:
order = design["blocks"][sample // design["cold_every"]]["order"]
for arm, records in (("base", base), ("head", head)):
if records[sample].get("position") != order.index(arm):
raise ValueError("randomization: observed order differs from schedule")
if set(base) != expected or set(head) != expected:
if allow_incomplete:
return None
raise ValueError("incomplete planned randomization blocks")
return None if allow_incomplete else design