manus-skills/.kortix/opencode/skills/data/experimental-design/scripts/randomization.py

"""Reproducible randomization / allocation schedules for experiments and trials.

Randomization is what licenses causal inference: it breaks the link between
treatment assignment and any confounder, measured or not. But "I shuffled it"
is not enough — the *method* matters (simple vs. blocked vs. stratified) and the
schedule must be reproducible (seeded) and auditable. This module produces
allocation tables as pandas DataFrames, with a fixed seed so the exact schedule
can be regenerated and archived.

Functions:
  simple_randomization        independent coin-flip per unit (can yield imbalance)
  block_randomization         permuted blocks -> balance throughout enrollment
  stratified_block_randomization   blocks within strata -> balance per subgroup
  cluster_randomization       randomize whole clusters (sites/classes), not units
  assign_factorial_runs       randomize the RUN ORDER of a list of design rows

Requires: numpy, pandas.
"""

from __future__ import annotations

import numpy as np
import pandas as pd


def _normalize_ratio(arms, ratio):
    """Turn arms + integer ratio into a block template list, e.g.
    arms=['A','B'], ratio=(2,1) -> ['A','A','B']."""
    if ratio is None:
        ratio = [1] * len(arms)
    if len(ratio) != len(arms):
        raise ValueError("ratio must have one entry per arm")
    if any(r <= 0 for r in ratio):
        raise ValueError("ratio entries must be positive integers")
    template = []
    for arm, r in zip(arms, ratio):
        template += [arm] * int(r)
    return template


def simple_randomization(n, arms=("treatment", "control"), ratio=None, seed=0):
    """Independent random assignment per unit.

    Simplest method; with small n it can produce noticeable arm-size imbalance
    (like flipping few coins). Fine for large n. Use block_randomization when you
    need balance, especially for n < ~100 or sequential enrollment.
    """
    rng = np.random.default_rng(seed)
    arms = list(arms)
    template = _normalize_ratio(arms, ratio)
    probs = np.array([template.count(a) for a in arms], dtype=float)
    probs /= probs.sum()
    assign = rng.choice(arms, size=n, p=probs)
    return pd.DataFrame({"unit_id": np.arange(1, n + 1), "arm": assign})


def block_randomization(n, arms=("treatment", "control"), block_size=None,
                        ratio=None, seed=0):
    """Permuted-block randomization: balance is maintained throughout enrollment.

    Within each block every arm appears in the specified ratio; the order inside
    a block is shuffled. block_size must be a multiple of sum(ratio). Leaving it
    None picks a small valid size. Mild caveat: fixed small blocks are slightly
    predictable in unblinded trials — vary block size if that matters.
    """
    rng = np.random.default_rng(seed)
    arms = list(arms)
    template = _normalize_ratio(arms, ratio)
    unit = len(template)
    if block_size is None:
        block_size = unit * 2  # two of each ratio-unit per block
    if block_size % unit != 0:
        raise ValueError(f"block_size ({block_size}) must be a multiple of "
                         f"sum(ratio)={unit}")
    reps = block_size // unit

    out = []
    block_id = 0
    while len(out) < n:
        block = template * reps
        rng.shuffle(block)
        for a in block:
            out.append((len(out) + 1, block_id, a))
        block_id += 1
    df = pd.DataFrame(out[:n], columns=["unit_id", "block", "arm"])
    return df


def stratified_block_randomization(strata, arms=("treatment", "control"),
                                   block_size=None, ratio=None, seed=0):
    """Block-randomize independently within each stratum.

    Use when a prognostic variable (site, sex, disease stage) must be balanced
    across arms. `strata` is a dict {stratum_label: n_in_that_stratum} or a
    sequence of stratum labels (one per unit). Each stratum gets its own permuted
    blocks, guaranteeing balance within every subgroup.
    """
    if isinstance(strata, dict):
        items = list(strata.items())
    else:  # sequence of labels
        s = pd.Series(list(strata))
        items = list(s.value_counts().sort_index().items())

    frames = []
    for i, (label, count) in enumerate(items):
        df = block_randomization(count, arms=arms, block_size=block_size,
                                 ratio=ratio, seed=seed + 1 + i)
        df.insert(1, "stratum", label)
        frames.append(df)
    out = pd.concat(frames, ignore_index=True)
    out["unit_id"] = np.arange(1, len(out) + 1)
    return out


def cluster_randomization(clusters, arms=("treatment", "control"), ratio=None,
                          block_size=None, seed=0):
    """Randomize whole clusters (clinics, schools, litters) to arms.

    The cluster — not the individual — is the unit of randomization AND the unit
    of analysis-level independence. `clusters` is a list of cluster IDs (or an int
    count). Returns one row per cluster. Analyze with a method that accounts for
    clustering (mixed model / GEE); treating members as independent is
    pseudoreplication. Uses blocking across clusters for arm balance.
    """
    if isinstance(clusters, int):
        clusters = [f"cluster_{i+1}" for i in range(clusters)]
    clusters = list(clusters)
    df = block_randomization(len(clusters), arms=arms, ratio=ratio,
                             block_size=block_size, seed=seed)
    df = df.drop(columns=["unit_id"])
    df.insert(0, "cluster_id", clusters)
    return df


def assign_factorial_runs(design_df, seed=0):
    """Randomize the execution order of a set of design runs (e.g. a DOE matrix).

    Run order matters: executing a factorial design in a systematic order
    confounds the factors with time/drift (the machine warms up, the reagent
    degrades). Randomizing run order protects against that. Returns the design
    with a 'run_order' column and rows sorted by it.
    """
    rng = np.random.default_rng(seed)
    df = design_df.copy().reset_index(drop=True)
    order = rng.permutation(len(df)) + 1
    df["run_order"] = order
    return df.sort_values("run_order").reset_index(drop=True)


def arm_balance(df, arm_col="arm", by=None):
    """Quick check: counts per arm (optionally within each stratum/block)."""
    if by:
        return df.groupby([by, arm_col]).size().unstack(fill_value=0)
    return df[arm_col].value_counts()


if __name__ == "__main__":
    print("== simple (n=10) ==")
    print(arm_balance(simple_randomization(10, seed=1)).to_dict())

    print("\n== permuted blocks, 2:1 treatment:control, n=12 ==")
    d = block_randomization(12, arms=["treatment", "control"], ratio=(2, 1), seed=1)
    print(d.to_string(index=False))
    print("balance:", arm_balance(d).to_dict())

    print("\n== stratified by site (A=8, B=6) ==")
    d = stratified_block_randomization({"siteA": 8, "siteB": 6}, seed=1)
    print(arm_balance(d, by="stratum").to_string())

    print("\n== cluster randomization (6 clinics) ==")
    print(cluster_randomization(6, seed=1).to_string(index=False))