Blocking and clustering

Every design so far has treated all $n$ units the same way. Real experiments rarely do. Two departures cover most of practice:

Blocking. Split the units into $K$ strata by something you know in advance — clinic, school, pre-treatment risk decile — and run a separate completely randomized experiment inside each one. Balance on the blocking variable stops being a matter of luck.
Clustering. You cannot treat half a village, so you randomize villages and every unit inside a treated village is treated.

They look like different subjects. In Lean they are the same construction seen from two sides, and neither one requires re-proving unbiasedness. What each requires is a propensity score, and once you have that, HT_unbiased from chapter 8 does the rest. That is the modularity the library was built for: one inclusion-probability computation per design.

A design on one space induces a design on another

Start with the general tool, because clustering is nothing but a special case of it and blocking borrows its proof idea.

Suppose $D$ is a design on $m$ units and $f$ turns an assignment of those $m$ units into an assignment of $n$ units. Then “draw $w$ from $D$ , report $f (w)$ ” is a design on $n$ units: the pushforward $f_{*} D$ . Its probability at $z$ is the total probability of the fibre $f^{- 1} (z)$ .

The pushforward of a designNew tabOpens with 1,051 lines of library preamble — the code above is at the bottom of the editor.

/-- The **pushforward** of a design `D` on `m` units along a map `f` of assignments: draw an
assignment `w` from `D`, and report `f w`.

The probability of seeing `z` is the total probability of all the `w` that `f` sends to `z`,
i.e. of the *fibre* over `z`.  Summing the fibres recovers the whole sample space, which is
why the probabilities still add up to one (`Finset.sum_fiberwise`). -/
noncomputable def map (f : Assignment m → Assignment n) (D : Design m) : Design n where
  prob z := ∑ w ∈ univ.filter (fun w => f w = z), D.prob w
  nonneg _ := Finset.sum_nonneg fun w _ => D.nonneg w
  sum_one := (Finset.sum_fiberwise univ f D.prob).trans D.sum_one

A Design has three fields — the pmf, its nonnegativity, and that it sums to one — so defining one means supplying three things. Nonnegativity is a sum of nonnegative terms. Summing to one is the only real obligation, and it is a Mathlib lemma: Finset.sum_fiberwise univ f D.prob says that summing over the fibres of $f$ and then over the values of $f$ recovers the sum over the whole domain. That is the definition of a partition into fibres, and it will show up twice more in this chapter.

The one property we need of the pushforward is the change-of-variables formula.

Change of variablesNew tabOpens with 1,065 lines of library preamble — the code above is at the bottom of the editor.

/-- **Change of variables.**  Taking an expectation under the pushforward design is the same
as taking the expectation of the composed statistic under the original design:
`E_{D.map f}[g] = E_D[g ∘ f]`.

Everything we prove about `clusterDesign` below is an application of this one lemma: a
statistic of the *unit-level* assignment is a statistic of the *cluster-level* assignment. -/
theorem expect_map (f : Assignment m → Assignment n) (D : Design m) (g : Assignment n → ℝ) :
    (D.map f).expect g = D.expect (fun w => g (f w)) := by
  have key : ∀ z : Assignment n,
      (D.map f).prob z * g z = ∑ w ∈ univ.filter (fun w => f w = z), D.prob w * g (f w) := by
    intro z
    rw [map_prob, Finset.sum_mul]
    exact Finset.sum_congr rfl fun w hw => by rw [(Finset.mem_filter.mp hw).2]
  calc (D.map f).expect g
      = ∑ z : Assignment n, ∑ w ∈ univ.filter (fun w => f w = z), D.prob w * g (f w) :=
        Finset.sum_congr rfl fun z _ => key z
    _ = D.expect (fun w => g (f w)) := Finset.sum_fiberwise univ f _

Read the statement first. (D.map f).expect g is an expectation over assignments of $n$ units; D.expect (fun w => g (f w)) is an expectation over assignments of $m$ units. The claim is that a statistic of the downstream assignment is a statistic of the upstream one, with the same mean — exactly $E_{f_{*} D} [g] = E_{D} [g \circ f]$ .

The proof is two steps in a calc block, which is Lean’s way of writing a chain of equalities with a justification for each link. The have key step distributes one probability across its fibre: Finset.sum_mul moves the factor g z inside the sum, and inside the fibre f w = z holds by definition of the filter, so g z can be rewritten as g (f w). That turns the outer sum into a double sum over fibres, and Finset.sum_fiberwise collapses it.

Stratified complete randomization

Fix a blocking function b : Fin n → Fin K: unit $i$ lives in block b i. The block itself is the fibre.

Blocks are fibres of the blocking functionNew tabOpens with 1,121 lines of library preamble — the code above is at the bottom of the editor.

/-- The **block** (stratum) with label `k`: the set of units the blocking function `b` sends
to `k`.  The blocks partition the population by construction, since every unit has exactly
one value of `b i`. -/
def block (b : Fin n → Fin K) (k : Fin K) : Finset (Fin n) :=
  univ.filter (fun i => b i = k)

Writing the block as univ.filter (fun i => b i = k) rather than as a preimage or an image is deliberate: it is syntactically what Finset.sum_fiberwise produces, so “group this sum by block” is a rewrite rather than an induction. Write $N_{k} = # (block k)$ for the block sizes; they sum to $n$ .

The experimenter fixes the quotas n₁ : Fin K → ℕ in advance, and the design is uniform on the assignments meeting all $K$ quotas at once.

The admissible assignmentsNew tabOpens with 1,217 lines of library preamble — the code above is at the bottom of the editor.

/-- The **support of the stratified design**: the assignments that treat exactly `n₁ k` units
inside block `k`, for every `k` at once. -/
def stratSupport (b : Fin n → Fin K) (n₁ : Fin K → ℕ) : Finset (Assignment n) :=
  univ.filter (fun z => ∀ k, #(z ∩ block b k) = n₁ k)

Counting: an assignment is a tuple of within-block choices

Here is the one place where blocking costs real work. We need to know how many admissible assignments there are, and the answer — one $n_{1 k}$ -subset per block, chosen freely — is the combinatorial content of “the blocks are independent experiments”.

The counting principle behind blockingNew tabOpens with 1,182 lines of library preamble — the code above is at the bottom of the editor.

/-- **The counting principle behind blocking.**  Suppose that for each block `k` we are given
a family `t k` of admissible *within-block* treated sets (each contained in block `k`).  Then
an assignment is admissible in every block simultaneously if and only if it is obtained by
choosing one admissible set per block, so

`#{S | ∀ k, S ∩ block k ∈ t k} = ∏ k, #(t k)`.

Nothing about the `t k` is used except that they consist of subsets of their own block: the
choices in different blocks are *free of each other*, which is exactly what "the blocks are
independent experiments" means combinatorially.

The proof is the explicit bijection `S ↦ (fun k => S ∩ block k)` with inverse
`g ↦ ⋃ k, g k`. -/
theorem card_blockwise (b : Fin n → Fin K) (t : Fin K → Finset (Assignment n))
    (ht : ∀ k, ∀ T ∈ t k, T ⊆ block b k) :
    #(univ.filter (fun S : Assignment n => ∀ k, S ∩ block b k ∈ t k)) = ∏ k, #(t k) := by
  classical
  rw [← Fintype.card_piFinset t]
  refine Finset.card_nbij' (fun S k => S ∩ block b k) (fun g => univ.biUnion g)
    ?_ ?_ ?_ ?_
  · intro S hS
    simp only [Finset.coe_filter, Set.mem_ofPred_eq, Finset.mem_univ, true_and] at hS
    simpa only [Finset.mem_coe, Fintype.mem_piFinset] using hS
  · intro g hg
    simp only [Finset.mem_coe, Fintype.mem_piFinset] at hg
    have hsub : ∀ k, g k ⊆ block b k := fun k => ht k (g k) (hg k)
    simp only [Finset.coe_filter, Set.mem_ofPred_eq, Finset.mem_univ, true_and]
    exact fun k => by rw [inter_block_biUnion hsub k]; exact hg k
  · exact fun S _ => biUnion_inter_block b S
  · intro g hg
    simp only [Finset.mem_coe, Fintype.mem_piFinset] at hg
    exact funext fun k => inter_block_biUnion (fun k => ht k (g k) (hg k)) k

The statement is more general than we need, and that is what makes it easy: nothing is assumed about the admissible families t k except that each of their members sits inside its own block. The proof exhibits the bijection

S ⟼ (S \cap block k)_{k}, g ⟼ ⋃_{k} g (k),

and hands it to Finset.card_nbij', which asks for four things: that each map lands in the right set, and that the two composites are the identity. Those four bullets are the four ?_ holes in the refine. Fintype.card_piFinset supplies the count of the tuple side, $\prod_{k} # (t k)$ .

Specializing to t k = powersetCard (n₁ k) (block b k) gives the count.

How many blocked assignments there areNew tabOpens with 1,237 lines of library preamble — the code above is at the bottom of the editor.

/-- **How many blocked assignments are there?**  One `n₁ k`-subset per block, independently:
`∏ k, C(N_k, n₁ k)`, where `N_k = #(block k)`.

Contrast with complete randomisation, which has the single binomial coefficient
`C(n, n₁)`.  Blocking is exactly the act of throwing away the assignments that are unbalanced
across blocks, and this product counts what survives. -/
theorem card_stratSupport (b : Fin n → Fin K) (n₁ : Fin K → ℕ) :
    #(stratSupport b n₁) = ∏ k, (#(block b k)).choose (n₁ k) := by
  rw [stratSupport_eq_filter,
    card_blockwise b (fun k => (block b k).powersetCard (n₁ k))
      (fun k T hT => (Finset.mem_powersetCard.mp hT).1)]
  exact Finset.prod_congr rfl fun k _ => Finset.card_powersetCard _ _

So the design has $\prod_{k} (\binom{N_{k}}{n_{1 k}})$ assignments in its support, where complete randomization has $(\binom{n}{n_{1}})$ . Blocking is precisely the act of deleting the assignments that are unbalanced across blocks, and this product counts the survivors.

The design

Stratified complete randomizationNew tabOpens with 1,257 lines of library preamble — the code above is at the bottom of the editor.

/-- **Stratified (blocked) complete randomisation.**  Run an independent completely
randomised experiment inside each block: treat exactly `n₁ k` of the `N_k` units of block
`k`, uniformly at random, independently across blocks.

Equivalently -- and this is the definition we take, because it needs no product construction
-- put the *uniform* distribution on the set of assignments satisfying all `K` constraints at
once.  The two descriptions agree precisely because of `card_blockwise`: an admissible
assignment is the same thing as a tuple of within-block choices.

The hypothesis `h` says each block really can supply the units asked of it; without it the
support could be empty and there would be no uniform distribution to speak of. -/
noncomputable def stratifiedDesign (b : Fin n → Fin K) (n₁ : Fin K → ℕ)
    (h : ∀ k, n₁ k ≤ #(block b k)) : Design n where
  prob z := if (∀ k, #(z ∩ block b k) = n₁ k) then (#(stratSupport b n₁) : ℝ)⁻¹ else 0
  nonneg _ := by
    split
    · positivity
    · exact le_refl 0
  sum_one := by
    have hne : ((#(stratSupport b n₁) : ℕ) : ℝ) ≠ 0 :=
      Nat.cast_ne_zero.mpr (card_stratSupport_pos h).ne'
    calc ∑ z : Assignment n,
          (if (∀ k, #(z ∩ block b k) = n₁ k) then (#(stratSupport b n₁) : ℝ)⁻¹ else 0)
        = ∑ _z ∈ stratSupport b n₁, (#(stratSupport b n₁) : ℝ)⁻¹ := by
          rw [← Finset.sum_filter]; rfl
      _ = ((#(stratSupport b n₁) : ℕ) : ℝ) * (#(stratSupport b n₁) : ℝ)⁻¹ := by
          rw [Finset.sum_const, nsmul_eq_mul]
      _ = 1 := mul_inv_cancel₀ hne

Note what this is not. There is no product of designs here, and there could not be: a Design n is a pmf on Finset (Fin n), and $K$ separate experiments do not live in that type without a construction to glue them. So the definition takes the direct route — uniform on the admissible set — and the product structure enters only as a theorem, card_blockwise. The two descriptions agree because an admissible assignment is a tuple of within-block choices.

The hypothesis h : ∀ k, n₁ k ≤ #(block b k) is what makes the support nonempty; without it there would be no uniform distribution to normalize.

Balance is then exact, not approximate:

Exact balance in every blockNew tabOpens with 1,292 lines of library preamble — the code above is at the bottom of the editor.

/-- **Balance is exact, not approximate.**  On every assignment the design can produce, the
number of treated units inside block `k` is exactly `n₁ k` -- not on average, but always.
This is the whole point of blocking, and it is what lets the blocked estimator below replace
random denominators by known constants. -/
theorem stratified_card_inter_block {z : Assignment n}
    (hz : (stratifiedDesign b n₁ h).prob z ≠ 0) (k : Fin K) :
    #(z ∩ block b k) = n₁ k := by
  by_contra hc
  refine hz ?_
  rw [stratifiedDesign_prob]
  exact ite_eq_right (fun hall => hc (hall k))

by_contra assumes the negation of the goal and derives a contradiction; here the contradiction is that the design’s if condition fails, so the probability of z is 0, contradicting hz. Under complete randomization $# (z \cap block k)$ is a random variable with mean $n_{1} N_{k} / n$ ; here it is the constant $n_{1 k}$ on every assignment the design can produce. That is the entire point of blocking, and it is what will let the blocked estimator divide by known constants.

Inclusion probabilities and a non-constant propensity

Because the design is uniform on its support, every inclusion probability is a ratio of two counts, and the numerator factorizes over blocks by the same card_blockwise:

Inclusion counts factorize over blocksNew tabOpens with 1,317 lines of library preamble — the code above is at the bottom of the editor.

/-- **The blocks really are separate experiments.**  The number of blocked assignments that
treat all of a given set `S` factorises over the blocks: inside block `k` one must pick an
`n₁ k`-subset of that block containing `S`'s piece of it, and the `K` picks constrain each
other not at all.

Taking `S = ∅` recovers the total count `card_stratSupport`, `S = {i}` gives the propensity
score, and `S = {i, j}` the joint propensity -- a single computation covering all of the
design's inclusion probabilities, in the style of §3.4 of `SPEC.md`. -/
theorem card_stratSupport_filter_subset (b : Fin n → Fin K) (n₁ : Fin K → ℕ)
    (S : Assignment n) :
    #((stratSupport b n₁).filter (fun z => S ⊆ z))
      = ∏ k, #(((block b k).powersetCard (n₁ k)).filter (fun T => S ∩ block b k ⊆ T)) := by
  have hsub : ∀ k, ∀ T ∈ ((block b k).powersetCard (n₁ k)).filter (fun T => S ∩ block b k ⊆ T),
      T ⊆ block b k :=
    fun k T hT => (Finset.mem_powersetCard.mp (Finset.mem_filter.mp hT).1).1
  have hset : (stratSupport b n₁).filter (fun z => S ⊆ z)
      = univ.filter (fun z : Assignment n => ∀ k,
          z ∩ block b k ∈ ((block b k).powersetCard (n₁ k)).filter
            (fun T => S ∩ block b k ⊆ T)) := by
    ext z
    simp only [Finset.mem_filter, mem_stratSupport, Finset.mem_univ, true_and,
      Finset.mem_powersetCard, subset_iff_forall_inter_block b S z]
    constructor
    · rintro ⟨hc, hS⟩ k
      exact ⟨⟨Finset.inter_subset_right, hc k⟩, hS k⟩
    · exact fun hall => ⟨fun k => (hall k).1.2, fun k => (hall k).2⟩
  rw [hset, card_blockwise b _ hsub]

Take $S = {i}$ . Every block other than $i$ ‘s own contributes the same factor $(\binom{N_{k}}{n_{1 k}})$ to the numerator and the denominator, so it cancels, and what survives is the complete-randomization ratio inside a single block.

TheoremThe propensity score of a blocked design

A unit’s chance of treatment depends on its own block and nothing else:

π_{i} = \frac{n_{1, b (i)}}{N_{b (i)}} .

The only hypothesis is h, that each block can supply the units asked of it. And note what the statement does not say: it does not say $π$ is constant. Two units in the same block share a propensity; two units in different blocks generally do not.

New tabOpens with 1,411 lines of library preamble — the code above is at the bottom of the editor.

/-- **The propensity score of a blocked design is the within-block treated fraction:**
`π i = n₁ (b i) / N_(b i)`.

A unit's chance of treatment depends only on *its own* block -- the other blocks contribute
the same factor to the numerator and the denominator and cancel.  Inside a block the formula
is the complete-randomisation propensity `n₁ / N`, which is what "an independent completely
randomised experiment per block" should mean.

Note that the propensity is generally *not* constant across units: blocks with a higher
treated fraction have higher propensities.  That is why difference-in-means is biased under
blocking and has to be replaced by the block-weighted estimator below. -/
theorem stratified_propensity (i : Fin n) :
    (stratifiedDesign b n₁ h).propensity i
      = (n₁ (b i) : ℝ) / (#(block b (b i)) : ℝ) := by
  -- Step 1: the propensity is a ratio of two counts.
  have hcount : (stratifiedDesign b n₁ h).propensity i
      = (#((stratSupport b n₁).filter (fun z => i ∈ z)) : ℝ)
        * ((#(stratSupport b n₁) : ℕ) : ℝ)⁻¹ := by
    rw [Design.propensity_eq_inclusion, stratified_inclusion, ← stratSupport_filter_mem_eq]
  -- Step 2: both counts are products over blocks, differing only in `i`'s own factor.
  have hQ : ((∏ k ∈ univ.erase (b i), (#(block b k)).choose (n₁ k) : ℕ) : ℝ) ≠ 0 :=
    Nat.cast_ne_zero.mpr (Finset.prod_pos fun k _ => Nat.choose_pos (h k)).ne'
  have hcprod : (#(stratSupport b n₁) : ℕ)
      = (#(block b (b i))).choose (n₁ (b i))
        * ∏ k ∈ univ.erase (b i), (#(block b k)).choose (n₁ k) := by
    rw [card_stratSupport]
    exact (Finset.mul_prod_erase univ _ (Finset.mem_univ (b i))).symm
  rw [hcount, ← div_eq_mul_inv, card_stratSupport_filter_mem b n₁ i, hcprod,
    Nat.cast_mul, Nat.cast_mul, mul_div_mul_right _ _ hQ]
  -- Step 3: the surviving ratio is the complete-randomisation counting fact.
  rcases Nat.eq_zero_or_pos (n₁ (b i)) with hzero | hpos
  · rw [hzero, card_filter_mem_powersetCard_zero]
    simp only [Nat.cast_zero, zero_div]
  · rw [card_filter_mem_powersetCard (mem_block_self b i) hpos]
    exact_mod_cast choose_ratio_one hpos (h (b i))

Three things to notice in that proof. The rcases at the end splits on whether $n_{1 k}$ is zero, because choose_ratio_one needs $1 \leq n_{1 k}$ and the degenerate case has to be handled by hand — both sides are $0$ there, so the theorem needs no lower bound on the quota. exact_mod_cast closes a goal that matches up to ℕ → ℝ coercions, which is most of the friction in this library. And mul_div_mul_right is where the other blocks cancel: it needs the product over the other blocks to be nonzero, which is Nat.choose_pos applied to h.

And now the punchline, which is a sentence of statistics rather than of Lean:

Why plain difference in means goes wrong

In chapter 8 the argument for DiM_unbiased_completeRandomization was: on the support of the design the denominators $# z$ and $# z^{c}$ are the constants $n_{1}$ and $n_{0}$ , the weights $1 / π_{i}$ are the same for every unit, and therefore difference in means and Horvitz–Thompson are literally the same function there.

Under blocking the first half of that still holds — the total number treated is $N_{1} = \sum_{k} n_{1 k}$ on every admissible assignment — and the second half fails. So DiM is a fixed linear combination of the indicators and its expectation is available in closed form.

What difference in means actually estimates under blockingNew tabOpens with 1,728 lines of library preamble — the code above is at the bottom of the editor.

/-- **What the plain difference in means actually estimates under blocking.**

Under a blocked design the two denominators of `DiM` are *not* random: every admissible
assignment treats `∑ k, n₁ k` units in total, so the estimator is a fixed linear combination
of the indicators and its expectation is available in closed form,

`E[DiM] = (∑ᵢ πᵢ y1ᵢ) / N₁  −  (∑ᵢ (1 − πᵢ) y0ᵢ) / (n − N₁)`,  `N₁ = ∑ k, n₁ k`.

If the propensity `π` is the *constant* `N₁ / n`, the weights collapse and this is
`mean y1 − mean y0 = τ`.  If it is not — which is the normal case under blocking — the two
terms are weighted averages with the wrong weights, and the estimator is biased. -/
theorem expect_DiM_stratified {b : Fin n → Fin K} {n₁ : Fin K → ℕ}
    (h : ∀ k, n₁ k ≤ #(block b k)) (P : Population n) :
    (stratifiedDesign b n₁ h).expect (DiM P)
      = (∑ i, (stratifiedDesign b n₁ h).propensity i * P.y1 i) / ((∑ k, n₁ k : ℕ) : ℝ)
        - (∑ i, (1 - (stratifiedDesign b n₁ h).propensity i) * P.y0 i)
            / ((n : ℝ) - ((∑ k, n₁ k : ℕ) : ℝ)) := by
  have hle : ∑ k, n₁ k ≤ n := sum_n₁_le h
  have hsup : ∀ z, (stratifiedDesign b n₁ h).prob z ≠ 0 →
      DiM P z = (∑ i, Z z i * P.y1 i) / ((∑ k, n₁ k : ℕ) : ℝ)
        - (∑ i, (1 - Z z i) * P.y0 i) / ((n : ℝ) - ((∑ k, n₁ k : ℕ) : ℝ)) := by
    intro z hz
    have hcard : #z = ∑ k, n₁ k := stratified_card_eq_of_prob_ne_zero hz
    have hcompl : ((#zᶜ : ℕ) : ℝ) = (n : ℝ) - ((∑ k, n₁ k : ℕ) : ℝ) := by
      rw [card_compl, hcard, Nat.cast_sub hle]
    rw [DiM, sum_mem_eq_sum_Z_mul, sum_compl_eq_sum_one_sub_Z_mul, hcard, hcompl]
  rw [(stratifiedDesign b n₁ h).expect_congr_of_support hsup]
  simp only [Design.expect_sub, Design.expect_div_const, Design.expect_sum,
    Design.expect_mul_const, Design.expect_Z, Design.expect_const]

E [\hat{DiM}] = \frac{\sum_{i} π_{i} y_{1 i}}{N_{1}} - \frac{\sum_{i} (1 - π_{i}) y_{0 i}}{n - N_{1}} .

If $π_{i}$ is the constant $N_{1} / n$ the weights collapse and this is $\overset{―}{y_{1}} - \overset{―}{y_{0}} = τ$ . Otherwise it is a pair of weighted averages with the wrong weights.

The proof is the standard two-step. expect_congr_of_support lets us replace DiM P by a function that agrees with it wherever the probability is nonzero, which is where stratified_card_eq_of_prob_ne_zero (the fixed total) and card_compl (so $# z^{c} = n - N_{1}$ ) get used; then simp only with the linearity lemmas of Design.expect pushes the expectation through the sum and lands on Design.expect_Z, which is π by definition. There are no positivity hypotheses anywhere: at $N_{1} = 0$ both sides use Lean’s $x / 0 = 0$ and the identity still holds.

That formula makes the bias concrete. Five units, blocks ${0, 1}$ and ${2, 3, 4}$ , quotas $1$ and $2$ — treated fractions $1 / 2$ and $2 / 3$ :

A blocked design under which difference in means is biasedNew tabOpens with 1,779 lines of library preamble — the code above is at the bottom of the editor.

/-- Five units in two blocks: `{0, 1}` and `{2, 3, 4}`. -/
def bEx : Fin 5 → Fin 2 := ![0, 0, 1, 1, 1]

/-- One of the two units of the first block is treated, two of the three units of the second:
treated fractions `1/2` and `2/3`, so the propensity score is *not* constant. -/
def n₁Ex : Fin 2 → ℕ := ![1, 2]

theorem hEx : ∀ k, n₁Ex k ≤ #(block bEx k) := by decide

/-- A population whose potential outcomes are `y1 ≡ 0` everywhere, and `y0 = 1` on the first
block, `y0 = 0` on the second.  So the treatment effect is `-1` in block `0` and `0` in block
`1`, and `τ = (2 · (-1) + 3 · 0) / 5 = -2/5`. -/
def PEx : Population 5 := ⟨fun _ => 0, fun i => if bEx i = 0 then 1 else 0⟩

/-- **Plain difference in means is biased under blocking.**

`τ = -2/5`, but `E[DiM] = -1/2`.  Nothing pathological is going on: every block has both a
treated and a control unit, the design is a perfectly ordinary stratified randomisation, and
the estimator is the one everybody writes down first.  It is simply weighting units by the
sizes of the *realised* treated and control groups rather than by the propensities that
produced them. -/
theorem DiM_biased_stratified :
    (stratifiedDesign bEx n₁Ex hEx).expect (DiM PEx) ≠ PEx.tau := by
  have hprop : ∀ i : Fin 5, (stratifiedDesign bEx n₁Ex hEx).propensity i
      = (n₁Ex (bEx i) : ℝ) / (#(block bEx (bEx i)) : ℝ) := stratified_propensity
  have hN : (∑ k, n₁Ex k) = 3 := by decide
  have hc : ∀ i : Fin 5, #(block bEx (bEx i)) = if bEx i = 0 then 2 else 3 := by decide
  have hm : ∀ i : Fin 5, n₁Ex (bEx i) = if bEx i = 0 then 1 else 2 := by decide
  have hp : ∀ i : Fin 5, (stratifiedDesign bEx n₁Ex hEx).propensity i
      = ((if bEx i = 0 then 1 else 2 : ℕ) : ℝ) / ((if bEx i = 0 then 2 else 3 : ℕ) : ℝ) :=
    fun i => by rw [hprop i, hc i, hm i]
  have b0 : bEx 0 = 0 := by decide
  have b1 : bEx 1 = 0 := by decide
  have b2 : bEx 2 = 1 := by decide
  have b3 : bEx 3 = 1 := by decide
  have b4 : bEx 4 = 1 := by decide
  rw [expect_DiM_stratified hEx PEx, hN, Population.tau_eq_mean_sub_mean, mean, mean,
    Fin.sum_univ_five, Fin.sum_univ_five, Fin.sum_univ_five, Fin.sum_univ_five,
    hp 0, hp 1, hp 2, hp 3, hp 4]
  simp only [PEx, b0, b1, b2, b3, b4]
  norm_num

Here $τ = - 2 / 5$ and $E [\hat{DiM}] = - 1 / 2$ . Nothing degenerate is happening: every block contains both a treated and a control unit, the design is an ordinary stratified randomization, and the estimator is the one everyone writes down first. decide proves the small arithmetic facts about Fin and Finset by evaluating them in the kernel; norm_num finishes the rational arithmetic.

The cure is visible in the same formula. If the quotas are proportional — $n_{1 k} / N_{k} = N_{1} / n$ in every block — the propensity is the constant $N_{1} / n$ , the weights collapse, and difference in means is unbiased after all. Two of the twelve exercises at the end of the chapter prove exactly that, in that order. So the honest statement is not “blocking breaks difference in means”; it is “disproportionate allocation breaks difference in means, and blocking is the usual way to end up with it”.

The blocked estimator

The fix is the one every applied paper uses: estimate within each block, then average the block estimates with weights proportional to block size.

Within-block DiM, and the blocked estimatorNew tabOpens with 1,478 lines of library preamble — the code above is at the bottom of the editor.

/-- The **within-block difference in means** for block `k`: the treated-minus-control
comparison computed using only the units of that block. -/
noncomputable def blockDiM (b : Fin n → Fin K) (P : Population n) (z : Assignment n)
    (k : Fin K) : ℝ :=
  (∑ i ∈ z ∩ block b k, P.y1 i) / (#(z ∩ block b k) : ℝ)
    - (∑ i ∈ zᶜ ∩ block b k, P.y0 i) / (#(zᶜ ∩ block b k) : ℝ)

/-- The **blocked (post-stratified) estimator**: run difference-in-means separately inside
each block, then average the block estimates with weights proportional to block *size*.

`blockedDiM = ∑ k, (N_k / n) · DiM_k`.

The weights are the population shares of the blocks, not the sample shares -- they are fixed
numbers known before the experiment.  Weighting by block size is what makes the estimator
target the population ATE `τ` rather than some block-weighted variant of it. -/
noncomputable def blockedDiM (b : Fin n → Fin K) (P : Population n) (z : Assignment n) : ℝ :=
  ∑ k, ((#(block b k) : ℝ) / (n : ℝ)) * blockDiM b P z k

{\hat{τ}}_{blocked} = \sum_{k} \frac{N_{k}}{n} (\frac{\sum_{i \in z \cap B_{k}} y_{1 i}}{# (z \cap B_{k})} - \frac{\sum_{i \in z^{c} \cap B_{k}} y_{0 i}}{# (z^{c} \cap B_{k})}),

where on the support the two denominators are the known constants $n_{1 k}$ and $N_{k} - n_{1 k}$ . The weights are the population shares $N_{k} / n$ , fixed before the experiment, not the realized sample shares. That is what makes the estimand $τ$ and not some block-weighted variant of it.

And now the payoff is a repeat of chapter 8’s, one level up. On the support of the design, the blocked estimator is Horvitz–Thompson: the within-block denominators are the constants $n_{1 k}$ and $N_{k} - n_{1 k}$ , the weights $1 / π_{i} = N_{k} / n_{1 k}$ are constant within a block, and grouping the Horvitz–Thompson sum block by block (sum_eq_sum_block, i.e. Finset.sum_fiberwise again) reassembles exactly the block-weighted difference in means. That is blockedDiM_eq_HT_of_mem_support, the stretch exercise posed at the end of the chapter. Given it, unbiasedness is a short corollary:

TheoremThe blocked estimator is unbiased

If every block contains at least one treated and at least one control unit, then $E [{\hat{τ}}_{blocked}] = τ$ under stratified complete randomization — for every population, with no condition on the potential outcomes and no requirement that the treated fractions agree across blocks.

The expectation is a finite sum over the $\prod_{k} (\binom{N_{k}}{n_{1 k}})$ assignments the design can produce, so “unbiased” here is an exact statement about a finite object, not a limit.

New tabOpens with 1,543 lines of library preamble — the code above is at the bottom of the editor.

/-- **The blocked estimator is unbiased for the ATE under stratified randomisation.**

`0 < n₁ k < N_k` in every block -- each block must contain at least one treated and at least
one control unit -- and then `E[blockedDiM] = τ`.

As with `DiM_unbiased_completeRandomization`, nothing is recomputed: the estimator coincides
with Horvitz--Thompson on every assignment the design can produce, and `HT_unbiased` does the
rest.  All that blocking changed was the propensity score. -/
theorem blockedDiM_unbiased (P : Population n)
    (hpos : ∀ k, 0 < n₁ k) (hlt : ∀ k, n₁ k < #(block b k)) :
    (stratifiedDesign b n₁ h).expect (blockedDiM b P) = P.tau := by
  have hprop : ∀ i : Fin n, (stratifiedDesign b n₁ h).propensity i
      = (n₁ (b i) : ℝ) / (#(block b (b i)) : ℝ) := stratified_propensity
  have hsupport : ∀ z, (stratifiedDesign b n₁ h).prob z ≠ 0 →
      blockedDiM b P z = HT (stratifiedDesign b n₁ h) P z := fun z hzprob =>
    blockedDiM_eq_HT_of_mem_support P hpos hlt (stratified_card_inter_block hzprob)
  rw [(stratifiedDesign b n₁ h).expect_congr_of_support hsupport]
  refine HT_unbiased _ P (fun i => ?_) (fun i => ?_)
  · rw [hprop i]
    exact div_pos (by exact_mod_cast hpos (b i))
      (by exact_mod_cast block_card_pos b i)
  · rw [hprop i]
    exact (div_lt_one (by exact_mod_cast block_card_pos b i)).mpr
      (by exact_mod_cast hlt (b i))

expect_congr_of_support swaps the estimator for HT on the support; HT_unbiased needs $0 < π_{i} < 1$ , which after rw [hprop i] is $0 < n_{1 k} / N_{k} < 1$ , i.e. exactly the hypotheses hpos and hlt. Nothing about blocking is re-derived. All that changed was the propensity score.

What is missing: the variance

The curriculum promised the variance of the blocked estimator, and the library does not have it. Informally it is

Var ({\hat{τ}}_{blocked}) = \sum_{k} {(\frac{N_{k}}{n})}^{2} (\frac{S_{1 k}^{2}}{n_{1 k}} + \frac{S_{0 k}^{2}}{N_{k} - n_{1 k}} - \frac{S_{τ k}^{2}}{N_{k}}),

i.e. Neyman’s formula from chapter 9 applied inside each block and combined with squared weights, with no cross-block covariance terms. Every ingredient is already available, and I want to be precise about what is unproven: the missing step is the second-order inclusion probabilities of the blocked design. Within a block, $π_{i j} = n_{1 k} (n_{1 k} - 1) / (N_{k} (N_{k} - 1))$ ; across blocks, $π_{i j} = (n_{1 k} / N_{k}) (n_{1 l} / N_{l})$ — and that second identity is the statement that the blocks are independent, which is what kills the covariance terms. Both follow from card_stratSupport_filter_subset with $S = {i, j}$ , plus choose_ratio_two for the within-block case. What was not written is the case analysis: ${i, j}$ can touch one block or two, and the two cases factor differently.

This is a genuinely good exercise for a reader who has got this far, and it is the one place in this chapter where you should not believe a formula because I wrote it down.

Cluster randomization

Now the map c : Fin n → Fin K groups units into clusters, and treatment is assigned to whole clusters. The unit-level design is the pushforward of a cluster-level design.

Cluster randomization is a pushforwardNew tabOpens with 1,604 lines of library preamble — the code above is at the bottom of the editor.

/-- The unit-level assignment induced by a set `T` of treated clusters: every unit whose
cluster lies in `T` is treated. -/
def clusterAssignment (c : Fin n → Fin K) (T : Assignment K) : Assignment n :=
  univ.filter (fun i => c i ∈ T)

/-- **Cluster randomisation.**  A design `Dc` on the `K` clusters induces a design on the `n`
units by treating every member of every treated cluster.  It is the pushforward of `Dc` along
`clusterAssignment`.

The reason experimenters do this is usually logistical (you cannot treat half a village); the
statistical consequence is that the randomness lives on `K` objects, not `n`. -/
noncomputable def clusterDesign (c : Fin n → Fin K) (Dc : Design K) : Design n :=
  Dc.map (clusterAssignment c)

Everything else follows from one identity: a unit’s treatment indicator is its cluster’s treatment indicator.

A unit's indicator is its cluster's indicatorNew tabOpens with 1,622 lines of library preamble — the code above is at the bottom of the editor.

/-- A unit's treatment indicator is its *cluster's* treatment indicator.  Everything about
cluster randomisation follows from this identity. -/
theorem Z_clusterAssignment (c : Fin n → Fin K) (T : Assignment K) (i : Fin n) :
    Z (clusterAssignment c T) i = Z T (c i) := by
  unfold Z
  by_cases hi : c i ∈ T <;> simp [hi]

unfold Z replaces Z by its definition, if i ∈ z then 1 else 0; by_cases then splits on whether c i ∈ T and both branches are simp. From it, the defining property of cluster randomization:

Treatment is all-or-nothing within a clusterNew tabOpens with 1,629 lines of library preamble — the code above is at the bottom of the editor.

/-- **Treatment is all-or-nothing within a cluster.**  Two units of the same cluster are
either both treated or both untreated -- under *every* assignment the design can produce, and
in fact under every assignment in the image of `clusterAssignment`.

This is the formal content of "the cluster is the unit of randomisation", and the reason the
effective sample size is `K` rather than `n`. -/
theorem cluster_all_or_nothing {c : Fin n → Fin K} {Dc : Design K} {z : Assignment n}
    (hz : (clusterDesign c Dc).prob z ≠ 0) {i j : Fin n} (hij : c i = c j) :
    Z z i = Z z j := by
  obtain ⟨T, hT, -⟩ := Design.exists_of_map_prob_ne_zero hz
  subst hT
  rw [Z_clusterAssignment, Z_clusterAssignment, hij]

The proof uses Design.exists_of_map_prob_ne_zero: an assignment in the support of a pushforward comes from some upstream assignment, and obtain names it T. subst then replaces z by clusterAssignment c T everywhere, at which point both indicators are cluster-level and hij closes it.

A unit's propensity is its cluster's propensityNew tabOpens with 1,642 lines of library preamble — the code above is at the bottom of the editor.

/-- **A unit's propensity is its cluster's propensity.**  There is no dilution and no
inflation: `π i = π_cluster (c i)`.

So a cluster design has *exactly the same first-order inclusion probabilities* as the
cluster-level design it came from.  Everything that distinguishes cluster randomisation from
unit randomisation -- the loss of precision -- is invisible at this level and lives entirely
in the joint propensities: two units of the same cluster have `π_ij = π_i`, perfect positive
dependence. -/
theorem cluster_propensity (c : Fin n → Fin K) (Dc : Design K) (i : Fin n) :
    (clusterDesign c Dc).propensity i = Dc.propensity (c i) := by
  rw [clusterDesign, Design.map_propensity]
  exact Dc.expect_congr fun T => Z_clusterAssignment c T i

There is no dilution and no inflation: $π_{i} = π_{c (i)}^{cluster}$ . So a cluster design has exactly the same first-order inclusion probabilities as the cluster-level design it came from, and unbiasedness transfers with no new counting at all:

Horvitz–Thompson under cluster randomizationNew tabOpens with 1,675 lines of library preamble — the code above is at the bottom of the editor.

/-- **Horvitz--Thompson is unbiased under cluster randomisation.**

Take the cluster-level design to be complete randomisation of `K₁` out of `K` clusters, with
`0 < K₁ < K`.  Then every unit has propensity `K₁ / K`, which is strictly between `0` and
`1`, and `HT_unbiased` applies verbatim.

Unbiasedness is therefore *free*: clustering costs precision, never validity.  (The variance
is a different story, and is where the design effect / intra-cluster correlation lives.) -/
theorem clusterHT_unbiased (c : Fin n → Fin K) {K₁ : ℕ} (hle : K₁ ≤ K) (P : Population n)
    (hpos : 0 < K₁) (hlt : K₁ < K) :
    (clusterDesign c (completeRandomization K K₁ hle)).expect
      (HT (clusterDesign c (completeRandomization K K₁ hle)) P) = P.tau := by
  have hKR : (0 : ℝ) < (K : ℝ) := by exact_mod_cast Nat.lt_of_lt_of_le hpos hle
  have hK1R : (0 : ℝ) < (K₁ : ℝ) := by exact_mod_cast hpos
  have hltR : (K₁ : ℝ) < (K : ℝ) := by exact_mod_cast hlt
  have hprop : ∀ i : Fin n,
      (clusterDesign c (completeRandomization K K₁ hle)).propensity i = (K₁ : ℝ) / (K : ℝ) :=
    fun i => by rw [cluster_propensity, completeRandomization_propensity]
  refine HT_unbiased _ P (fun i => ?_) (fun i => ?_)
  · rw [hprop i]; exact div_pos hK1R hKR
  · rw [hprop i]; exact (div_lt_one hKR).mpr hltR

Read that proof and notice what is absent. There is no combinatorics, no Nat.choose, no bijection. cluster_propensity reduces the propensity to completeRandomization_propensity, which was proved in chapter 7, and HT_unbiased finishes. Clustering costs precision, never validity — and the precision loss is invisible at this level, because it lives entirely in the joint propensities: two units of the same cluster have $π_{i j} = π_{i}$ , perfect positive dependence, which is where the design effect and the intra-cluster correlation come from. Three of the exercises below make that precise — the last of them computes $Cov (Z_{i}, Z_{j}) = π_{i} (1 - π_{i})$ for a within-cluster pair, the largest a covariance of two indicators with that propensity can be. The variance calculation that turns those covariances into a design effect is not in this library.

Exercises

ExerciseThe quotas fit (warm-up)

A blocked design cannot ask for more treated units than exist. Chain Finset.sum_le_sum (which needs fun k _ => h k, the hypothesis at each block) with sum_card_block, using calc for the two links.

New tabOpens with 1,720 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (warm-up).  A blocked design cannot ask for more treated units than there are
units: the block quotas add up to at most the population size.  Every statement about the
*total* number of treated units under blocking needs this. -/
theorem sum_n₁_le {b : Fin n → Fin K} {n₁ : Fin K → ℕ} (h : ∀ k, n₁ k ≤ #(block b k)) :
    ∑ k, n₁ k ≤ n := by
  sorry

Show solution

New tabOpens with 1,720 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (warm-up).  A blocked design cannot ask for more treated units than there are
units: the block quotas add up to at most the population size.  Every statement about the
*total* number of treated units under blocking needs this. -/
theorem sum_n₁_le {b : Fin n → Fin K} {n₁ : Fin K → ℕ} (h : ∀ k, n₁ k ≤ #(block b k)) :
    ∑ k, n₁ k ≤ n := by
  calc ∑ k, n₁ k ≤ ∑ k, #(block b k) := Finset.sum_le_sum fun k _ => h k
    _ = n := sum_card_block b

ExerciseCounting the survivors (warm-up)

The five-unit blocked design of the biased example has $(\binom{2}{1}) (\binom{3}{2}) = 6$ assignments in its support, whereas complete randomization of three units out of five has $(\binom{5}{3}) = 10$ . Deleting the four unbalanced assignments is the blocking. card_stratSupport turns the count into the product; decide evaluates it.

New tabOpens with 1,821 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (warm-up).  **Blocking deletes assignments.**  Complete randomisation of three
units out of five has `C(5, 3) = 10` assignments in its support; the blocked design of the
previous example has only `C(2,1) · C(3,2) = 6`.  The four that are gone are exactly the ones
that are unbalanced across the two blocks.

Hint: `card_stratSupport` turns the count into a product of binomial coefficients, and
`decide` evaluates the rest in the kernel. -/
theorem card_stratSupport_example : #(stratSupport bEx n₁Ex) = 6 := by
  sorry

Show solution

New tabOpens with 1,821 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (warm-up).  **Blocking deletes assignments.**  Complete randomisation of three
units out of five has `C(5, 3) = 10` assignments in its support; the blocked design of the
previous example has only `C(2,1) · C(3,2) = 6`.  The four that are gone are exactly the ones
that are unbalanced across the two blocks.

Hint: `card_stratSupport` turns the count into a product of binomial coefficients, and
`decide` evaluates the rest in the kernel. -/
theorem card_stratSupport_example : #(stratSupport bEx n₁Ex) = 6 := by
  rw [card_stratSupport]
  decide

ExerciseProportional allocation makes the propensity constant (core)

The hypothesis ∀ k, n₁ k * n = (∑ l, n₁ l) * #(block b k) is the cross-multiplied form of $n_{1 k} / N_{k} = N_{1} / n$ : every block gets its population share. Show the propensity is then the constant $N_{1} / n$ , which is precisely the hbal that the stretch exercise below assumes.

Rewrite with stratified_propensity i, then div_eq_div_iff — whose two side conditions are block_card_pos and Fin.pos i, since a unit lives in a nonempty block and in a nonempty population — and close with exact_mod_cast, because the hypothesis is about ℕ and the goal about ℝ.

New tabOpens with 1,834 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (core).  **Proportional allocation is what makes the propensity constant.**

`stratified_propensity` says `π i = n₁ (b i) / N_(b i)`, a number that varies by block.  It is
constant — and equal to the overall treated fraction `N₁ / n`, `N₁ = ∑ k, n₁ k` — exactly when
every block is allocated its population share, i.e. when `n₁ k · n = N₁ · N_k` for every `k`.
That is the hypothesis an applied paper writes as "we allocated treatment proportionally".

Combined with `expect_DiM_stratified_of_balanced`, this is the precise sense in which
proportional allocation rescues the plain difference in means: blocking does not break it,
disproportionate allocation does.

Hint: `stratified_propensity i` rewrites the goal into an equality of two ratios; then
`div_eq_div_iff` (both denominators are positive — `block_card_pos` and `Fin.pos i`) turns it
into the cross-multiplied form, which is `hprop (b i)` up to `ℕ → ℝ` coercions. -/
theorem stratified_propensity_of_proportional {b : Fin n → Fin K} {n₁ : Fin K → ℕ}
    (h : ∀ k, n₁ k ≤ #(block b k)) (hprop : ∀ k, n₁ k * n = (∑ l, n₁ l) * #(block b k))
    (i : Fin n) :
    (stratifiedDesign b n₁ h).propensity i = ((∑ k, n₁ k : ℕ) : ℝ) / (n : ℝ) := by
  sorry

Show solution

New tabOpens with 1,834 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (core).  **Proportional allocation is what makes the propensity constant.**

`stratified_propensity` says `π i = n₁ (b i) / N_(b i)`, a number that varies by block.  It is
constant — and equal to the overall treated fraction `N₁ / n`, `N₁ = ∑ k, n₁ k` — exactly when
every block is allocated its population share, i.e. when `n₁ k · n = N₁ · N_k` for every `k`.
That is the hypothesis an applied paper writes as "we allocated treatment proportionally".

Combined with `expect_DiM_stratified_of_balanced`, this is the precise sense in which
proportional allocation rescues the plain difference in means: blocking does not break it,
disproportionate allocation does.

Hint: `stratified_propensity i` rewrites the goal into an equality of two ratios; then
`div_eq_div_iff` (both denominators are positive — `block_card_pos` and `Fin.pos i`) turns it
into the cross-multiplied form, which is `hprop (b i)` up to `ℕ → ℝ` coercions. -/
theorem stratified_propensity_of_proportional {b : Fin n → Fin K} {n₁ : Fin K → ℕ}
    (h : ∀ k, n₁ k ≤ #(block b k)) (hprop : ∀ k, n₁ k * n = (∑ l, n₁ l) * #(block b k))
    (i : Fin n) :
    (stratifiedDesign b n₁ h).propensity i = ((∑ k, n₁ k : ℕ) : ℝ) / (n : ℝ) := by
  have hN : (0 : ℝ) < (#(block b (b i)) : ℝ) := by exact_mod_cast block_card_pos b i
  have hn : (0 : ℝ) < (n : ℝ) := by exact_mod_cast Fin.pos i
  rw [stratified_propensity i, div_eq_div_iff hN.ne' hn.ne']
  exact_mod_cast hprop (b i)

ExerciseThe price of clustering, in one covariance (core)

For two units of the same cluster, $Cov (Z_{i}, Z_{j}) = π (1 - π)$ — the maximum a covariance of two indicators with propensity $π$ can be, because they are the same random variable. This is the ingredient a design effect is made of.

Design.cov_eq_expect_mul_sub_mul turns the covariance into $E [Z_{i} Z_{j}] - E [Z_{i}] E [Z_{j}]$ . The three expectations are the joint propensity and the two propensities by definition, so each have is closed by rfl; then cluster_jointPropensity_same and cluster_propensity do the work.

New tabOpens with 1,859 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (core).  **Two units of the same cluster are perfectly dependent.**

Under any cluster-randomised design, the covariance of the treatment indicators of two units
sharing a cluster is `π (1 - π)`, the largest value a covariance of two indicators with
propensity `π` can take: they are the *same* Bernoulli variable, so the correlation is one.

This is where the design effect comes from.  Nothing in the chapter computes a variance, but
this lemma is the ingredient that a variance calculation would consume: the variance of a
total over a cluster design carries one such term for every within-cluster pair, and there is
no configuration of the design that can make those terms smaller.

Hint: `Design.cov_eq_expect_mul_sub_mul` rewrites the covariance as `E[Z_i Z_j] - E[Z_i]E[Z_j]`.
The three expectations are the joint propensity and the two propensities *by definition*, so
`rfl` proves each `have`; then `cluster_jointPropensity_same` and `cluster_propensity` finish. -/
theorem cluster_cov_Z_same (c : Fin n → Fin K) (Dc : Design K) {i j : Fin n} (hij : c i = c j) :
    (clusterDesign c Dc).cov (fun z => Z z i) (fun z => Z z j)
      = (clusterDesign c Dc).propensity i * (1 - (clusterDesign c Dc).propensity i) := by
  sorry

Show solution

New tabOpens with 1,859 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (core).  **Two units of the same cluster are perfectly dependent.**

Under any cluster-randomised design, the covariance of the treatment indicators of two units
sharing a cluster is `π (1 - π)`, the largest value a covariance of two indicators with
propensity `π` can take: they are the *same* Bernoulli variable, so the correlation is one.

This is where the design effect comes from.  Nothing in the chapter computes a variance, but
this lemma is the ingredient that a variance calculation would consume: the variance of a
total over a cluster design carries one such term for every within-cluster pair, and there is
no configuration of the design that can make those terms smaller.

Hint: `Design.cov_eq_expect_mul_sub_mul` rewrites the covariance as `E[Z_i Z_j] - E[Z_i]E[Z_j]`.
The three expectations are the joint propensity and the two propensities *by definition*, so
`rfl` proves each `have`; then `cluster_jointPropensity_same` and `cluster_propensity` finish. -/
theorem cluster_cov_Z_same (c : Fin n → Fin K) (Dc : Design K) {i j : Fin n} (hij : c i = c j) :
    (clusterDesign c Dc).cov (fun z => Z z i) (fun z => Z z j)
      = (clusterDesign c Dc).propensity i * (1 - (clusterDesign c Dc).propensity i) := by
  have hjoint : (clusterDesign c Dc).expect (fun z => Z z i * Z z j)
      = (clusterDesign c Dc).jointPropensity i j := rfl
  have hi : (clusterDesign c Dc).expect (fun z => Z z i)
      = (clusterDesign c Dc).propensity i := rfl
  have hj : (clusterDesign c Dc).expect (fun z => Z z j)
      = (clusterDesign c Dc).propensity j := rfl
  have hp : (clusterDesign c Dc).propensity j = (clusterDesign c Dc).propensity i := by
    rw [cluster_propensity, cluster_propensity, hij]
  rw [Design.cov_eq_expect_mul_sub_mul, hjoint, hi, hj, hp,
    cluster_jointPropensity_same c Dc hij]
  ring

ExerciseWhen difference in means survives blocking (stretch)

The converse reading of the expectation formula: if every unit has the same propensity $N_{1} / n$ , plain difference in means is unbiased after all. Rewrite with expect_DiM_stratified and Population.tau_eq_mean_sub_mean, pull the constant out of each sum with ← Finset.mul_sum, unfold mean, and let field_simp clear the three denominators — which is why you need $0 < N_{1}$ , $N_{1} < n$ as real-number facts first.

New tabOpens with 1,758 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (stretch).  The converse reading of `expect_DiM_stratified`: when the blocked
design happens to give every unit the *same* propensity — the treated fractions `n₁ k / N_k`
agree across blocks — the plain difference in means is unbiased after all.

Blocking does not break `DiM` by itself; unequal treated fractions do. -/
theorem expect_DiM_stratified_of_balanced {b : Fin n → Fin K} {n₁ : Fin K → ℕ}
    (h : ∀ k, n₁ k ≤ #(block b k)) (P : Population n)
    (hpos : 0 < ∑ k, n₁ k) (hlt : ∑ k, n₁ k < n)
    (hbal : ∀ i, (stratifiedDesign b n₁ h).propensity i = ((∑ k, n₁ k : ℕ) : ℝ) / (n : ℝ)) :
    (stratifiedDesign b n₁ h).expect (DiM P) = P.tau := by
  sorry

Show solution

New tabOpens with 1,758 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (stretch).  The converse reading of `expect_DiM_stratified`: when the blocked
design happens to give every unit the *same* propensity — the treated fractions `n₁ k / N_k`
agree across blocks — the plain difference in means is unbiased after all.

Blocking does not break `DiM` by itself; unequal treated fractions do. -/
theorem expect_DiM_stratified_of_balanced {b : Fin n → Fin K} {n₁ : Fin K → ℕ}
    (h : ∀ k, n₁ k ≤ #(block b k)) (P : Population n)
    (hpos : 0 < ∑ k, n₁ k) (hlt : ∑ k, n₁ k < n)
    (hbal : ∀ i, (stratifiedDesign b n₁ h).propensity i = ((∑ k, n₁ k : ℕ) : ℝ) / (n : ℝ)) :
    (stratifiedDesign b n₁ h).expect (DiM P) = P.tau := by
  have hnR : (0 : ℝ) < (n : ℝ) := by exact_mod_cast Nat.lt_of_le_of_lt (Nat.zero_le _) hlt
  have hN1 : (0 : ℝ) < ((∑ k, n₁ k : ℕ) : ℝ) := by exact_mod_cast hpos
  have hN0 : (0 : ℝ) < (n : ℝ) - ((∑ k, n₁ k : ℕ) : ℝ) := by
    have : ((∑ k, n₁ k : ℕ) : ℝ) < (n : ℝ) := by exact_mod_cast hlt
    linarith
  rw [expect_DiM_stratified h P, P.tau_eq_mean_sub_mean]
  simp only [hbal, ← Finset.mul_sum, mean]
  field_simp

Seven more come from lean/MrCLean/Blocking.lean itself — the library file the rest of the chapter quotes, which poses its own exercises inline.

ExerciseDistinct blocks are disjoint (warm-up)

Half of “the blocks partition the population”. Finset.disjoint_left reduces disjointness to: no i is in both. A unit in block b k has b i = k (mem_block), so being in two blocks makes k = l, contradicting hkl.

New tabOpens with 1,145 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (warm-up).  Distinct blocks are disjoint -- a unit has only one value of
`b i`.  Together with `sum_card_block` this is the statement that the blocks *partition* the
population. -/
theorem block_disjoint (b : Fin n → Fin K) {k l : Fin K} (hkl : k ≠ l) :
    Disjoint (block b k) (block b l) := by
  sorry

Show solution

New tabOpens with 1,145 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (warm-up).  Distinct blocks are disjoint -- a unit has only one value of
`b i`.  Together with `sum_card_block` this is the statement that the blocks *partition* the
population. -/
theorem block_disjoint (b : Fin n → Fin K) {k l : Fin K} (hkl : k ≠ l) :
    Disjoint (block b k) (block b l) := by
  refine Finset.disjoint_left.mpr fun i hik hil => ?_
  exact hkl ((mem_block.mp hik).symm.trans (mem_block.mp hil))

ExerciseThe block sizes add up (core)

The other half: $\sum_{k} N_{k} = n$ . Finset.card_eq_sum_card_fiberwise is the statement that counting a set by fibres of a function gives the same answer; instantiate it at b with s = t = univ, then simpa with block, Finset.card_univ and Fintype.card_fin to line the two sides up.

New tabOpens with 1,153 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (core).  The block sizes add up to the population size, `∑ k, N_k = n`. -/
theorem sum_card_block (b : Fin n → Fin K) : ∑ k, #(block b k) = n := by
  sorry

Show solution

New tabOpens with 1,153 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (core).  The block sizes add up to the population size, `∑ k, N_k = n`. -/
theorem sum_card_block (b : Fin n → Fin K) : ∑ k, #(block b k) = n := by
  have hcount := Finset.card_eq_sum_card_fiberwise (f := b) (s := (univ : Finset (Fin n)))
    (t := (univ : Finset (Fin K))) (fun i _ => Finset.mem_univ (b i))
  simpa only [Finset.card_univ, Fintype.card_fin, block] using hcount.symm

ExerciseHorvitz–Thompson under a blocked design (warm-up)

The last two lines of blockedDiM_unbiased, run on their own. Apply HT_unbiased and discharge its two hypotheses from stratified_propensity: $0 < n_{1 k} / N_{k}$ is div_pos, and $n_{1 k} / N_{k} < 1$ is div_lt_one. Both side conditions are $N$ -to- $R$ casts, so exact_mod_cast.

New tabOpens with 1,568 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (warm-up).  Under a blocked design, the Horvitz--Thompson estimator is also
unbiased -- a one-line consequence of `HT_unbiased` once the propensity score is known.  It is
the *same* statement as `blockedDiM_unbiased` in disguise, since the two estimators agree on
the support. -/
theorem stratifiedHT_unbiased (P : Population n)
    (hpos : ∀ k, 0 < n₁ k) (hlt : ∀ k, n₁ k < #(block b k)) :
    (stratifiedDesign b n₁ h).expect (HT (stratifiedDesign b n₁ h) P) = P.tau := by
  sorry

Show solution

New tabOpens with 1,568 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (warm-up).  Under a blocked design, the Horvitz--Thompson estimator is also
unbiased -- a one-line consequence of `HT_unbiased` once the propensity score is known.  It is
the *same* statement as `blockedDiM_unbiased` in disguise, since the two estimators agree on
the support. -/
theorem stratifiedHT_unbiased (P : Population n)
    (hpos : ∀ k, 0 < n₁ k) (hlt : ∀ k, n₁ k < #(block b k)) :
    (stratifiedDesign b n₁ h).expect (HT (stratifiedDesign b n₁ h) P) = P.tau := by
  refine HT_unbiased _ P (fun i => ?_) (fun i => ?_)
  · rw [stratified_propensity i]
    exact div_pos (by exact_mod_cast hpos (b i)) (by exact_mod_cast block_card_pos b i)
  · rw [stratified_propensity i]
    exact (div_lt_one (by exact_mod_cast block_card_pos b i)).mpr
      (by exact_mod_cast hlt (b i))

ExerciseThe total treated count is fixed too (core)

On the support, the overall number treated is $\sum_{k} n_{1 k}$ — a blocked design is a refinement of complete randomization. Group the count by block with Finset.sum_fiberwise applied to the indicator fun i => if i ∈ z then 1 else 0, then apply stratified_card_inter_block in each block.

New tabOpens with 1,582 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (core).  A sanity check on the design: the *total* number of treated units is
the same on every assignment the design can produce, namely `∑ k, n₁ k`.  Blocking fixes not
only the block-level counts but, a fortiori, the overall count -- a blocked design is a
refinement of complete randomisation. -/
theorem stratified_card_eq_of_prob_ne_zero {z : Assignment n}
    (hz : (stratifiedDesign b n₁ h).prob z ≠ 0) : #z = ∑ k, n₁ k := by
  sorry

Show solution

New tabOpens with 1,582 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (core).  A sanity check on the design: the *total* number of treated units is
the same on every assignment the design can produce, namely `∑ k, n₁ k`.  Blocking fixes not
only the block-level counts but, a fortiori, the overall count -- a blocked design is a
refinement of complete randomisation. -/
theorem stratified_card_eq_of_prob_ne_zero {z : Assignment n}
    (hz : (stratifiedDesign b n₁ h).prob z ≠ 0) : #z = ∑ k, n₁ k := by
  have hfib : ∑ k, #(z ∩ block b k) = #z := by
    have := Finset.sum_fiberwise (M := ℕ) univ b (fun i => if i ∈ z then 1 else 0)
    simpa only [Finset.sum_ite_mem, Finset.card_eq_sum_ones, block, Finset.univ_inter,
      Finset.inter_comm z] using this
  rw [← hfib]
  exact Finset.sum_congr rfl fun k _ => stratified_card_inter_block hz k

ExerciseJoint propensities pass to the clusters (warm-up)

Nothing is lost or gained by passing to the unit level. Unfold Design.jointPropensity and clusterDesign, push the expectation through the pushforward with Design.expect_map, and finish with Design.expect_congr and two uses of Z_clusterAssignment.

New tabOpens with 1,666 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (warm-up).  The joint propensity of two units is the joint propensity of their
two clusters -- again nothing is lost or gained by passing to the unit level.  For units in
different clusters this is the cluster-level second-order inclusion probability; for units in
the same cluster it degenerates, as `cluster_jointPropensity_same` records. -/
theorem cluster_jointPropensity_diff (c : Fin n → Fin K) (Dc : Design K) (i j : Fin n) :
    (clusterDesign c Dc).jointPropensity i j = Dc.jointPropensity (c i) (c j) := by
  sorry

Show solution

New tabOpens with 1,666 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (warm-up).  The joint propensity of two units is the joint propensity of their
two clusters -- again nothing is lost or gained by passing to the unit level.  For units in
different clusters this is the cluster-level second-order inclusion probability; for units in
the same cluster it degenerates, as `cluster_jointPropensity_same` records. -/
theorem cluster_jointPropensity_diff (c : Fin n → Fin K) (Dc : Design K) (i j : Fin n) :
    (clusterDesign c Dc).jointPropensity i j = Dc.jointPropensity (c i) (c j) := by
  rw [Design.jointPropensity, clusterDesign, Design.expect_map, Design.jointPropensity]
  exact Dc.expect_congr fun T => by rw [Z_clusterAssignment, Z_clusterAssignment]

ExerciseTwo units of one cluster (core)

$π_{i j} = π_{i}$ when $c (i) = c (j)$ : the joint propensity does not factorize at all. The same rewrite chain as the previous exercise, with one extra step — after hij turns both indicators into the same one, Z_mul_self says an indicator times itself is itself. These two are the pair that make the intra-cluster dependence precise.

New tabOpens with 1,655 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (core).  Two units in the same cluster are treated together with probability
equal to their common propensity: the joint propensity does not factorise at all.  Compare
`bernoulli_jointPropensity`, where it factorises exactly. -/
theorem cluster_jointPropensity_same (c : Fin n → Fin K) (Dc : Design K) {i j : Fin n}
    (hij : c i = c j) :
    (clusterDesign c Dc).jointPropensity i j = (clusterDesign c Dc).propensity i := by
  sorry

Show solution

New tabOpens with 1,655 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (core).  Two units in the same cluster are treated together with probability
equal to their common propensity: the joint propensity does not factorise at all.  Compare
`bernoulli_jointPropensity`, where it factorises exactly. -/
theorem cluster_jointPropensity_same (c : Fin n → Fin K) (Dc : Design K) {i j : Fin n}
    (hij : c i = c j) :
    (clusterDesign c Dc).jointPropensity i j = (clusterDesign c Dc).propensity i := by
  rw [cluster_propensity, Design.jointPropensity, clusterDesign, Design.expect_map,
    Design.propensity]
  exact Dc.expect_congr fun T => by
    rw [Z_clusterAssignment, Z_clusterAssignment, hij, Z_mul_self]

ExerciseThe blocked estimator is Horvitz–Thompson (stretch)

The one that carries the chapter: the blocked analogue of DiM_eq_HT_of_card, and the lemma blockedDiM_unbiased is a corollary of. Everything the proof needs is named in its docstring — stratified_card_inter_block for the within-block counts, stratified_propensity for the weights, sum_eq_sum_block for the regrouping. The work is bookkeeping with Nat.cast_sub and getting field_simp to accept the two inverse-probability weights.

New tabOpens with 1,496 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (stretch).  On the support of the blocked design, the blocked estimator **is**
the Horvitz--Thompson estimator.

This is the blocked analogue of `DiM_eq_HT_of_card`.  The within-block denominators `n₁ k`
and `N_k - n₁ k` are constants on the support (`stratified_card_inter_block`), and the
inverse-probability weights `1 / π i = N_k / n₁ k` are constant within a block
(`stratified_propensity`), so grouping the Horvitz--Thompson sum block by block
(`sum_eq_sum_block`) produces exactly the block-weighted difference in means. -/
theorem blockedDiM_eq_HT_of_mem_support (P : Population n)
    (hpos : ∀ k, 0 < n₁ k) (hlt : ∀ k, n₁ k < #(block b k))
    {z : Assignment n} (hz : ∀ k, #(z ∩ block b k) = n₁ k) :
    blockedDiM b P z = HT (stratifiedDesign b n₁ h) P z := by
  sorry

Show solution

New tabOpens with 1,496 lines of library preamble — the code above is at the bottom of the editor.

/-- Exercise (stretch).  On the support of the blocked design, the blocked estimator **is**
the Horvitz--Thompson estimator.

This is the blocked analogue of `DiM_eq_HT_of_card`.  The within-block denominators `n₁ k`
and `N_k - n₁ k` are constants on the support (`stratified_card_inter_block`), and the
inverse-probability weights `1 / π i = N_k / n₁ k` are constant within a block
(`stratified_propensity`), so grouping the Horvitz--Thompson sum block by block
(`sum_eq_sum_block`) produces exactly the block-weighted difference in means. -/
theorem blockedDiM_eq_HT_of_mem_support (P : Population n)
    (hpos : ∀ k, 0 < n₁ k) (hlt : ∀ k, n₁ k < #(block b k))
    {z : Assignment n} (hz : ∀ k, #(z ∩ block b k) = n₁ k) :
    blockedDiM b P z = HT (stratifiedDesign b n₁ h) P z := by
  have hNne : ∀ k, ((#(block b k) : ℕ) : ℝ) ≠ 0 := fun k =>
    Nat.cast_ne_zero.mpr (Nat.lt_of_le_of_lt (Nat.zero_le _) (hlt k)).ne'
  have hn1ne : ∀ k, ((n₁ k : ℕ) : ℝ) ≠ 0 := fun k => Nat.cast_ne_zero.mpr (hpos k).ne'
  have hdiff : ∀ k, ((#(block b k) : ℕ) : ℝ) - ((n₁ k : ℕ) : ℝ) ≠ 0 := fun k => by
    have : ((n₁ k : ℕ) : ℝ) < ((#(block b k) : ℕ) : ℝ) := by exact_mod_cast hlt k
    linarith
  have hcompl : ∀ k, ((#(zᶜ ∩ block b k) : ℕ) : ℝ)
      = ((#(block b k) : ℕ) : ℝ) - ((n₁ k : ℕ) : ℝ) := by
    intro k
    rw [card_compl_inter, hz k, Nat.cast_sub ((hlt k).le)]
  -- the Horvitz--Thompson summand, block by block
  have hterm : ∀ (k : Fin K) (i : Fin n), i ∈ block b k →
      Z z i * P.y1 i / (stratifiedDesign b n₁ h).propensity i
        - (1 - Z z i) * P.y0 i / (1 - (stratifiedDesign b n₁ h).propensity i)
      = ((#(block b k) : ℕ) : ℝ)
          * (Z z i * P.y1 i / (n₁ k : ℝ)
            - (1 - Z z i) * P.y0 i / (((#(block b k) : ℕ) : ℝ) - (n₁ k : ℝ))) := by
    intro k i hi
    have hbi : b i = k := mem_block.mp hi
    rw [stratified_propensity i, hbi]
    rw [div_div_eq_mul_div, one_sub_div (hNne k)]
    field_simp
  have hblock : ∀ k : Fin K,
      ∑ i ∈ block b k, (Z z i * P.y1 i / (stratifiedDesign b n₁ h).propensity i
        - (1 - Z z i) * P.y0 i / (1 - (stratifiedDesign b n₁ h).propensity i))
      = ((#(block b k) : ℕ) : ℝ) * blockDiM b P z k := by
    intro k
    rw [Finset.sum_congr rfl fun i hi => hterm k i hi, ← Finset.mul_sum,
      Finset.sum_sub_distrib, ← Finset.sum_div, ← Finset.sum_div,
      ← sum_inter_eq_sum_Z_mul, ← sum_compl_inter_eq_sum_one_sub_Z_mul]
    rw [blockDiM, hz k, hcompl k]
  unfold HT blockedDiM
  rw [← sum_eq_sum_block b, Finset.sum_congr rfl fun k _ => hblock k, Finset.mul_sum]
  exact Finset.sum_congr rfl fun k _ => by ring

Reading autoformalized Lean

Three ways to get this chapter wrong while compiling. The first two are statements that look like the results above and are not; the third is a statement that is true and says nothing.

The blocking function that does nothing

-- WRONG: the "blocked" design is complete randomization in disguise
noncomputable def stratifiedDesign' (b : Fin n → Fin K) (n₁ : Fin K → ℕ)
    (h : ∀ k, n₁ k ≤ #(block b k)) : Design n where
  prob z := if #z = ∑ k, n₁ k then (n.choose (∑ k, n₁ k) : ℝ)⁻¹ else 0
  nonneg := ...
  sum_one := ...

theorem stratified'_propensity (i : Fin n) :
    (stratifiedDesign' b n₁ h).propensity i = (∑ k, n₁ k : ℝ) / n := ...

Fill in the two elided proofs and everything type-checks, the theorem is true, and b appears in the signature three times. It appears nowhere in the body of the definition. The support constrains only the total number treated, so this is completeRandomization n (∑ k, n₁ k) with a blocking function bolted on for decoration — and the propensity comes out constant, which is the tell. A real blocked design has ∀ k, #(z ∩ block b k) = n₁ k, one constraint per block; the chapter’s stratified_card_eq_of_prob_ne_zero shows that the total-count condition is a consequence of it, never the other way round.

Clusters that pretend to be independent

-- WRONG: no hypothesis relating i and j
theorem cluster_jointPropensity' (c : Fin n → Fin K) (Dc : Design K) (i j : Fin n) :
    (clusterDesign c Dc).jointPropensity i j
      = Dc.propensity (c i) * Dc.propensity (c j) := ...

This is the statement everybody wants to be true, because it is what makes the variance of a cluster experiment look like the variance of a unit experiment. It is false, and it is false in exactly the case that matters: when $c (i) = c (j)$ the two indicators are the same random variable, so $π_{i j} = E [Z_{i}^{2}] = E [Z_{i}] = π_{i}$ , which equals $π_{i}^{2}$ only when $π_{i} \in {0, 1}$ . The library states the two cases separately — cluster_jointPropensity_same and cluster_jointPropensity_diff — and neither claims factorization; factorization is a property of the cluster-level design, and only bernoulli_jointPropensity actually has it.

A hypothesis that is satisfied by nothing

Finally, a milder version of the same disease. Consider blockedDiM_unbiased’s hypothesis ∀ k, n₁ k < #(block b k). If any block is empty this says n₁ k < 0 in ℕ, which is unsatisfiable, so for such a b the theorem is vacuously true and tells you nothing. The theorem is not wrong — an empty stratum identifies nothing, and refusing to talk about it is correct — but “this compiles” and “this applies to my experiment” are different claims, and the second one requires you to check that the hypotheses can be met by the design you actually ran.