Blocking and clustering

The designs considered so far assign treatment directly to individual units. Two common alternatives are:

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.

Both designs can reuse the Horvitz–Thompson theorem from chapter 8. The new work is to construct each design and compute its inclusion probabilities. Unbiasedness then follows from the general theorem.

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 810 lines of library preamble — the code above is at the bottom of the editor.

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.

```
prob z := ∑ w ∈ univ.filter (fun w => f w = z), D.prob w
```
The pmf. univ.filter (fun w => f w = z) is the fibre $f^{- 1} (z)$ written as a Finset, so the probability the pushforward assigns to $z$ is the total upstream probability of the assignments that map to $z$ .
```
nonneg _ := Finset.sum_nonneg fun w _ => D.nonneg w
```
Nonnegativity, which is free: Finset.sum_nonneg says a sum of nonnegative terms is nonnegative, and every term is nonnegative because D is already a design. The underscore is the assignment $z$ , which this argument never looks at.
```
sum_one := (Finset.sum_fiberwise univ f D.prob).trans D.sum_one
```
Summing to one, the only real obligation, and it is a Mathlib lemma. Finset.sum_fiberwise univ f D.prob states $\sum_{z} \sum_{w : f (w) = z} D (w) = \sum_{w} D (w)$ : 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. .trans chains that equation with D.sum_one, which says the right-hand side is $1$ .

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

Proposition (change of variables). Let $D$ be a design on $m$ units, let $f$ map assignments of those $m$ units to assignments of $n$ units, and let $g$ be any statistic of an $n$ -unit assignment. Then

E_{f_{*} D} [g] = E_{D} [g \circ f] .

Proof.

Write $F_{z} = {w : f (w) = z}$ for the fibre over $z$ . By definition $(f_{*} D) (z) = \sum_{w \in F_{z}} D (w)$ , so

E_{f_{*} D} [g] = \sum_{z} (f_{*} D) (z) g (z) = \sum_{z} \sum_{w \in F_{z}} D (w) g (z) .

Inside the fibre $F_{z}$ we have $f (w) = z$ , so $g (z) = g (f (w))$ there and the inner sum is $\sum_{w \in F_{z}} D (w) g (f (w))$ . The fibres partition the $m$ -unit assignments, so summing over $z$ and then over $F_{z}$ is summing over all $w$ once, and the double sum collapses to $\sum_{w} D (w) g (f (w)) = E_{D} [g \circ f]$ .∎

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

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.

The Lean follows the informal proof exactly: have key is the per- $z$ rewriting of one term, and the calc block — Lean’s way of writing a chain of equalities with a justification for each link — is the two displayed equations.

intro z

Fix the downstream assignment $z$ ; the claim of key is that one term of the outer sum, the probability of $z$ times $g (z)$ , is the fibre sum with $g$ evaluated upstream.

Goal

f : Assignment m → Assignment n
D : Design m
g : Assignment n → ℝ
z : Assignment n
⊢ (Design.map f D).prob z * g z = ∑ w with f w = z, D.prob w * g (f w)

```
rw [map_prob, …]
```
The first of the two rewrites: map_prob is the @[simp] lemma recording that the pushforward’s pmf is the fibre sum, so the left-hand side becomes (∑ w ∈ fibre, D.prob w) * g z.
Goal
```
⊢ (∑ w with f w = z, D.prob w) * g z = ∑ w with f w = z, D.prob w * g (f w)
```
```
rw […, Finset.sum_mul]
```
The second: Finset.sum_mul, which states $(\sum_{i} a_{i}) c = \sum_{i} a_{i} c$ , pushes the factor g z inside the sum. The two sides now differ only in whether $g$ is applied to z or to f w.
Goal
```
⊢ ∑ i with f i = z, D.prob i * g z = ∑ w with f w = z, D.prob w * g (f w)
```
```
exact Finset.sum_congr rfl fun w hw => by rw [(Finset.mem_filter.mp hw).2]
```
Finset.sum_congr rfl reduces an equality of two sums over the same Finset to an equality of their terms, giving one w in the fibre at a time. Finset.mem_filter.mp hw splits membership in univ.filter (f · = z) into “in univ” and the filter condition, and .2 is the second half, f w = z; rewriting with it turns g z into g (f w). This is the step where the fibre does its work.
```
= ∑ 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
```
First calc link: the expectation is by definition $\sum_{z} (f_{*} D) (z) g (z)$ , so applying key termwise turns it into a double sum over fibres. Again Finset.sum_congr rfl does the termwise part.
```
_ = D.expect (fun w => g (f w)) := Finset.sum_fiberwise univ f _
```
Second link, and the same partition lemma that made map a design in the first place: Finset.sum_fiberwise collapses “over the values of $f$ , then over the fibre” into “over the whole domain once”.

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 868 lines of library preamble — the code above is at the bottom of the editor.

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 $B_{k}$ for block $k$ and $N_{k} = | B_{k} |$ for the block sizes; since the blocks partition the units, $\sum_{k = 1}^{K} N_{k} = 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 944 lines of library preamble — the code above is at the bottom of the editor.

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”.

Lemma (blockwise counting). For each $k = 1, \dots, K$ let $T_{k}$ be a family of subsets of the population, every member of which is contained in $B_{k}$ . Then

| {S \subseteq {1, \dots, n} : S \cap B_{k} \in T_{k} for every k} | = \prod_{k = 1}^{K} | T_{k} | .

Proof.

The map $S \mapsto (S \cap B_{1}, \dots, S \cap B_{K})$ sends the left-hand set to $T_{1} \times \dots \times T_{K}$ , and $(T_{1}, \dots, T_{K}) \mapsto T_{1} \cup \dots \cup T_{K}$ sends the product back. The two maps are mutually inverse. Cutting a set down to each block and reassembling the pieces returns the set, because the blocks cover the population; and reassembling a tuple and cutting back down to $B_{j}$ returns $T_{j}$ , because each $T_{k}$ lies inside $B_{k}$ and distinct blocks are disjoint. A bijection between finite sets equates their cardinalities, and the product set has $\prod_{k} | T_{k} |$ elements.∎

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

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 the hypothesis ht, that each of their members sits inside its own block.

```
classical
```
Makes every proposition decidable, so that no instance search behind a filter or a biUnion can fail. Defensive rather than necessary: the predicate ∀ k, S ∩ block b k ∈ t k is decidable anyway and the proof goes through without the line. It changes nothing mathematically either way.
```
rw [← Fintype.card_piFinset t]
```
Fintype.card_piFinset t states $| T_{1} \times \dots \times T_{K} | = \prod_{k} | T_{k} |$ , where Fintype.piFinset t is the Finset of dependent tuples $(T_{k})_{k}$ with $T_{k} \in T_{k}$ . Used right-to-left it replaces the product on the right by that single cardinality, so the goal becomes an equality of two counts — the shape a bijection can close.
Goal
```
b : Fin n → Fin K
t : Fin K → Finset (Assignment n)
ht : ∀ (k : Fin K), ∀ T ∈ t k, T ⊆ block b k
⊢ #{S | ∀ (k : Fin K), S ∩ block b k ∈ t k} = #(Fintype.piFinset t)
```

refine Finset.card_nbij' (fun S k => S ∩ block b k) (fun g => univ.biUnion g) ?_ ?_ ?_ ?_

Finset.card_nbij' equates two cardinalities from a pair of mutually inverse maps. The forward map is $S \mapsto (S \cap B_{k})_{k}$ and the backward map is $g \mapsto ⋃_{k} g (k)$ ; the four ?_ holes are its four obligations, in order: each map lands in the right set, and each composite is the identity on its own side.

Goal

case refine_1
⊢ Set.MapsTo (fun S k => S ∩ block b k) ↑{S | ∀ (k : Fin K), S ∩ block b k ∈ t k}
    ↑(Fintype.piFinset t)
case refine_2
⊢ Set.MapsTo (fun g => univ.biUnion g) ↑(Fintype.piFinset t)
    ↑{S | ∀ (k : Fin K), S ∩ block b k ∈ t k}
case refine_3
⊢ Set.LeftInvOn (fun g => univ.biUnion g) (fun S k => S ∩ block b k)
    ↑{S | ∀ (k : Fin K), S ∩ block b k ∈ t k}
case refine_4
⊢ Set.RightInvOn (fun g => univ.biUnion g) (fun S k => S ∩ block b k)
    ↑(Fintype.piFinset t)

```
intro S hS; simp only [Finset.coe_filter, Set.mem_ofPred_eq, Finset.mem_univ, true_and] at hS
```
First obligation, forward map lands in the tuple set. The hypothesis hS says $S$ is in the coerced filter; Finset.coe_filter turns that coercion into a Set comprehension, Set.mem_ofPred_eq unfolds membership in it to the predicate, and Finset.mem_univ/true_and discard the vacuous S ∈ univ conjunct. What is left is exactly ∀ k, S ∩ block b k ∈ t k.
```
simpa only [Finset.mem_coe, Fintype.mem_piFinset] using hS
```
And that is the goal, once membership in the tuple set is unfolded the same way: Finset.mem_coe strips the coercion and Fintype.mem_piFinset says a tuple lies in piFinset t exactly when each component lies in its t k.
```
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)
```
Second obligation, backward map lands in the filter set. Unfolding hg the same way gives ∀ k, g k ∈ t k, and feeding that to the hypothesis ht gives hsub: every piece of the tuple sits inside its own block. That is the only thing the reassembly argument needs.
```
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
```
Unfold the goal’s membership to ∀ k, (univ.biUnion g) ∩ block b k ∈ t k, then inter_block_biUnion hsub k — the disjointness half of “the blocks partition the population” — rewrites $(⋃_{l} g (l)) \cap B_{k}$ as $g (k)$ , and hg k says $g (k) \in T_{k}$ .
```
exact fun S _ => biUnion_inter_block b S
```
Third obligation, cut then reassemble is the identity: biUnion_inter_block b S states $⋃_{k} (S \cap B_{k}) = S$ , which is the covering half of the same partition fact.
```
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
```
Fourth obligation, reassemble then cut is the identity. The two tuples are equal when they agree at every index, which is what funext reduces the goal to, and at index $k$ that is inter_block_biUnion again.

Specializing to $T_{k} = {T \subseteq B_{k} : | T | = n_{1 k}}$ gives the count we are after.

Corollary. The blocked design’s support — the assignments treating exactly $n_{1 k}$ units in every block — has

\prod_{k = 1}^{K} (\binom{N_{k}}{n_{1 k}})

elements.

Proof.

An assignment $z$ satisfies $| z \cap B_{k} | = n_{1 k}$ for every $k$ exactly when each piece $z \cap B_{k}$ is an $n_{1 k}$ -element subset of $B_{k}$ ; the containment $z \cap B_{k} \subseteq B_{k}$ is automatic. So the support is the set the lemma counts, with $T_{k}$ the family of $n_{1 k}$ -element subsets of $B_{k}$ , and $| T_{k} | = (\binom{N_{k}}{n_{1 k}})$ .∎

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

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 _ _

```
rw [stratSupport_eq_filter, …]
```
stratSupport_eq_filter is the restatement in the shape the lemma expects: it trades the condition #(z ∩ block b k) = n₁ k for z ∩ block b k ∈ (block b k).powersetCard (n₁ k), where Finset.powersetCard m s is the Finset of $m$ -element subsets of s. The two are equivalent because Finset.inter_subset_right supplies the containment for free — that is the one-line iff in its own proof.
Goal
```
b : Fin n → Fin K
n₁ : Fin K → ℕ
⊢ #{S | ∀ (k : Fin K), S ∩ block b k ∈ powersetCard (n₁ k) (block b k)}
    = ∏ k, (#(block b k)).choose (n₁ k)
```
```
rw […, card_blockwise b (fun k => (block b k).powersetCard (n₁ k)) (fun k T hT => (Finset.mem_powersetCard.mp hT).1)]
```
Apply the lemma at $T_{k} =$ the $n_{1 k}$ -subsets of $B_{k}$ . Its hypothesis ht is discharged inline: Finset.mem_powersetCard.mp hT splits “T is an n₁ k-subset of block b k” into the containment and the count, and .1 is the containment. The count is now a product of cardinalities of powerset-card families.
Goal
```
⊢ ∏ k, #(powersetCard (n₁ k) (block b k)) = ∏ k, (#(block b k)).choose (n₁ k)
```
```
exact Finset.prod_congr rfl fun k _ => Finset.card_powersetCard _ _
```
Finset.prod_congr rfl compares the two products factor by factor, and Finset.card_powersetCard is the binomial coefficient itself: $| {T \subseteq s : | T | = m} | = (\binom{| s |}{m})$ .

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

Definition (stratified complete randomization). Fix quotas $n_{1 k} \leq N_{k}$ for $k = 1, \dots, K$ . Stratified complete randomization is the uniform distribution on the set of assignments $z$ with $| z \cap B_{k} | = n_{1 k}$ for every $k$ ; that is,

Pr (Z = z) = {\begin{cases} (\prod_{k} (\binom{N_{k}}{n_{1 k}}))^{- 1} & if | z \cap B_{k} | = n_{1 k} for every k, \\ 0 & otherwise. \end{cases}

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

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. The three fields again:

```
prob z := if (∀ k, #(z ∩ block b k) = n₁ k) then (#(stratSupport b n₁) : ℝ)⁻¹ else 0
```
The pmf: the reciprocal of the number of admissible assignments on the admissible ones, zero elsewhere. Note that the normalizing constant is written as a cardinality, not as $\prod_{k} (\binom{N_{k}}{n_{1 k}})$ ; the product form is the corollary above, and keeping the count abstract here is what lets the definition go through before any combinatorics is done.
```
nonneg _ := by
  split
  · positivity
  · exact le_refl 0
```
Nonnegativity, one branch of the if at a time. split produces both branches; in the first the value is the inverse of a natural-number cast, which positivity handles, and in the second it is literally $0$ .
```
have hne : ((#(stratSupport b n₁) : ℕ) : ℝ) ≠ 0 := Nat.cast_ne_zero.mpr (card_stratSupport_pos h).ne'
```
Summing to one begins with the fact that makes the normalization legal at all. card_stratSupport_pos h is where the hypothesis h is used: the count is a product of $(\binom{N_{k}}{n_{1 k}})$ , each positive by Nat.choose_pos because $n_{1 k} \leq N_{k}$ . .ne' turns $0 < c$ into $c \neq 0$ and Nat.cast_ne_zero.mpr moves that from $N$ to $R$ — the denominator is nonzero, so the reciprocal is a real number and not Lean’s $1 / 0 = 0$ .
Goal
```
hne : ↑(#(stratSupport b n₁)) ≠ 0
⊢ (∑ z, if ∀ (k : Fin K), #(z ∩ block b k) = n₁ k
      then (↑(#(stratSupport b n₁)))⁻¹ else 0) = 1
```
```
= ∑ _z ∈ stratSupport b n₁, (#(stratSupport b n₁) : ℝ)⁻¹ := by rw [← Finset.sum_filter]; rfl
```
First calc link: throw away the zero terms. Finset.sum_filter says that summing $f$ over the members of $s$ that satisfy $p$ is the same as summing if p a then f a else 0 over all of $s$ , so used right-to-left it turns the sum of an if over all assignments into a sum over the filtered set. rfl then identifies that filtered set with stratSupport b n₁, which is what stratSupport is defined to be.
Goal
```
⊢ ∑ a with ∀ (k : Fin K), #(a ∩ block b k) = n₁ k, (↑(#(stratSupport b n₁)))⁻¹
    = ∑ _z ∈ stratSupport b n₁, (↑(#(stratSupport b n₁)))⁻¹
```

_ = ((#(stratSupport b n₁) : ℕ) : ℝ) * (#(stratSupport b n₁) : ℝ)⁻¹ := by rw [Finset.sum_const, nsmul_eq_mul]

Second link: the summand no longer depends on $z$ , so Finset.sum_const replaces the sum by #s • c, the cardinality acting on the constant by repeated addition. nsmul_eq_mul converts that $N$ -action into an honest multiplication by the cast.

Goal

⊢ #(stratSupport b n₁) • (↑(#(stratSupport b n₁)))⁻¹
    = ↑(#(stratSupport b n₁)) * (↑(#(stratSupport b n₁)))⁻¹

```
_ = 1 := mul_inv_cancel₀ hne
```
Third link: $c \cdot c^{- 1} = 1$ for $c \neq 0$ , which is exactly what hne supplies. The design is normalized.

Balance is then exact, not approximate:

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

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 hc
```
by_contra assumes the negation of the goal — here hc, that the count in block $k$ is not $n_{1 k}$ — and asks for a contradiction. The proof is by contradiction because the design’s pmf is an if, and the natural thing to say about an if is which branch it took.
Goal
```
hz : (stratifiedDesign b n₁ h).prob z ≠ 0
k : Fin K
hc : ¬#(z ∩ block b k) = n₁ k
⊢ False
```
```
refine hz ?_
```
hz says the probability of $z$ is nonzero, i.e. it is a function from “the probability is $0$ ” to False. Applying it turns the goal False into the goal (stratifiedDesign b n₁ h).prob z = 0.

rw [stratifiedDesign_prob]

Unfold the pmf with its @[simp] lemma, so the goal is that the if expression evaluates to $0$ — that is, that its condition is false.

Goal

hc : ¬#(z ∩ block b k) = n₁ k
⊢ (if ∀ (k : Fin K), #(z ∩ block b k) = n₁ k
      then (↑(#(stratSupport b n₁)))⁻¹ else 0) = 0

```
exact ite_eq_right (fun hall => hc (hall k))
```
ite_eq_right states that if c then t else e equals e when ¬c, so it wants a refutation of the condition ∀ k, #(z ∩ block b k) = n₁ k. Given such an hall, instantiating it at our block $k$ contradicts hc — one instance is enough to kill the universal.

Under complete randomization $| z \cap B_{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 counting lemma.

Lemma (inclusion counts factorize). Let $S$ be any set of units. The number of admissible assignments that treat all of $S$ is

| {z : | z \cap B_{k} | = n_{1 k} \forall k, S \subseteq z} | = \prod_{k = 1}^{K} | {T \subseteq B_{k} : | T | = n_{1 k}, S \cap B_{k} \subseteq T} | .

Proof.

The constraint $S \subseteq z$ is local after all: $S \subseteq z$ holds exactly when $S \cap B_{k} \subseteq z \cap B_{k}$ for every $k$ , since every unit lies in exactly one block. So the set on the left is again a set of the form the blockwise counting lemma handles, with $T_{k}$ the $n_{1 k}$ -element subsets of $B_{k}$ that contain $S \cap B_{k}$ ; and its hypothesis holds because those subsets lie inside $B_{k}$ by construction.∎

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

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]

The Lean proof is that argument with the set identity done explicitly, because the two descriptions of the same collection of assignments are not syntactically the same Finset.

```
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
```
The hypothesis the counting lemma will need, discharged first. Finset.mem_filter.mp hT splits membership in the filtered family into “an $n_{1 k}$ -subset of $B_{k}$ ” and “contains $S \cap B_{k}$ ”; .1 keeps the former, Finset.mem_powersetCard.mp splits that into containment and count, and the outer .1 is the containment $T \subseteq B_{k}$ .

ext z

Now hset, the set identity. Two Finsets are equal when they have the same members, and ext z reduces the goal to that biconditional for an arbitrary $z$ .

Goal

z : Assignment n
⊢ z ∈ {z ∈ stratSupport b n₁ | S ⊆ z} ↔
    z ∈ {z | ∀ (k : Fin K), z ∩ block b k ∈
      {T ∈ powersetCard (n₁ k) (block b k) | S ∩ block b k ⊆ T}}

```
simp only [Finset.mem_filter, mem_stratSupport, Finset.mem_univ, true_and, Finset.mem_powersetCard, subset_iff_forall_inter_block b S z]
```
Unfold both sides into elementary statements about the blocks. Finset.mem_filter and Finset.mem_univ/true_and strip the filters, mem_stratSupport replaces membership in the support by ∀ k, #(z ∩ block b k) = n₁ k, Finset.mem_powersetCard splits the powerset-card membership into containment and count, and subset_iff_forall_inter_block is the mathematical content: $S \subseteq z$ iff $S \cap B_{k} \subseteq z \cap B_{k}$ for every $k$ . What is left is a purely logical rearrangement — a pair of $\forall$ s on the left, one $\forall$ over a conjunction on the right.
Goal
```
⊢ ((∀ (k : Fin K), #(z ∩ block b k) = n₁ k)
      ∧ ∀ (k : Fin K), S ∩ block b k ⊆ z ∩ block b k) ↔
    ∀ (k : Fin K), (z ∩ block b k ⊆ block b k ∧ #(z ∩ block b k) = n₁ k)
      ∧ S ∩ block b k ⊆ z ∩ block b k
```
```
constructor
· rintro ⟨hc, hS⟩ k
  exact ⟨⟨Finset.inter_subset_right, hc k⟩, hS k⟩
```
Left to right. rintro ⟨hc, hS⟩ names the two halves, k fixes a block, and the tuple is reassembled at that block; the one piece not in the hypotheses is $z \cap B_{k} \subseteq B_{k}$ , which is Finset.inter_subset_right.
```
· exact fun hall => ⟨fun k => (hall k).1.2, fun k => (hall k).2⟩
```
Right to left, pure projection: from the conjunction at each block take the count (.1.2) to rebuild the support condition and the containment (.2) to rebuild $S \subseteq z$ . The stray containment $z \cap B_{k} \subseteq B_{k}$ is simply dropped.
```
rw [hset, card_blockwise b _ hsub]
```
Rewrite with the identity just proved and apply the counting lemma at the filtered families, with hsub as its hypothesis and the family itself left to unification. That closes the goal.

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,102 lines of library preamble — the code above is at the bottom of the editor.

theorem stratified_propensity (i : Fin n) :
    (stratifiedDesign b n₁ h).propensity i
      = (n₁ (b i) : ℝ) / (#(block b (b i)) : ℝ) := by
  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]
  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]
  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))

Proof.

The design is uniform on its support, so for any unit $i$

π_{i} = Pr (i \in z) = \frac{| {z \in supp : i \in z} |}{| supp |},

and both counts factorize over blocks by the previous lemma, taken at $S = {i}$ . In the numerator, a block $k \neq b (i)$ meets ${i}$ in the empty set, so its factor is unconstrained and equals $(\binom{N_{k}}{n_{1 k}})$ — the same factor it contributes to the denominator. Only $i$ ‘s own block differs: there the factor counts the $n_{1, b (i)}$ -element subsets of $B_{b (i)}$ that contain $i$ , of which there are $(\binom{N_{b (i)} - 1}{n_{1, b (i)} - 1})$ . Cancelling the other $K - 1$ factors,

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

the last step being the identity $(\binom{N - 1}{m - 1}) / (\binom{N}{m}) = m / N$ for $1 \leq m \leq N$ , which is the same one that gave the propensity of complete randomization in chapter 7. If $n_{1, b (i)} = 0$ no admissible subset of $B_{b (i)}$ contains $i$ , so the numerator is $0$ and both sides of the claim are $0$ .∎

The Lean is that computation with the cancellation and the degenerate case written out.

```
have hcount : … = (#((stratSupport b n₁).filter (fun z => i ∈ z)) : ℝ) * ((#(stratSupport b n₁) : ℕ) : ℝ)⁻¹ := by rw [Design.propensity_eq_inclusion, stratified_inclusion, ← stratSupport_filter_mem_eq]
```
The first of three haves: the propensity is a ratio of counts. Design.propensity_eq_inclusion says $π_{i}$ is the inclusion probability of the singleton ${i}$ ; stratified_inclusion evaluates any inclusion probability of this design as (count of admissible assignments containing $S$ ) $\times | supp |^{- 1}$ , which is what “uniform on its support” means quantitatively; and stratSupport_filter_mem_eq, used right-to-left, trades {i} ⊆ z back for the more readable i ∈ z.
Goal
```
⊢ (stratifiedDesign b n₁ h).inclusion {i}
    = ↑(#({z ∈ stratSupport b n₁ | i ∈ z})) * (↑(#(stratSupport b n₁)))⁻¹
```
```
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'
```
The second: the product over the other blocks is nonzero. Finset.prod_pos reduces positivity of a product to positivity of each factor, and Nat.choose_pos (h k) supplies it because $n_{1 k} \leq N_{k}$ , which is what licenses the cancellation two steps below.
```
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
```
The third: the denominator, split into $i$ ‘s own block times the rest. card_stratSupport is the corollary proved above, and Finset.mul_prod_erase states $f (a) \cdot \prod_{x \in s ∖ {a}} f (x) = \prod_{x \in s} f (x)$ for $a \in s$ ; reversed, it peels the factor belonging to $b (i)$ off the front of the product. card_stratSupport_filter_mem, quoted in the next step, does the same peeling for the numerator.

rw [hcount, ← div_eq_mul_inv, card_stratSupport_filter_mem b n₁ i, hcprod, Nat.cast_mul, Nat.cast_mul, …]

Assemble. hcount puts the propensity in count form and div_eq_mul_inv, reversed, rewrites $a \cdot c^{- 1}$ as $a / c$ ; card_stratSupport_filter_mem factorizes the numerator into ( $i$ ‘s block’s count) $\times$ (the other blocks’ $(\binom{N_{k}}{n_{1 k}})$ ), which is where the previous lemma is actually consumed; hcprod does the denominator; and the two Nat.cast_muls push the $N \to R$ cast inside both products so that the common factor is literally the same real number top and bottom.

Goal

⊢ ↑(#({T ∈ powersetCard (n₁ (b i)) (block b (b i)) | i ∈ T}))
      * ↑(∏ k ∈ univ.erase (b i), (#(block b k)).choose (n₁ k))
    / (↑((#(block b (b i))).choose (n₁ (b i)))
      * ↑(∏ k ∈ univ.erase (b i), (#(block b k)).choose (n₁ k)))
    = ↑(n₁ (b i)) / ↑(#(block b (b i)))

rw […, mul_div_mul_right _ _ hQ]

The cancellation: mul_div_mul_right states $a c / (b c) = a / b$ for $c \neq 0$ , and hQ is the $c \neq 0$ . Every block but $i$ ‘s has now disappeared, and what is left is a within-block ratio of binomial coefficients.

Goal

⊢ ↑(#({T ∈ powersetCard (n₁ (b i)) (block b (b i)) | i ∈ T}))
      / ↑((#(block b (b i))).choose (n₁ (b i)))
    = ↑(n₁ (b i)) / ↑(#(block b (b i)))

```
rcases Nat.eq_zero_or_pos (n₁ (b i)) with hzero | hpos
```
The case split the informal proof also has to make: choose_ratio_one needs $1 \leq n_{1, b (i)}$ , so the empty quota is handled separately. The theorem itself carries no lower bound on the quota, which is the point of doing this here rather than assuming it away.
```
· rw [hzero, card_filter_mem_powersetCard_zero]; simp only [Nat.cast_zero, zero_div]
```
The degenerate branch. With the quota $0$ the only admissible within-block set is $\emptyset$ , which does not contain $i$ , so the numerator is $0$ — that is card_filter_mem_powersetCard_zero. Both sides are then $0$ divided by something, which Nat.cast_zero and zero_div finish.

· rw [card_filter_mem_powersetCard (mem_block_self b i) hpos]

The real branch. card_filter_mem_powersetCard counts the $m$ -subsets of a set containing a fixed element of it: there are $(\binom{| s | - 1}{m - 1})$ of them. Its two hypotheses are mem_block_self b i (a unit lies in its own block) and hpos.

Goal

hpos : n₁ (b i) > 0
⊢ ↑((#(block b (b i)) - 1).choose (n₁ (b i) - 1)) / ↑((#(block b (b i))).choose (n₁ (b i)))
    = ↑(n₁ (b i)) / ↑(#(block b (b i)))

```
exact_mod_cast choose_ratio_one hpos (h (b i))
```
And that ratio is the quota over the block size, by choose_ratio_one from chapter 7 — the same lemma that gave complete randomization its propensity. exact_mod_cast closes a goal that matches the lemma up to ℕ → ℝ coercions, which is most of the friction in this library.

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 $\hat{DiM}$ is a fixed linear combination of the indicators and its expectation is available in closed form.

Proposition. Under stratified complete randomization with quotas $n_{1 k}$ and $n_{1} = \sum_{k} n_{1 k}$ ,

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

Proof.

On every assignment the design can produce, $| z | = n_{1}$ and hence $| z^{c} | = n - n_{1}$ : the quotas are met in each block, and the block counts add up. So on the support

\hat{DiM} = \frac{\sum_{i \in z} y_{i} (1)}{| z |} - \frac{\sum_{i \notin z} y_{i} (0)}{| z^{c} |} = \frac{\sum_{i} Z_{i} y_{i} (1)}{n_{1}} - \frac{\sum_{i} (1 - Z_{i}) y_{i} (0)}{n - n_{1}},

where the second equality replaces each restricted sum by an indicator-weighted sum over the whole population. Two functions that agree wherever the design puts positive probability have the same expectation, so we may take the expectation of the right-hand side instead. It is a fixed linear combination of the indicators, and by linearity of expectation with $E [Z_{i}] = π_{i}$ the claim follows.∎

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

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]

If $π_{i}$ is the constant $n_{1} / n$ the weights collapse and this is $\bar{y} (1) - \bar{y} (0) = τ$ . Otherwise it is a pair of weighted averages with the wrong weights.

```
have hle : ∑ k, n₁ k ≤ n := sum_n₁_le h
```
The quotas fit inside the population — needed only so that the cast $n - n_{1}$ from $N$ to $R$ can be taken without a truncated subtraction. It is the first warm-up exercise below.
```
intro z hz; have hcard : #z = ∑ k, n₁ k := stratified_card_eq_of_prob_ne_zero hz
```
Now hsup, the “on the support” rewriting. Fix an assignment $z$ with nonzero probability; stratified_card_eq_of_prob_ne_zero (a core exercise below) says the overall number treated is the sum of the quotas — the fixed total that the first half of chapter 8’s argument needed.
Goal
```
hcard : #z = ∑ k, n₁ k
⊢ 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))
```
```
have hcompl : ((#zᶜ : ℕ) : ℝ) = (n : ℝ) - ((∑ k, n₁ k : ℕ) : ℝ) := by rw [card_compl, hcard, Nat.cast_sub hle]
```
And so is the control count. card_compl gives $| z^{c} | = n - | z |$ in $N$ , hcard substitutes the total, and Nat.cast_sub hle is what lets the truncated $N$ -subtraction be pushed through the cast as an honest real subtraction; that is the only work hle does.
```
rw [DiM, sum_mem_eq_sum_Z_mul, sum_compl_eq_sum_one_sub_Z_mul, hcard, hcompl]
```
Unfold the estimator and rewrite it into indicator form: sum_mem_eq_sum_Z_mul states $\sum_{i \in z} v_{i} = \sum_{i} Z_{i} v_{i}$ and sum_compl_eq_sum_one_sub_Z_mul its control-side twin $\sum_{i \notin z} v_{i} = \sum_{i} (1 - Z_{i}) v_{i}$ . Then hcard and hcompl replace the two random denominators by constants. This is the displayed equation in the informal proof.

rw [(stratifiedDesign b n₁ h).expect_congr_of_support hsup]

expect_congr_of_support is the formal version of “two functions that agree on the support have the same expectation”, so the goal is now about the indicator-form function, whose denominators are constants.

Goal

⊢ ((stratifiedDesign b n₁ h).expect fun z =>
      (∑ i, Z z i * P.y1 i) / ↑(∑ k, n₁ k)
        - (∑ i, (1 - Z z i) * P.y0 i) / (↑n - ↑(∑ k, n₁ k)))
    = (∑ 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))

```
simp only [Design.expect_sub, Design.expect_div_const, Design.expect_sum, Design.expect_mul_const, Design.expect_Z, Design.expect_const]
```
Linearity, applied outside in: expect_sub splits the difference of the two terms, expect_div_const pulls each constant denominator out of the expectation, expect_sum exchanges the expectation with the sum over units, and expect_mul_const pulls each $y_{i} (1)$ or $y_{i} (0)$ out of its term. What is left inside is $E [Z_{i}]$ on the treated side, which expect_Z replaces by $π_{i}$ — that is what Design.propensity is defined to be — and $E [1 - Z_{i}]$ on the control side, which expect_sub splits a second time so that expect_const gives the $1$ and expect_Z the $π_{i}$ . Note that 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,348 lines of library preamble — the code above is at the bottom of the editor.

def bEx : Fin 5 → Fin 2 := ![0, 0, 1, 1, 1]

def n₁Ex : Fin 2 → ℕ := ![1, 2]

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

def PEx : Population 5 := ⟨fun _ => 0, fun i => if bEx i = 0 then 1 else 0⟩

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. Both potential outcomes are known here because the population is a thought experiment: $y_{i} (1) = 0$ for everyone, and $y_{i} (0) = 1$ in the first block and $0$ in the second, so $τ = \frac{1}{5} (0 - 2) = - 2 / 5$ .

The proof is a chain of small evaluations rather than an argument; it is worth walking once because every numeric example in this book has this shape.

```
have hprop : ∀ i : Fin 5, (stratifiedDesign bEx n₁Ex hEx).propensity i = (n₁Ex (bEx i) : ℝ) / (#(block bEx (bEx i)) : ℝ) := stratified_propensity
```
Instantiate the propensity theorem at this design. Nothing here is specific to the example; it is the general theorem with the concrete bEx and n₁Ex substituted in.
```
have hN : (∑ k, n₁Ex k) = 3 := by decide
```
The total quota. decide proves a decidable arithmetic fact about Fin and Finset by evaluating it in the kernel, which is exactly what a claim about five explicit units is.
```
have hc : ∀ i : Fin 5, #(block bEx (bEx i)) = if bEx i = 0 then 2 else 3 := by decide
```
The block sizes, as a function of the unit rather than of the block index — two units in block $0$ , three in block $1$ . Writing it as an if on bEx i is what lets the next few steps stay uniform in $i$ .

have hm : ∀ i : Fin 5, n₁Ex (bEx i) = if bEx i = 0 then 1 else 2 := by decide

Likewise the quotas: one treated unit in block $0$ , two in block $1$ .

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]

Combine the three: each unit’s propensity is $1 / 2$ in the first block and $2 / 3$ in the second. This is the disproportionality that will produce the bias, now in a form the arithmetic can chew on.

```
have b0 : bEx 0 = 0 := by decide
```
Five decides of this shape record which block each unit is in. They are needed because the if bEx i = 0 conditions above have to be resolved unit by unit before norm_num can see numbers.
```
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]
```
The one substantive line. expect_DiM_stratified replaces the left-hand expectation by the closed form just proved, hN puts $3$ in for $n_{1}$ , Population.tau_eq_mean_sub_mean rewrites $τ$ as $\bar{y} (1) - \bar{y} (0)$ , the two means unfold to sums over $/ n$ , the four Fin.sum_univ_fives expand the four sums over Fin 5 into five explicit terms each, and the five hps substitute the propensities.
```
simp only [PEx, b0, b1, b2, b3, b4]
```
Substitute the population’s two outcome vectors and resolve every if bEx i = 0 using the five block facts. After this the goal is a statement about explicit rational numbers.
```
norm_num
```
And norm_num does that arithmetic: $- 1 / 2 \neq - 2 / 5$ .

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,155 lines of library preamble — the code above is at the bottom of the editor.

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) : ℝ)

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_{i} (1)}{| z \cap B_{k} |} - \frac{\sum_{i \in z^{c} \cap B_{k}} y_{i} (0)}{| 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,201 lines of library preamble — the code above is at the bottom of the editor.

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))

Proof.

On the support of the design the blocked estimator equals the Horvitz–Thompson estimator, by the identity just described. Two estimators that agree wherever the design puts positive probability have the same expectation. And Horvitz–Thompson is unbiased for $τ$ whenever $0 < π_{i} < 1$ for every unit, which here reads $0 < n_{1, b (i)} / N_{b (i)} < 1$ : the left inequality is $n_{1 k} \geq 1$ and the right is $n_{1 k} < N_{k}$ , which are exactly the hypotheses that every block has a treated and a control unit.∎

Nothing about blocking is re-derived; all that changed was the propensity score.

```
have hprop : ∀ i : Fin n, (stratifiedDesign b n₁ h).propensity i = (n₁ (b i) : ℝ) / (#(block b (b i)) : ℝ) := stratified_propensity
```
The propensity theorem, named so that the two positivity goals at the end can rewrite with it. This have restates a fact already proved; it buys readability, not content.
```
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)
```
The “on the support” fact, assembled from two pieces: the stretch exercise blockedDiM_eq_HT_of_mem_support needs to know that every block hits its quota on $z$ , and stratified_card_inter_block — the exact-balance theorem above — extracts that from the mere assumption that $z$ has nonzero probability.

rw [(stratifiedDesign b n₁ h).expect_congr_of_support hsupport]

Swap the estimator for Horvitz–Thompson under the expectation. The goal is now the general Horvitz–Thompson theorem’s conclusion verbatim.

Goal

hsupport : ∀ (z : Assignment n), (stratifiedDesign b n₁ h).prob z ≠ 0 →
  blockedDiM b P z = HT (stratifiedDesign b n₁ h) P z
⊢ (stratifiedDesign b n₁ h).expect (HT (stratifiedDesign b n₁ h) P) = P.tau

```
refine HT_unbiased _ P (fun i => ?_) (fun i => ?_)
```
Apply HT_unbiased from chapter 8 and leave its two hypotheses open, one per ?_: the propensity is positive, and it is below one. Both are now goals about a single unit $i$ .
Goal
```
i : Fin n
⊢ 0 < (stratifiedDesign b n₁ h).propensity i
```
```
rw [hprop i]
```
Replace the propensity by the ratio. In the first goal this leaves $0 < n_{1, b (i)} / N_{b (i)}$ ; in the second, after the same rewrite, $n_{1, b (i)} / N_{b (i)} < 1$ .
Goal
```
i : Fin n
⊢ 0 < ↑(n₁ (b i)) / ↑(#(block b (b i)))
```
```
exact div_pos (by exact_mod_cast hpos (b i)) (by exact_mod_cast block_card_pos b i)
```
Positivity: div_pos wants both numerator and denominator positive. hpos (b i) gives $0 < n_{1, b (i)}$ and block_card_pos b i gives $0 < N_{b (i)}$ because a unit lies in its own block; exact_mod_cast moves each from $N$ to $R$ .
```
exact (div_lt_one (by exact_mod_cast block_card_pos b i)).mpr (by exact_mod_cast hlt (b i))
```
And below one: div_lt_one, given a positive denominator, says $a / b < 1$ iff $a < b$ , and hlt (b i) is $n_{1, b (i)} < N_{b (i)}$ . This is where the “a control in every block” hypothesis finally earns its place — it is also handed on to the stretch exercise in step 2 — and the callout below is about what it costs.

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,246 lines of library preamble — the code above is at the bottom of the editor.

def clusterAssignment (c : Fin n → Fin K) (T : Assignment K) : Assignment n :=
  univ.filter (fun i => c i ∈ T)

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, $Z_{i} = Z_{c (i)}^{cluster}$ .

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

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
```
Replace both indicators by their definition, if i ∈ z then 1 else 0. The goal is now an equality of two ifs, whose conditions are “unit $i$ is treated” and “cluster $c (i)$ is treated”.
Goal
```
c : Fin n → Fin K
T : Assignment K
i : Fin n
⊢ (if i ∈ clusterAssignment c T then 1 else 0) = if c i ∈ T then 1 else 0
```
```
by_cases hi : c i ∈ T <;> simp [hi]
```
Split on whether $i$ ‘s cluster is treated. In each branch simp [hi] resolves both conditions at once, because mem_clusterAssignment — a @[simp] lemma — says $i \in$ clusterAssignment c T iff $c (i) \in T$ , which is the same condition. Both branches are then $1 = 1$ or $0 = 0$ .

From it, the defining property of cluster randomization: two units of the same cluster always get the same treatment.

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

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]

```
obtain ⟨T, hT, -⟩ := Design.exists_of_map_prob_ne_zero hz
```
Design.exists_of_map_prob_ne_zero says an assignment in the support of a pushforward comes from some upstream assignment: if the fibre sum is nonzero then some term of it is. obtain names that cluster-level assignment T and the equation hT : clusterAssignment c T = z; the third component, that T itself has nonzero probability, is discarded with - because we do not need it.
Goal
```
hij : c i = c j
T : Assignment K
hT : clusterAssignment c T = z
⊢ Z z i = Z z j
```
```
subst hT
```
Replace z by clusterAssignment c T everywhere. This is the step that makes the argument possible: the hypothesis only said $z$ was in the support, and now $z$ is displayed as coming from a cluster-level assignment.
Goal
```
hij : c i = c j
⊢ Z (clusterAssignment c T) i = Z (clusterAssignment c T) j
```
```
rw [Z_clusterAssignment, Z_clusterAssignment, hij]
```
Both indicators are now cluster-level by the previous lemma, applied once on each side, and hij : c i = c j makes the two cluster indices the same. The goal closes as a reflexivity.

Proposition. Under cluster randomization, a unit’s propensity is its cluster’s: $π_{i} = π_{c (i)}^{cluster}$ .

Proof.

By the change-of-variables formula, the expectation of any statistic of the unit-level assignment is the expectation of the same statistic composed with “treated clusters $\mapsto$ treated units”. Apply it to the statistic $Z_{i}$ : $π_{i} = E [Z_{i}]$ becomes an expectation over cluster-level assignments, and $Z_{i} = Z_{c (i)}^{cluster}$ pointwise, so the expectation is $E [Z_{c (i)}^{cluster}] = π_{c (i)}^{cluster}$ .∎

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

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

```
rw [clusterDesign, …]
```
Unfold clusterDesign to what it is: the pushforward of the cluster-level design along clusterAssignment c. Nothing about clusters is left in the goal except that map.
Goal
```
⊢ (Design.map (clusterAssignment c) Dc).propensity i = Dc.propensity (c i)
```
```
rw […, Design.map_propensity]
```
Design.map_propensity is the change-of-variables formula specialized to the statistic $Z_{i}$ : the propensity of a pushforward is an expectation computed upstream — the average, over cluster-level assignments $T$ , of the indicator that unit $i$ is treated when the treated clusters are exactly $T$ .
Goal
```
⊢ (Dc.expect fun w => Z (clusterAssignment c w) i) = Dc.propensity (c i)
```
```
exact Dc.expect_congr fun T => Z_clusterAssignment c T i
```
Dc.expect_congr says two statistics that agree at every assignment have the same expectation, so it is enough to compare integrands: at each $T$ , the previous lemma turns $Z_{i}$ into $Z_{c (i)}^{cluster}$ , and the expectation of that is the cluster’s propensity by definition.

There is no dilution and no inflation. 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,286 lines of library preamble — the code above is at the bottom of the editor.

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

```
have hKR : (0 : ℝ) < (K : ℝ) := by exact_mod_cast Nat.lt_of_lt_of_le hpos hle
```
Three real-number facts first, because div_pos and div_lt_one below want inequalities in $R$ and the hypotheses are in $N$ . This one is $0 < K$ , chained from $0 < K_{1}$ and $K_{1} \leq K$ .

have hK1R : (0 : ℝ) < (K₁ : ℝ) := by exact_mod_cast hpos

$0 < K_{1}$ : at least one cluster is treated.

have hltR : (K₁ : ℝ) < (K : ℝ) := by exact_mod_cast hlt

And $K_{1} < K$ : at least one cluster is not. These are the cluster-level analogues of “both groups non-empty”, stated about clusters rather than units: a first hint that it is the number of clusters that matters here, not the number of units.

Goal

hKR : 0 < ↑K
hK1R : 0 < ↑K₁
hltR : ↑K₁ < ↑K
⊢ (clusterDesign c (completeRandomization K K₁ hle)).expect
    (HT (clusterDesign c (completeRandomization K K₁ hle)) P) = P.tau

```
have hprop : ∀ i : Fin n, (clusterDesign c (completeRandomization K K₁ hle)).propensity i = (K₁ : ℝ) / (K : ℝ) := fun i => by rw [cluster_propensity, completeRandomization_propensity]
```
The whole content of the theorem, in one line: cluster_propensity reduces a unit’s propensity to its cluster’s, and completeRandomization_propensity (chapter 7) evaluates that as $K_{1} / K$ . Every unit has the same propensity, whatever the cluster sizes are.
```
refine HT_unbiased _ P (fun i => ?_) (fun i => ?_)
```
Apply the general theorem and leave its two side conditions open.
```
· rw [hprop i]; exact div_pos hK1R hKR
```
$0 < K_{1} / K$ from the two positivity facts.

· rw [hprop i]; exact (div_lt_one hKR).mpr hltR

And $K_{1} / K < 1$ from $K_{1} < K$ with $K > 0$ .

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 — one 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: the block quotas add up to at most the population size.

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

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

Hint

A calc with two links: Finset.sum_le_sum (fed fun k _ => h k, the hypothesis at each block) bounds $\sum_{k} n_{1 k}$ by $\sum_{k} N_{k}$ , and sum_card_block says that sum is $n$ .

Show solution

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

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.

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

theorem card_stratSupport_example : #(stratSupport bEx n₁Ex) = 6 := by
  sorry

Hint

rw [card_stratSupport] turns the count into a product of binomial coefficients; decide evaluates it in the kernel.

Show solution

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

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.

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

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

Hint

Rewrite with stratified_propensity i, then div_eq_div_iff — whose two side conditions are the positivity of $N_{b (i)}$ and of $n$ , from block_card_pos and Fin.pos i cast to ℝ, since a unit lives in a nonempty block and in a nonempty population — leaves the cross-multiplied form, which is hprop (b i) up to ℕ → ℝ coercions; exact_mod_cast closes it.

Show solution

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

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.

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

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

Hint

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_propensity (twice, with hij) makes the two propensities agree, cluster_jointPropensity_same replaces $π_{i j}$ by $π_{i}$ , and ring finishes.

Show solution

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

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. Blocking does not break DiM by itself; unequal treated fractions do.

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

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

Hint

Rewrite with expect_DiM_stratified and Population.tau_eq_mean_sub_mean, then simp only with hbal, ← Finset.mul_sum (to pull the constant propensity out of each sum) and mean, and let field_simp clear the denominators — which is why you need $0 < n$ , $0 < n_{1}$ and $0 < n - n_{1}$ as real-number haves first.

Show solution

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

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 the library itself, which poses its own exercises inline.

ExerciseDistinct blocks are disjoint (warm-up)

Distinct blocks are disjoint — a unit has only one value of b i. This is half of “the blocks partition the population”.

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

theorem block_disjoint (b : Fin n → Fin K) {k l : Fin K} (hkl : k ≠ l) :
    Disjoint (block b k) (block b l) := by
  sorry

Hint

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.

Show solution

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

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: the block sizes add up to the population size, $\sum_{k} N_{k} = n$ .

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

theorem sum_card_block (b : Fin n → Fin K) : ∑ k, #(block b k) = n := by
  sorry

Hint

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 f := b with s = t = univ, then simpa only with Finset.card_univ, Fintype.card_fin and block lines the two sides up (the lemma states the equation the other way round, so use .symm).

Show solution

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

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)

Under a blocked design the Horvitz–Thompson estimator is unbiased too. It is the same statement as blockedDiM_unbiased in disguise, since the two estimators agree on the support.

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

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

Hint

The last two lines of blockedDiM_unbiased, run on their own. refine HT_unbiased _ P (fun i => ?_) (fun i => ?_), then in each goal rw [stratified_propensity i]: $0 < n_{1 k} / N_{k}$ is div_pos, and $n_{1 k} / N_{k} < 1$ is div_lt_one, with the $N$ -to- $R$ side conditions (hpos, hlt, block_card_pos) discharged by exact_mod_cast.

Show solution

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

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.

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

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

Hint

First establish ∑ k, #(z ∩ block b k) = #z as a have: it is Finset.sum_fiberwise applied to the indicator fun i => if i ∈ z then 1 else 0, cleaned up by simpa only with Finset.sum_ite_mem, Finset.card_eq_sum_ones, block, Finset.univ_inter and Finset.inter_comm. Rewrite the goal with it backwards, and Finset.sum_congr rfl with stratified_card_inter_block hz k handles each block.

Show solution

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

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)

The joint propensity of two units is the joint propensity of their two clusters: nothing is lost or gained by passing to the unit level.

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

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

Hint

Unfold Design.jointPropensity and clusterDesign, push the expectation through the pushforward with Design.expect_map, unfold Design.jointPropensity again on the cluster side, and finish with Dc.expect_congr and two rewrites by Z_clusterAssignment.

Show solution

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

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. This and the previous exercise are the pair that make the intra-cluster dependence precise.

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

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

Hint

The same rewrite chain as the previous exercise — cluster_propensity, Design.jointPropensity, clusterDesign, Design.expect_map, then Design.propensity on the cluster side — with one extra step inside Dc.expect_congr: after two Z_clusterAssignment rewrites and hij turn both indicators into the same one, Z_mul_self says an indicator times itself is itself.

Show solution

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

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: on the support of the blocked design, the blocked estimator is the Horvitz–Thompson estimator. It is the blocked analogue of DiM_eq_HT_of_card, and the lemma blockedDiM_unbiased is a corollary of.

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

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

Hint

Work one summand at a time first: for i ∈ block b k, stratified_propensity i and mem_block turn the Horvitz–Thompson summand into $N_{k}$ times Z z i * y1 i / n₁ k - (1 - Z z i) * y0 i / (N_k - n₁ k), with div_div_eq_mul_div, one_sub_div and field_simp doing the algebra. Summing that over a block (Finset.mul_sum, Finset.sum_sub_distrib, Finset.sum_div, sum_inter_eq_sum_Z_mul, sum_compl_inter_eq_sum_one_sub_Z_mul) gives N_k * blockDiM b P z k, once hz k and card_compl_inter with Nat.cast_sub identify the two denominators. Then unfold HT blockedDiM, regroup the Horvitz–Thompson sum with ← sum_eq_sum_block b, and ring closes each block’s term; the nonvanishing of $n_{1 k}$ , $N_{k}$ and $N_{k} - n_{1 k}$ has to be on hand as haves throughout.

Show solution

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

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)]
  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

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

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.