Randomisation, Stratification and Balance

23.1 Three ways to flip coins (40 min)

All three are unbiased. Stratification removes the chance that, say, the treated group gets more heavy shoppers. That imbalance doesn't bias the estimate on average, but it adds noise to any single experiment. If strata are built from variables that predict the outcome (here: loyalty and pre-period revenue), the noise they remove is substantial.

23.2 Comparing designs over 400 replays (60 min)

Create day23/three_designs.py:

# day23/three_designs.py   (python -m day23.three_designs)
"""
Three ways to randomise the same 40,000 customers. All three are unbiased;
they differ in how much the estimate jumps around from one experiment to the next.
"""
import numpy as np
import polars as pl

from causalkit.lab import population, run, truth

pop = population()
n = pop.height
rng = np.random.default_rng(23)

# Strata: loyalty x quintile of pre-period revenue (10 strata)
quint = pl.Series(pop["pre_revenue"]).rank("ordinal").to_numpy() * 5 // (n + 1)
stratum = pop["loyalty"].to_numpy() * 5 + quint


def bernoulli() -> np.ndarray:            # a coin per customer: group sizes vary
    return (rng.random(n) < 0.5).astype(np.int8)


def complete() -> np.ndarray:             # exactly half treated
    return rng.permutation(np.repeat(np.array([0, 1], dtype=np.int8), n // 2))


def stratified() -> np.ndarray:           # exactly half treated WITHIN each stratum
    z = np.zeros(n, dtype=np.int8)
    for s in np.unique(stratum):
        idx = np.flatnonzero(stratum == s)
        z[rng.choice(idx, size=len(idx) // 2, replace=False)] = 1
    return z


results = {name: [] for name in ["bernoulli", "complete", "stratified"]}
sizes = {name: [] for name in results}
for r in range(400):
    for name, design in [("bernoulli", bernoulli), ("complete", complete), ("stratified", stratified)]:
        z = design()
        y = run(pop, z, "email_v2", seed=r)["revenue"].to_numpy()
        results[name].append(y[z == 1].mean() - y[z == 0].mean())
        sizes[name].append(z.sum())

true_effect = truth("email_v2")["revenue"]
rows = [{"design": k, "mean estimate": round(np.mean(v), 3), "SD of estimate": round(np.std(v), 3),
         "treated group size (min-max)": f"{min(sizes[k]):,}-{max(sizes[k]):,}"} for k, v in results.items()]
print(pl.DataFrame(rows))
print(f"True effect: {true_effect:+.3f}")

python -m day23.three_designs

│ design     ┆ mean estimate ┆ SD of estimate ┆ treated group size (min-max) │
│ bernoulli  ┆ 3.507         ┆ 1.018          ┆ 19,655-20,245                │
│ complete   ┆ 3.459         ┆ 0.979          ┆ 20,000-20,000                │
│ stratified ┆ 3.561         ┆ 0.874          ┆ 19,997-19,997                │
True effect: +3.480

• All three centre on the truth (differences of ±0.08 are simulation noise: 400 replays with SD ≈ 1).
• Stratification cuts the estimate's standard deviation by 11% compared with complete randomisation: free precision, worth about 25% more customers.
• Six of the 10 strata have an odd number of customers, and each gives its extra customer to control, so the treated group has 19,997 rather than 20,000. That's by design, not a bug.

23.3 Balance checks: what they can and can't tell you (45 min)

A balance table compares treated and control groups on pre-treatment characteristics, using the standardised mean difference (SMD, Day 13). In a randomised experiment it serves two purposes:

1. Catching bugs. If one group is clearly different (|SMD| > 0.1 with thousands of customers), assignment is probably broken: an eligibility filter applied to one group only, a cached old assignment, a bot filter that removed more of one group.
2. Reassuring readers. Stakeholders who see similar groups trust the result more.

What it can't do: prove the experiment worked. Balance on observed covariates says nothing directly about unobserved ones; in a randomised experiment, randomisation (not the table) is what protects those. And don't run significance tests on every covariate: with 20 covariates, one will be "significant" by chance (Day 26).

Try it. Compare simple hash assignment (Day 21) with stratified assignment (23.4) for five different salts: how big is the SMD of pre_revenue in each?

# day23/compare_balance.py   (python -m day23.compare_balance)
import numpy as np
import polars as pl

from causalkit.assignment import assign_stratified, balance_table
from causalkit.lab import population
from day21.hash_assignment import assign

pop = population()
rows = []
for salt in ["test_a", "test_b", "test_c", "test_d", "test_e"]:
    simple = balance_table(pop, assign(pop["customer_id"], salt))
    strat = balance_table(pop, assign_stratified(pop, salt))
    pick = lambda t: t.filter(pl.col("covariate") == "pre_revenue")["smd"].item()
    rows.append({"salt": salt, "smd_simple": pick(simple), "smd_stratified": pick(strat)})
print(pl.DataFrame(rows))

│ salt   ┆ smd_simple ┆ smd_stratified │
│ test_a ┆ 0.013      ┆ 0.002          │
│ test_b ┆ -0.0004    ┆ -0.0043        │
│ test_c ┆ 0.0121     ┆ -0.0031        │
│ test_d ┆ 0.0136     ┆ -0.0027        │
│ test_e ┆ -0.007     ┆ 0.0013         │

23.4 Building an assignment service (40 min)

Concept. Combine Day 21's hashing with stratification: within each stratum, order customers by their hash and treat every second one. The assignment stays deterministic (same salt, same groups), needs no stored table, and is exactly balanced within strata.

Create causalkit/assignment.py. It also holds the SRM check and the balance table, so every later experiment uses the same audit:

"""
NorthMart's assignment service (Day 23 mini-project): deterministic, stratified, auditable.

    from causalkit.assignment import assign_stratified, srm_pvalue, balance_table
"""
import hashlib

import numpy as np
import polars as pl
from scipy.stats import chisquare


def _bucket(customer_id: str, salt: str) -> int:
    return int(hashlib.sha256(f"{salt}:{customer_id}".encode()).hexdigest()[:12], 16)


def add_strata(pop: pl.DataFrame) -> pl.DataFrame:
    """Loyalty x quintile of pre-period revenue: 10 strata."""
    return pop.with_columns(
        stratum=pl.col("loyalty").cast(pl.Int32) * 5
        + (pl.col("pre_revenue").rank("ordinal") * 5 // (pl.len() + 1)).cast(pl.Int32)
    )


def assign_stratified(pop: pl.DataFrame, salt: str) -> np.ndarray:
    """Within each stratum, order customers by their hash and treat every second one.
    Same salt, same answer; exact 50/50 inside every stratum (+/- 1 customer)."""
    df = add_strata(pop).with_columns(
        h=pl.Series([_bucket(c, salt) for c in pop["customer_id"]], dtype=pl.UInt64),
        row=pl.int_range(pl.len()),
    )
    ranked = df.sort("stratum", "h").with_columns(
        z=(pl.int_range(pl.len()).over("stratum") % 2).cast(pl.Int8))
    return ranked.sort("row")["z"].to_numpy()


def srm_pvalue(z: np.ndarray, treat_share: float = 0.5) -> float:
    counts = np.bincount(z, minlength=2)
    return chisquare(counts, f_exp=[counts.sum() * (1 - treat_share), counts.sum() * treat_share]).pvalue


def balance_table(pop: pl.DataFrame, z: np.ndarray) -> pl.DataFrame:
    """Standardised mean differences on pre-treatment covariates; |SMD| > 0.1 is flagged."""
    df = pop.with_columns(z=pl.Series(z), mobile=(pl.col("device") == "mobile").cast(pl.Int8))
    rows = []
    for col in ["pre_orders", "pre_revenue", "loyalty", "mobile"]:
        a = df.filter(pl.col("z") == 1)[col].cast(pl.Float64)
        b = df.filter(pl.col("z") == 0)[col].cast(pl.Float64)
        smd = (a.mean() - b.mean()) / np.sqrt((a.var() + b.var()) / 2)
        rows.append({"covariate": col, "treated": round(a.mean(), 3), "control": round(b.mean(), 3),
                     "smd": round(smd, 4), "flag": bool(abs(smd) > 0.1)})
    return pl.DataFrame(rows)

23.5 Day 23 mini-project: assign and audit email_v2 (70 min)

Create day23/assignment_report.py:

# day23/assignment_report.py   (python -m day23.assignment_report)
"""Mini-project: assign the email_v2 test with the stratified service, and audit the result."""
import polars as pl

from causalkit.assignment import add_strata, assign_stratified, balance_table, srm_pvalue
from causalkit.lab import population

pop = population()
z = assign_stratified(pop, "email_v2")

print("Deterministic:", bool((assign_stratified(pop, "email_v2") == z).all()))
print(f"Treated: {z.sum():,} of {len(z):,}   SRM p-value: {srm_pvalue(z):.3f}")

per_stratum = (add_strata(pop).with_columns(z=pl.Series(z))
               .group_by("stratum").agg(customers=pl.len(), treated_share=pl.col("z").mean().round(4))
               .sort("stratum"))
print(per_stratum)
print(balance_table(pop, z))

python -m day23.assignment_report

Deterministic: True
Treated: 19,997 of 40,000   SRM p-value: 0.976
│ stratum ┆ customers ┆ treated_share │
│ 0       ┆ 6497      ┆ 0.4999        │
│ 1       ┆ 6126      ┆ 0.5           │
...
│ 9       ┆ 3536      ┆ 0.5           │
│ covariate   ┆ treated ┆ control ┆ smd     ┆ flag  │
│ pre_orders  ┆ 2.275   ┆ 2.262   ┆ 0.006   ┆ false │
│ pre_revenue ┆ 138.178 ┆ 138.536 ┆ -0.0021 ┆ false │
│ loyalty     ┆ 0.3     ┆ 0.3     ┆ -0.0001 ┆ false │
│ mobile      ┆ 0.621   ┆ 0.62    ┆ 0.0022  ┆ false │

Add an Assignment section to docs/EXPERIMENT_PLAN_email_v2.md: the salt (email_v2), the strata definition, and this audit. Anyone can now reproduce the exact groups from the plan alone.

Strata 0–4 are non-members by revenue quintile, 5–9 are members. Notice how unequal they are: loyalty members are 30% of customers but are concentrated in the top revenue quintiles.

23.6 Save your work with Git (15 min)

git add .
git commit -m "Day 23: stratified hash assignment service, balance and SRM audit"
git push

Coming next (Days 24–27). The email_v2 test is designed, sized and assigned. Days 24–27 analyse it: regression and robust standard errors, CUPED variance reduction, and the two most common ways teams fool themselves, testing many metrics and peeking early.

Day 23 checklist

• I can explain Bernoulli, complete and stratified randomisation, and when stratification helps
• I can show by simulation that stratification reduces the estimate's variance
• I can read a balance table and explain what it can and can't prove
• I have a deterministic, stratified assignment service with SRM and balance audits