Skills › Scientific Agent Skills › experimental-design › randomization.py
randomization.py
228 lines · skills/experimental-design/scripts/randomization.py at 5f2bc03
1"""Reproducible randomization / allocation schedules for experiments and trials.
3Randomization is what licenses causal inference: it breaks the link between
4treatment assignment and any confounder, measured or not. But "I shuffled it"
5is not enough — the *method* matters (simple vs. blocked vs. stratified) and the
6schedule must be reproducible (seeded) and auditable. This module produces
7allocation tables as pandas DataFrames, with a fixed seed so the exact schedule
8can be regenerated and archived.
10Functions:
11 simple_randomization independent coin-flip per unit (can yield imbalance)
12 block_randomization permuted blocks -> balance at complete-block boundaries
13 stratified_block_randomization blocks within strata -> balance per subgroup
14 cluster_randomization randomize whole clusters (sites/classes), not units
15 assign_factorial_runs randomize the RUN ORDER of a list of design rows
17Requires: numpy, pandas.
18"""
20from __future__ import annotations
22import operator
24import numpy as np
25import pandas as pd
28def _count(value, name, minimum=0):
29 """Validate counts before allocation; zero-sized blocks would never advance."""
30 try:
31 count = operator.index(value)
32 except TypeError:
33 raise ValueError(f"{name} must be an integer >= {minimum}") from None
34 if isinstance(value, (bool, np.bool_)) or count < minimum:
35 raise ValueError(f"{name} must be an integer >= {minimum}")
36 return count
39def _labels(values, name, unique=False, minimum=0):
40 values = list(values)
41 try:
42 valid = (len(values) >= minimum
43 and not pd.Series(values, dtype=object).isna().any()
44 and all(isinstance(hash(value), int) for value in values))
45 if unique:
46 valid = valid and len(set(values)) == len(values)
47 except (TypeError, ValueError):
48 valid = False
49 if not valid:
50 raise ValueError(f"{name} must contain non-missing, hashable labels"
51 + (" that are unique" if unique else "")
52 + f" (at least {minimum})")
53 return values
56def _normalize_ratio(arms, ratio):
57 """Turn arms + integer ratio into a block template list, e.g.
58 arms=['A','B'], ratio=(2,1) -> ['A','A','B']."""
59 arms = _labels(arms, "arms", unique=True, minimum=2)
60 if ratio is None:
61 ratio = [1] * len(arms)
62 if len(ratio) != len(arms):
63 raise ValueError("ratio must have one entry per arm")
64 try:
65 invalid = any(isinstance(r, (bool, np.bool_)) or r <= 0 or int(r) != r
66 for r in ratio)
67 except (TypeError, ValueError, OverflowError):
68 invalid = True
69 if invalid:
70 raise ValueError("ratio entries must be positive integers")
71 template = []
72 for arm, r in zip(arms, ratio):
73 template += [arm] * int(r)
74 return template
77def simple_randomization(n, arms=("treatment", "control"), ratio=None, seed=0):
78 """Independent random assignment per unit.
80 Simplest method; with small n it can produce noticeable arm-size imbalance
81 (like flipping few coins). Fine for large n. Use block_randomization when you
82 need tighter control of arm counts or have sequential enrollment.
83 """
84 n = _count(n, "n")
85 rng = np.random.default_rng(seed)
86 arms = list(arms)
87 template = _normalize_ratio(arms, ratio)
88 probs = np.array([template.count(a) for a in arms], dtype=float)
89 probs /= probs.sum()
90 # Preserve labels when arms mix numeric and string values.
91 labels = np.empty(len(arms), dtype=object)
92 labels[:] = arms
93 assign = rng.choice(labels, size=n, p=probs)
94 return pd.DataFrame({"unit_id": np.arange(1, n + 1), "arm": assign})
97def block_randomization(n, arms=("treatment", "control"), block_size=None,
98 ratio=None, seed=0):
99 """Permuted-block randomization: exact ratio at complete-block boundaries.
101 Within each block every arm appears in the specified ratio; the order inside
102 a block is shuffled. block_size must be a multiple of sum(ratio). Leaving it
103 None picks a small valid size. Partial blocks need not balance. This helper
104 uses a fixed block size; it does not implement variable block sizes or
105 allocation concealment. Keep the seed/schedule away from enrolling staff.
106 """
107 n = _count(n, "n")
108 rng = np.random.default_rng(seed)
109 arms = list(arms)
110 template = _normalize_ratio(arms, ratio)
111 unit = len(template)
112 if block_size is None:
113 block_size = unit * 2 # two of each ratio-unit per block
114 block_size = _count(block_size, "block_size", minimum=1)
115 if block_size % unit != 0:
116 raise ValueError(f"block_size ({block_size}) must be a multiple of "
117 f"sum(ratio)={unit}")
118 reps = block_size // unit
120 out = []
121 block_id = 0
122 while len(out) < n:
123 block = template * reps
124 rng.shuffle(block)
125 for a in block:
126 out.append((len(out) + 1, block_id, a))
127 block_id += 1
128 df = pd.DataFrame(out[:n], columns=["unit_id", "block", "arm"])
129 return df
132def stratified_block_randomization(strata, arms=("treatment", "control"),
133 block_size=None, ratio=None, seed=0):
134 """Block-randomize independently within each stratum.
136 Use when a prognostic variable (site, sex, disease stage) must be balanced
137 across arms. `strata` is a dict {stratum_label: n_in_that_stratum} or a
138 sequence of stratum labels (one per unit). Sequence inputs retain their input
139 order: unit_id is the original one-based position. Map real subject IDs to
140 these positions before use. Dict inputs produce grouped planning schedules.
141 Balance is exact only for complete blocks within each stratum.
142 """
143 arms = list(arms)
144 _normalize_ratio(arms, ratio)
145 seed = _count(seed, "seed")
146 if isinstance(strata, dict):
147 _labels(strata, "strata")
148 items = [(label, _count(count, "stratum count"), None)
149 for label, count in strata.items()]
150 else: # sequence of labels, with original unit identities retained
151 labels = _labels(strata, "strata")
152 positions = {}
153 for position, label in enumerate(labels, start=1):
154 positions.setdefault(label, []).append(position)
155 items = [(label, len(ids), ids) for label, ids in positions.items()]
157 frames = []
158 next_id = 1
159 for i, (label, count, ids) in enumerate(items):
160 df = block_randomization(count, arms=arms, block_size=block_size,
161 ratio=ratio, seed=seed + 1 + i)
162 df["unit_id"] = ids if ids is not None else np.arange(next_id, next_id + count)
163 next_id += count
164 df.insert(1, "stratum", pd.Series([label] * count, dtype=object))
165 frames.append(df)
166 if not frames:
167 return pd.DataFrame(columns=["unit_id", "stratum", "block", "arm"])
168 out = pd.concat(frames, ignore_index=True)
169 return out.sort_values("unit_id").reset_index(drop=True)
172def cluster_randomization(clusters, arms=("treatment", "control"), ratio=None,
173 block_size=None, seed=0):
174 """Randomize whole clusters (clinics, schools, litters) to arms.
176 The cluster — not the individual — is the unit of randomization AND the unit
177 of analysis-level independence. `clusters` is a list of cluster IDs (or an int
178 count). Returns one row per cluster. Analyze with a method that accounts for
179 clustering (mixed model / GEE); treating members as independent is
180 pseudoreplication. Uses blocking across clusters for arm balance.
181 """
182 if isinstance(clusters, (int, np.integer)):
183 clusters = [f"cluster_{i+1}" for i in range(_count(clusters, "clusters"))]
184 clusters = _labels(clusters, "clusters", unique=True)
185 df = block_randomization(len(clusters), arms=arms, ratio=ratio,
186 block_size=block_size, seed=seed)
187 df = df.drop(columns=["unit_id"])
188 df.insert(0, "cluster_id", clusters)
189 return df
192def assign_factorial_runs(design_df, seed=0):
193 """Randomize the execution order of a set of design runs (e.g. a DOE matrix).
195 Run order matters: executing a factorial design in a systematic order
196 confounds the factors with time/drift (the machine warms up, the reagent
197 degrades). Randomizing run order protects against that. Returns the design
198 with a 'run_order' column and rows sorted by it.
199 """
200 rng = np.random.default_rng(seed)
201 df = design_df.copy().reset_index(drop=True)
202 order = rng.permutation(len(df)) + 1
203 df["run_order"] = order
204 return df.sort_values("run_order").reset_index(drop=True)
207def arm_balance(df, arm_col="arm", by=None):
208 """Quick check: counts per arm (optionally within each stratum/block)."""
209 if by:
210 return df.groupby([by, arm_col]).size().unstack(fill_value=0)
211 return df[arm_col].value_counts()
214if __name__ == "__main__":
215 print("== simple (n=10) ==")
216 print(arm_balance(simple_randomization(10, seed=1)).to_dict())
218 print("\n== permuted blocks, 2:1 treatment:control, n=12 ==")
219 d = block_randomization(12, arms=["treatment", "control"], ratio=(2, 1), seed=1)
220 print(d.to_string(index=False))
221 print("balance:", arm_balance(d).to_dict())
223 print("\n== stratified by site (A=8, B=6) ==")
224 d = stratified_block_randomization({"siteA": 8, "siteB": 6}, seed=1)
225 print(arm_balance(d, by="stratum").to_string())
227 print("\n== cluster randomization (6 clinics) ==")
228 print(cluster_randomization(6, seed=1).to_string(index=False))