Designs

In the design-based model, the population and potential outcomes are fixed. The treatment assignment is random according to a known design: a probability mass function on the $2^n$ possible treated sets.

This chapter defines designs, expectations, marginal and joint propensities, complete randomization, and Bernoulli randomization. Because the assignment space is finite, all probabilities and expectations are finite sums.

A design is its pmf

Informally: a design assigns to each possible treated set $z\subseteq \{1,\dots ,n\}$ a probability $P(z)\ge 0$, and those probabilities sum to one. There is no more to it.

structure Design (n : ℕ) where
  prob : Assignment n → ℝ
  nonneg : ∀ z, 0 ≤ prob z
  sum_one : ∑ z, prob z = 1

The definition illustrates three patterns used throughout the library.

It is a structure, so the axioms travel with the data. A design is not “a function prob, plus two hypotheses I will remember to write down in every theorem”. It is one bundle with three fields, two of which happen to be proofs. You cannot construct a Design without supplying them, and once you have a D : Design n you can use them: D.sum_one is the proof that the probabilities sum to one, ready to be rewritten with.

The proof fields cost nothing at the point of use. nonneg and sum_one live in Prop, and Lean’s proof irrelevance means any two proofs of the same proposition are interchangeable. So a design carries exactly one piece of information, its pmf: two designs with the same prob are the same design.¹That statement is itself a theorem, Design.ext, and it lives in the bridge chapter rather than here because nothing before it needs to compare two designs for equality.

There is no sample space beyond the treated sets. Assignment n is an abbreviation for Finset (Fin n), the type of finite sets of units, which Lean knows is a finite type with $2^n$ elements. So ∑ z, D.prob z in the sum_one field is a sum over all $2^n$ assignments, with no index set written down: ∑ z, f z means ∑ z ∈ Finset.univ, f z, and univ here is the set of all treated sets. A “random variable” is then any function Assignment n → ℝ, and no measurability condition is needed, because on a finite type there is nothing to measure.

Expectation is a weighted sum, and that is all

noncomputable def expect : ℝ := ∑ z, D.prob z * f z

This is the definition a statistician already has in mind: average the statistic over the randomization, weighting by the design. f ranges over functions of the assignment alone, which is exactly right for a design-based theory — an estimator, an indicator, a squared deviation are all functions of $z$ once the population is fixed.

Everything the tutorial needs about expectation follows from three facts about finite sums. Here is linearity, in full.

Proposition (linearity of expectation). For every design and every pair of statistics $f,g$, $\mathbb{E}[f+g]=\mathbb{E}[f]+\mathbb{E}[g]$.

theorem expect_add : D.expect (fun z => f z + g z) = D.expect f + D.expect g := by
  simp only [expect, mul_add, Finset.sum_add_distrib]

The proof is one tactic with three rewrite rules, and it is worth walking through slowly, because simp only is the workhorse of the whole library. simp only [h₁, h₂, …] rewrites the goal left-to-right with the given lemmas, repeatedly, and with nothing else — unlike a bare simp, which is free to use several thousand lemmas you did not name and can therefore break when Mathlib changes.

The goal starts as

D.expect (fun z => f z + g z) = D.expect f + D.expect g

Naming a definition in a simp only list unfolds it (Lean generates an equation lemma expect D f = ∑ z, D.prob z * f z when the definition is made), so after expect the goal is

∑ z, D.prob z * (f z + g z) = (∑ z, D.prob z * f z) + ∑ z, D.prob z * g z

mul_add : a * (b + c) = a * b + a * c then rewrites the summand. Note that it rewrites under the binder: the body of ∑ z, … is a function of z, and simp descends into it using congruence lemmas. The goal becomes

∑ z, (D.prob z * f z + D.prob z * g z) = (∑ z, D.prob z * f z) + ∑ z, D.prob z * g z

and Finset.sum_add_distrib : ∑ i ∈ s, (f i + g i) = ∑ i ∈ s, f i + ∑ i ∈ s, g i turns the left-hand side into the right-hand side. The two sides are now syntactically equal, and simp only closes any goal that becomes rfl.

That is the entire content of “expectation is linear” in a finite population: distributivity, plus the fact that a finite sum of sums is a sum of finite sums. No integrability, no dominated convergence, no hypotheses at all — expect_add holds for every design and every pair of statistics.

The second axiom of a design first earns its keep in one small lemma.

Lemma. For every design and every constant $c$, $\mathbb{E}[c]=c$.

theorem expect_const (c : ℝ) : D.expect (fun _ => c) = c := by
  simp only [expect, ← Finset.sum_mul, D.sum_one, one_mul]

Four entries, one per step of that display. expect unfolds the definition, leaving ∑ z, D.prob z * c. ← Finset.sum_mul pulls the constant c out of the sum (the lemma is stated as (∑ i ∈ s, f i) * b = ∑ i ∈ s, f i * b, so the arrow uses it backwards) to give (∑ z, D.prob z) * c. D.sum_one — the structure field, used as a rewrite rule — replaces ∑ z, D.prob z by 1, and one_mul finishes. If you ever wondered what work the “probabilities sum to one” axiom actually does in a proof of unbiasedness, this is it: it makes $\mathbb{E}[c]=c$, and everything else is linearity.

The form of linearity that unbiasedness proofs actually use is the one for a sum over units, since every estimator in this tutorial is a sum over units.

Lemma. For any finite index set and any family of statistics $f_1,\dots ,f_m$, $\mathbb{E}[\sum _i{f}_i]=\sum _i\mathbb{E}[f_i]$.

theorem expect_sum {ι : Type*} (s : Finset ι) (F : ι → Assignment n → ℝ) :
    D.expect (fun z => ∑ i ∈ s, F i z) = ∑ i ∈ s, D.expect (F i) := by
  simp only [expect, Finset.mul_sum]
  exact Finset.sum_comm

1. simp only [expect, Finset.mul_sum]

   expect unfolds both sides, and Finset.mul_sum : b * ∑ i ∈ s, f i = ∑ i ∈ s, b * f i pushes the factor D.prob z inside the sum over units — the “taking $P(z)$ inside” of the informal proof. What is left is one double sum on each side, summed in opposite orders.

   Goal

   ⊢ ∑ x, ∑ i ∈ s, D.prob x * F i x = ∑ i ∈ s, ∑ z, D.prob z * F i z

2. exact Finset.sum_comm

   Finset.sum_comm : ∑ x ∈ s, ∑ y ∈ t, f x y = ∑ y ∈ t, ∑ x ∈ s, f x y is the exchange, and it matches the goal exactly. That exchange is the only “theorem” in the whole of our probability theory, and in a finite population it is free.

theorem expect_const_mul (c : ℝ) : D.expect (fun z => c * f z) = c * D.expect f := by
  sorry

theorem expect_const_mul (c : ℝ) : D.expect (fun z => c * f z) = c * D.expect f := by
  simp only [expect, Finset.mul_sum]
  exact Finset.sum_congr rfl fun z _ => by ring

theorem expect_congr_of_support {f g : Assignment n → ℝ}
    (h : ∀ z, D.prob z ≠ 0 → f z = g z) : D.expect f = D.expect g := by
  sorry

theorem expect_congr_of_support {f g : Assignment n → ℝ}
    (h : ∀ z, D.prob z ≠ 0 → f z = g z) : D.expect f = D.expect g := by
  refine Finset.sum_congr rfl fun z _ => ?_
  by_cases hz : D.prob z = 0
  · rw [hz, zero_mul, zero_mul]
  · rw [h z hz]

Indicators, events, probabilities

The bridge between algebra and probability is that an event’s probability is the expectation of its indicator. Both sides exist in the library:

noncomputable def probOf (p : Assignment n → Prop) [DecidablePred p] : ℝ :=
  ∑ z ∈ univ.filter p, D.prob z

Lemma. For every event $A$ — every set of assignments — $\mathbb{E}[\mathbf{1}\{z\in A\}]=Pr(A)$.

theorem expect_indicator (p : Assignment n → Prop) [DecidablePred p] :
    D.expect (fun z => if p z then 1 else 0) = D.probOf p := by
  simp only [expect, probOf, mul_ite, mul_one, mul_zero]
  exact (Finset.sum_filter _ _).symm

probOf sums the pmf over the assignments satisfying a decidable predicate; expect_indicator says that summing P(z) · 1{p(z)} over everything gives the same number.

1. simp only [expect, probOf, mul_ite, mul_one, mul_zero]

   Both definitions are unfolded, and then three lemmas move the multiplication inside the if: mul_ite : a * (if c then b else d) = if c then a * b else a * d distributes it, and mul_one/mul_zero evaluate the two branches, so the summand D.prob z * (if p z then 1 else 0) becomes if p z then D.prob z else 0. That is the informal “zero off the event, P(z) on it”.

   Goal

   ⊢ (∑ x, if p x then D.prob x else 0) = ∑ z with p z, D.prob z

2. exact (Finset.sum_filter _ _).symm

   Finset.sum_filter : ∑ i ∈ s with p i, f i = ∑ i ∈ s, if p i then f i else 0 states exactly the remaining equality, but the other way round, so .symm flips it. Dropping the zero summands from a finite sum is the one and only content of the lemma.

The [DecidablePred p] in square brackets is an instance argument: univ.filter p has to know which assignments to keep, so Lean finds, by itself, the algorithm that decides p. For predicates like fun z => i ∈ z it always can, and you will never have to supply it.

Propensities

A design is a distribution over $2^n$ sets, which is far too much information to work with directly. Almost everything an estimator cares about is captured by two summaries.

noncomputable def propensity (i : Fin n) : ℝ := D.expect (fun z => Z z i)

noncomputable def jointPropensity (i j : Fin n) : ℝ := D.expect (fun z => Z z i * Z z j)

In a design-based setting the propensity score ${\pi }_i=\mathbb{E}[Z_i]=P(i\text{ treated})$ is not a nuisance function to be estimated from covariates. It is a number the experimenter chose. The joint propensity ${\pi }_{ij}=P(i\text{ and }j\text{ both treated})$ is what distinguishes designs that agree on every ${\pi }_i$, and it is what drives every variance formula in chapter 9.

That these two are the same object at different arities is worth making explicit, and the library does so with the inclusion probability of an arbitrary set of units:

noncomputable def inclusion (S : Finset (Fin n)) : ℝ :=
  D.expect (fun z => ∏ i ∈ S, Z z i)

The inclusion probability $\pi (S)=\mathbb{E}[\prod _{i\in S}{Z}_i]$ is the probability that all of $S$ is treated, written as an expectation of a product of indicators — which is legitimate precisely because Z is real-valued, so a product of indicators is again a number and prod_Z : ∏ i ∈ S, Z z i = if S ⊆ z then 1 else 0 turns it into a counting problem. Specializing gives the two summaries back: the propensity is the inclusion probability of the singleton $\{i\}$, because a product over one index, $\prod _{k\in \{i\}}{Z}_k$, is $Z_i$.

theorem propensity_eq_inclusion (i : Fin n) : D.propensity i = D.inclusion {i} := by
  simp only [propensity, inclusion, Finset.prod_singleton]

1. simp only [propensity, inclusion, Finset.prod_singleton]

   The goal shown is the state after the first two names have unfolded the two definitions: both sides are expectations, of Z z i on the left and of a one-factor product on the right. Finset.prod_singleton : ∏ x ∈ {a}, f x = f a collapses that product — under the binder, since it sits inside D.expect (fun z => …) — and the two sides become syntactically equal, which is what simp only needs to finish.

   Goal

   ⊢ (D.expect fun z => Z z i) = D.expect fun z => ∏ i ∈ {i}, Z z i

The propensity is also a probability in the ordinary sense.

Lemma. ${\pi }_i=Pr(i\in z)$ for every design and every unit $i$.

theorem propensity_eq_probOf (i : Fin n) : D.propensity i = D.probOf (fun z => i ∈ z) := by
  unfold propensity Z
  exact D.expect_indicator (fun z => i ∈ z)

1. unfold propensity Z

   Unfolding propensity exposes the expectation, and unfolding Z turns its argument from an indicator function into a literal if p z then 1 else 0. That syntactic shape is the whole point: expect_indicator is stated about the if, and until Z is out of the way it cannot match.

   Goal

   ⊢ (D.expect fun z => if i ∈ z then 1 else 0) = D.probOf fun z => i ∈ z

2. exact D.expect_indicator (fun z => i ∈ z)

   The previous lemma, applied to the event “unit i is treated”. The predicate has to be given explicitly because Lean cannot guess how to split the goal’s if back into a predicate; everything else, including the DecidablePred instance, is found for it.

@[simp] theorem expect_one_sub_Z (i : Fin n) :
    D.expect (fun z => 1 - Z z i) = 1 - D.propensity i := by
  sorry

@[simp] theorem expect_one_sub_Z (i : Fin n) :
    D.expect (fun z => 1 - Z z i) = 1 - D.propensity i := by
  have hsub := D.expect_sub (fun _ => (1 : ℝ)) (fun z => Z z i)
  simpa only [D.expect_const 1, expect_Z] using hsub

theorem sum_propensity : ∑ i, D.propensity i = D.expect (fun z => (#z : ℝ)) := by
  sorry

theorem sum_propensity : ∑ i, D.propensity i = D.expect (fun z => (#z : ℝ)) := by
  simp only [propensity]
  rw [← D.expect_sum Finset.univ (fun i z => Z z i)]
  exact D.expect_congr fun z => sum_Z z

Variance and covariance

The two second-order summaries are defined the way they always are, as expectations of centred quantities.

noncomputable def cov : ℝ :=
  D.expect (fun z => (f z - D.expect f) * (g z - D.expect g))

noncomputable def var : ℝ := D.expect (fun z => (f z - D.expect f) ^ 2)

Nothing here is surprising; what matters is what the variance is over. It is over the randomization and nothing else. There is no superpopulation, so $\mathrm{Var}$ of a fixed potential outcome is not zero-because-large-n, it is zero because the potential outcome is a constant function of $z$ (var_const).

Proposition (the shortcut formula). For any design and any statistics $f,g$,

theorem cov_eq_expect_mul_sub_mul :
    D.cov f g = D.expect (fun z => f z * g z) - D.expect f * D.expect g := by
  unfold cov expect
  have h : ∀ z ∈ (univ : Finset (Assignment n)),
      D.prob z * ((f z - ∑ w, D.prob w * f w) * (g z - ∑ w, D.prob w * g w))
        = D.prob z * (f z * g z)
          - (∑ w, D.prob w * g w) * (D.prob z * f z)
          - (∑ w, D.prob w * f w) * (D.prob z * g z)
          + ((∑ w, D.prob w * f w) * (∑ w, D.prob w * g w)) * D.prob z :=
    fun z _ => by ring
  rw [Finset.sum_congr rfl h, Finset.sum_add_distrib, Finset.sum_sub_distrib,
    Finset.sum_sub_distrib, ← Finset.mul_sum, ← Finset.mul_sum, ← Finset.mul_sum, D.sum_one]
  ring

It is the longest proof in Design.lean and it is still a dozen lines, which is a fair summary of how much probability theory the finite-population view needs. The reason it is not three lines is that it does not use the expect lemmas at all: it drops to the sums underneath them, so that a single ring can do the expansion that the informal proof does in one line of algebra.

1. unfold cov expect

   Down to sums. Every D.expect on both sides — the outer one and the two inner centring constants — becomes an explicit ∑ z, D.prob z * …. The goal is ugly and that is deliberate: it is now a statement about finite sums of reals, and finite sums of reals are what Finset.sum_congr and ring know about.

   Goal

   ⊢ ∑ z, D.prob z * (fun z => (f z - ∑ z, D.prob z * f z) * (g z - ∑ z, D.prob z * g z)) z =
       ∑ z, D.prob z * (fun z => f z * g z) z - (∑ z, D.prob z * f z) * ∑ z, D.prob z * g z

2. have h : ∀ z ∈ univ, … := fun z _ => by ring

   The expansion, done once, at a fixed assignment. Writing $\mu$ for the sum ∑ w, D.prob w * f w and $\nu$ for ∑ w, D.prob w * g w, h says that the summand D.prob z * ((f z - μ) * (g z - ν)) equals the four-term expression whose terms are D.prob z * (f z * g z), then ν * (D.prob z * f z) and μ * (D.prob z * g z) subtracted, then (μ * ν) * D.prob z added — which is the informal $fg-\nu f-\mu g+\mu \nu$ multiplied through by $P(z)$, with each term arranged so that the constant stands to the left of a D.prob z * … factor. That arrangement is not fussiness: it is what lets the next two steps pull the constants out. Because $z$ is fixed, $\mu$ and $\nu$ are just numbers, and ring proves the identity without knowing anything about them. The goal shown is the state after Finset.sum_congr rfl h has replaced every summand by that expansion; the rfl there says the index set is unchanged.

   Goal

   ⊢ ∑ x,
         (D.prob x * (f x * g x) - (∑ w, D.prob w * g w) * (D.prob x * f x) - (∑ w, D.prob w * f w) * (D.prob x * g x) +
           ((∑ w, D.prob w * f w) * ∑ w, D.prob w * g w) * D.prob x) =
       ∑ z, D.prob z * (fun z => f z * g z) z - (∑ z, D.prob z * f z) * ∑ z, D.prob z * g z

3. Finset.sum_add_distrib, Finset.sum_sub_distrib, Finset.sum_sub_distrib

   Three splittings, applied outside in. The expression is $(A-B-C)+D$ at the top level, so sum_add_distrib splits off the + D first, and the two sum_sub_distribs then peel - C and - B. Four separate sums remain, one per term of the expansion.

   Goal

   ⊢ ∑ x, D.prob x * (f x * g x) - ∑ x, (∑ w, D.prob w * g w) * (D.prob x * f x) -
           ∑ x, (∑ w, D.prob w * f w) * (D.prob x * g x) +
         ∑ x, ((∑ w, D.prob w * f w) * ∑ w, D.prob w * g w) * D.prob x =
       ∑ z, D.prob z * (fun z => f z * g z) z - (∑ z, D.prob z * f z) * ∑ z, D.prob z * g z

4. ← Finset.mul_sum, ← Finset.mul_sum, ← Finset.mul_sum

   Finset.mul_sum : b * ∑ i ∈ s, f i = ∑ i ∈ s, b * f i, used backwards — the arrow — so each of the three sums with a constant in front of it has that constant pulled out. This is ”$\mu$ and $\nu$ are constants, so they come out of the expectation”, spelled at the level of sums.

   Goal

   ⊢ ∑ x, D.prob x * (f x * g x) - (∑ w, D.prob w * g w) * ∑ i, D.prob i * f i -
           (∑ w, D.prob w * f w) * ∑ i, D.prob i * g i +
         ((∑ w, D.prob w * f w) * ∑ w, D.prob w * g w) * ∑ i, D.prob i =
       ∑ z, D.prob z * (fun z => f z * g z) z - (∑ z, D.prob z * f z) * ∑ z, D.prob z * g z

5. D.sum_one

   The last of the three pulled-out sums is $\sum _{z}P(z)$ with nothing else in it, and the design’s second axiom replaces it by 1. This is the only place in the proof where the object being manipulated is a design rather than an arbitrary family of weights.

6. ring

   What remains is $\mathbb{E}[fg]-\nu \mu -\mu \nu +\mu \nu \cdot 1=\mathbb{E}[fg]-\mu \nu$ with $\mathbb{E}[fg]$, $\mu$ and $\nu$ standing for three particular sums. ring treats each sum as an opaque atom and does the cancellation, exactly as the last line of the informal proof does.

The variance form is then a corollary, and the two rewrites that get there are worth reading because the same manoeuvre — turn f z * f z into f z ^ 2 inside a binder — recurs whenever a squared statistic meets a product lemma.

theorem var_eq_expect_sq_sub_sq :
    D.var f = D.expect (fun z => f z ^ 2) - (D.expect f) ^ 2 := by
  rw [var_eq_cov_self, cov_eq_expect_mul_sub_mul]
  rw [D.expect_congr (f := fun z => f z * f z) (g := fun z => f z ^ 2) fun z => (sq (f z)).symm]
  rw [sq]

1. rw [var_eq_cov_self, cov_eq_expect_mul_sub_mul]

   var_eq_cov_self : D.var f = D.cov f f is proved just above it in the file by simp only [var, cov, sq]: a square is a product of a thing with itself, so the two definitions coincide. Rewriting with it and then with the shortcut formula turns the left-hand side into the covariance of f with itself, written out.

   Goal

   ⊢ (D.expect fun z => f z * f z) - D.expect f * D.expect f = (D.expect fun z => f z ^ 2) - D.expect f ^ 2

2. rw [D.expect_congr (f := fun z => f z * f z) (g := fun z => f z ^ 2) fun z => (sq (f z)).symm]

   expect_congr says two statistics that agree pointwise have the same expectation; here the pointwise fact is sq (f z) : f z ^ 2 = f z * f z, reversed with .symm, so the product inside the expectation becomes a square. The two statistics have to be named explicitly — the (f := …) syntax — because a plain rw cannot do this: rewriting under a binder is beyond rw, and expect_congr exists precisely to smuggle a pointwise rewrite past one.

   Goal

   ⊢ (D.expect fun z => f z ^ 2) - D.expect f * D.expect f = (D.expect fun z => f z ^ 2) - D.expect f ^ 2

3. rw [sq]

   The two sides now differ in one place only: D.expect f * D.expect f on the left against D.expect f ^ 2 on the right. rw [sq] rewrites the first occurrence of a square it can reach — not the f z ^ 2 under the binder, which it cannot get at — so it hits D.expect f ^ 2, turns it into the product, and the two sides become identical.

theorem ex_var_Z (D : Design n) (i : Fin n) :
    D.var (fun z => Z z i) = D.propensity i * (1 - D.propensity i) := by
  sorry

theorem ex_var_Z (D : Design n) (i : Fin n) :
    D.var (fun z => Z z i) = D.propensity i * (1 - D.propensity i) := by
  rw [Design.var_eq_expect_sq_sub_sq, D.expect_congr fun z => Z_sq z i, Design.expect_Z]
  ring

theorem ex_cov_Z (D : Design n) (i j : Fin n) :
    D.cov (fun z => Z z i) (fun z => Z z j)
      = D.jointPropensity i j - D.propensity i * D.propensity j := by
  sorry

theorem ex_cov_Z (D : Design n) (i j : Fin n) :
    D.cov (fun z => Z z i) (fun z => Z z j)
      = D.jointPropensity i j - D.propensity i * D.propensity j := by
  rw [Design.cov_eq_expect_mul_sub_mul, Design.expect_Z, Design.expect_Z]
  rfl

Computing with a design you can see

Three units, eight assignments. At this size a design is a table, and Lean will evaluate it for you.

#eval (univ : Finset (Assignment 3))

theorem card_assignments_three : Fintype.card (Assignment 3) = 8 := by decide

decide is the tactic for a goal that is decidable: Lean compiles the proposition to a Boolean-valued computation, runs it, and accepts the goal if the answer is true. It works here because everything in sight is finite — sets of Fin 3, cardinalities, memberships — and it fails immediately on anything involving ℝ, where equality is not decidable. Keep that division in mind: the combinatorial half of a design computation is decide’s job, the arithmetic half is norm_num’s.

To compute an expectation we need the eight-term expansion:

theorem sum_assignments_three (f : Assignment 3 → ℝ) :
    ∑ z, f z
      = f ∅ + f {0} + f {1} + f {2} + f {0, 1} + f {0, 2} + f {1, 2} + f {0, 1, 2} := by
  rw [show (univ : Finset (Assignment 3))
      = {∅, {0}, {1}, {2}, {0, 1}, {0, 2}, {1, 2}, {0, 1, 2}} from by decide]
  repeat rw [Finset.sum_insert (by decide)]
  rw [Finset.sum_singleton]
  ring

1. rw [show (univ : Finset (Assignment 3)) = {∅, {0}, {1}, {2}, {0, 1}, {0, 2}, {1, 2}, {0, 1, 2}} from by decide]

   show T from e is a proof of T supplied inline, and here T is the Finset equality “the eight assignments listed are all of them”. decide checks it by computation. Rewriting with it replaces univ in the goal’s sum by the explicit list, so the sum now has a syntactic head to peel.

2. repeat rw [Finset.sum_insert (by decide)]

   {a, b, c} is notation for insert a (insert b {c}), and Finset.sum_insert : a ∉ s → ∑ x ∈ insert a s, f x = f a + ∑ x ∈ s, f x peels one element off the front, provided that element is not in the rest — a decidable side condition, hence the inline (by decide). repeat applies the rewrite until it fails, which is after all seven insertions.

3. rw [Finset.sum_singleton]

   What is left of the sum is the last element on its own, and Finset.sum_singleton : ∑ x ∈ {a}, f x = f a removes the sum sign.

4. ring

   Both sides are now the same eight terms added up, but the peeling produced them right-nested, f ∅ + (f {0} + (…)), while the statement is written left-nested. ring reassociates. It is the only arithmetic in the proof, and it does no arithmetic.

Now a real design, written out as a table:

noncomputable def oneOfThree : Design 3 where
  prob z := if #z = 1 then 1 / 3 else 0
  nonneg z := by split <;> norm_num
  sum_one := by
    rw [sum_assignments_three fun z => if #z = 1 then (1 : ℝ) / 3 else 0]
    norm_num [Finset.card_insert_of_notMem, Fin.ext_iff]

@[simp] theorem oneOfThree_prob (z : Assignment 3) :
    oneOfThree.prob z = if #z = 1 then 1 / 3 else 0 := rfl

The two proof obligations are exactly the two axioms.

1. nonneg z := by split <;> norm_num

   split on the if in the pmf gives two goals, 0 ≤ 1 / 3 and 0 ≤ 0, and norm_num closes both. A field of a structure can be given by a tactic proof like any other goal; nonneg z := by … is the field, with the assignment z already in scope.

2. rw [sum_assignments_three fun z => if #z = 1 then (1 : ℝ) / 3 else 0]

   The sum_one field. The lemma just proved is applied to this particular summand, turning the sum over univ into the eight explicit terms. The lambda has to be written out because Lean will not guess which function of z is meant.

3. norm_num [Finset.card_insert_of_notMem, Fin.ext_iff]

   Eight ifs, each conditional on a cardinality: #∅ = 1 is false, #{0} = 1 is true, #{0, 1} = 1 is false, and so on. Two lemmas are needed to settle them. Finset.card_insert_of_notMem : a ∉ s → #(insert a s) = #s + 1 peels an element off a literal set, and its side condition is a disequality of Fin 3 values; Fin.ext_iff reduces (0 : Fin 3) ≠ 2 to 0 ≠ 2 in ℕ, which norm_num can decide. Without it the side conditions stick, the cardinalities never evaluate, and the sum stops half-way. With them, three terms are 1 / 3 and five are 0.

4. @[simp] theorem oneOfThree_prob (z : Assignment 3) : oneOfThree.prob z = if #z = 1 then 1 / 3 else 0 := rfl

   The pmf as a rewrite rule. It is true by rfl — prob is that if by definition — and it is tagged @[simp] so that later proofs can get at the table without unfolding the structure by hand.

Then the propensity is a computation rather than a theorem:

theorem oneOfThree_propensity (i : Fin 3) : oneOfThree.propensity i = 1 / 3 := by
  rw [Design.propensity, Design.expect,
    sum_assignments_three fun z => oneOfThree.prob z * Z z i]
  fin_cases i <;>
    norm_num [Z, Fin.ext_iff, Finset.card_insert_of_notMem]

theorem oneOfThree_prob_eq (z : Assignment 3) :
    oneOfThree.prob z = (completeRandomization 3 1 (by norm_num)).prob z := by
  rw [oneOfThree_prob, completeRandomization_prob]
  norm_num

1. rw [Design.propensity, Design.expect, sum_assignments_three fun z => oneOfThree.prob z * Z z i]

   Unfold the propensity to an expectation, the expectation to a sum, and expand that sum into its eight terms. Every quantity in the goal is now a number written out.

   Goal

   i : Fin 3
   ⊢ oneOfThree.prob ∅ * Z ∅ i + oneOfThree.prob {0} * Z {0} i + oneOfThree.prob {1} * Z {1} i +
                 oneOfThree.prob {2} * Z {2} i +
               oneOfThree.prob {0, 1} * Z {0, 1} i +
             oneOfThree.prob {0, 2} * Z {0, 2} i +
           oneOfThree.prob {1, 2} * Z {1, 2} i +
         oneOfThree.prob {0, 1, 2} * Z {0, 1, 2} i =
       1 / 3

2. fin_cases i

   i : Fin 3 still stands for an arbitrary unit, and Z {0, 1} i cannot be evaluated until it is known which unit that is. fin_cases produces the three concrete goals.

3. <;> norm_num [Z, Fin.ext_iff, Finset.card_insert_of_notMem]

   Each branch is eight probabilities times eight indicator evaluations. Z unfolds the indicator to its if, Fin.ext_iff decides the memberships ((0 : Fin 3) ∈ {1, 2} reduces to two disequalities in ℕ), and Finset.card_insert_of_notMem settles the cardinalities inside oneOfThree’s own if, as it did above. Exactly one of the eight terms survives: five assignments have probability 0 because their size is not 1, two of the three singletons do not contain i and so contribute 1 / 3 * 0, and the singleton {i} contributes 1 / 3 * 1.

4. rw [oneOfThree_prob, completeRandomization_prob]

   The second theorem, which checks the hand-written design against the library’s. Both rewrites replace a pmf by its defining if, leaving (if #z = 1 then 1 / 3 else 0) = if #z = 1 then (↑(Nat.choose 3 1))⁻¹ else 0.

5. norm_num

   $(\genfrac{}{}{0ex}{}{3}{1})=3$ and $3^{-1}=1/3$, so the two branches agree and the ifs are equal. Note what this second theorem buys: completeRandomization 3 1 has pmf $1/(\genfrac{}{}{0ex}{}{3}{1})$ on the singletons, so the general propensity theorem gives ${\pi }_i=1/3$ with no enumeration at all — and the brute-force computation above is then a check on the general theorem, not a substitute for it.

Splitting into the three units and evaluating eight indicators apiece is the honest way to check a design on a small example, and it is also a good way to convince yourself that a general theorem you are about to read says what you think it says.

The Bernoulli design

Flip an independent coin per unit. $P(z)=p^{|z|}(1-p{)}^{n-|z|}$.

noncomputable def bernoulliDesign (n : ℕ) (p : ℝ) (hp0 : 0 ≤ p) (hp1 : p ≤ 1) : Design n where
  prob z := p ^ #z * (1 - p) ^ (n - #z)
  nonneg z := mul_nonneg (pow_nonneg hp0 _) (pow_nonneg (by linarith) _)
  sum_one := by
    have h := sum_bernoulli_weight_mul_subset n p (∅ : Finset (Fin n))
    simpa using h

1. nonneg z := mul_nonneg (pow_nonneg hp0 _) (pow_nonneg (by linarith) _)

   The first axiom, in term mode. mul_nonneg needs both factors nonnegative, and each factor is a power, so pow_nonneg reduces each to a statement about its base: hp0 : 0 ≤ p for the first, and for the second 0 ≤ 1 - p, which linarith derives from hp1 : p ≤ 1. The _s are the exponents, which Lean reads off the goal.

2. have h := sum_bernoulli_weight_mul_subset n p (∅ : Finset (Fin n))

   The second axiom is the binomial theorem, and it is the $S=\mathrm{\varnothing }$ case of the one lemma this design needs. Instantiated there, that lemma reads ∑ z, p ^ #z * (1 - p) ^ (n - #z) * (if ∅ ⊆ z then 1 else 0) = p ^ #∅.

3. simpa using h

   simpa … using h simplifies both h and the goal and then matches them. Here the simplification empties out the empty set on both sides of h: the condition ∅ ⊆ z holds always, so the indicator is 1 and the factor disappears, and #∅ = 0 makes the right-hand side p ^ 0 = 1. What is left is the goal.

That one lemma is worth the page it takes.

theorem sum_bernoulli_weight_mul_subset (n : ℕ) (p : ℝ) (S : Finset (Fin n)) :
    ∑ z : Assignment n, p ^ #z * (1 - p) ^ (n - #z) * (if S ⊆ z then 1 else 0) = p ^ #S := by
  classical
  set g : Fin n → ℝ := fun i => if i ∈ S then 0 else (1 - p) with hg
  have key := Finset.prod_add (fun _ : Fin n => p) g univ
  rw [Finset.powerset_univ] at key
  have hL : ∏ _i : Fin n, (p + g _i) = p ^ #S := by
    have hpt : ∀ i : Fin n, p + g i = if i ∈ S then p else 1 := by
      intro i
      simp only [hg]
      by_cases h : i ∈ S <;> simp [h]
    rw [Finset.prod_congr rfl fun i _ => hpt i, Finset.prod_ite_mem, Finset.univ_inter,
      Finset.prod_const]
  have hR : ∀ t : Assignment n,
      (∏ _i ∈ t, p) * (∏ i ∈ univ \ t, g i)
        = p ^ #t * (1 - p) ^ (n - #t) * (if S ⊆ t then 1 else 0) := by
    intro t
    have hcard : #(univ \ t) = n - #t := by
      rw [← Finset.compl_eq_univ_sdiff, Finset.card_compl, Fintype.card_fin]
    rw [Finset.prod_const]
    by_cases h : S ⊆ t
    ·
      rw [ite_eq_left h, mul_one]
      have : ∀ i ∈ univ \ t, g i = 1 - p := by
        intro i hi
        simp only [hg]
        exact ite_eq_right fun hiS => absurd (h hiS) (Finset.mem_sdiff.mp hi).2
      rw [Finset.prod_congr rfl this, Finset.prod_const, hcard]
    ·
      rw [ite_eq_right h, mul_zero]
      refine mul_eq_zero_of_right _ ?_
      obtain ⟨i, hiS, hit⟩ := Finset.not_subset.mp h
      refine Finset.prod_eq_zero (Finset.mem_sdiff.mpr ⟨Finset.mem_univ i, hit⟩) ?_
      simp only [hg]
      exact ite_eq_left hiS
  rw [← hL, key]
  exact (Finset.sum_congr rfl fun t _ => hR t).symm

The mechanism is Finset.prod_add, which is the “multiply out the brackets” step of a first probability course:

A sum over subsets of the population is a sum over assignments — this is the payoff of representing an assignment as a treated set rather than as a vector of indicators. Take $a_i=p$ for every $i$ and

and the right-hand side becomes $\sum _z{p}^{|z|}(1-p{)}^{n-|z|}\mathbf{1}\{S\subseteq z\}$, because the factor $0$ kills exactly the assignments that fail to treat all of $S$. The left-hand side collapses to $\prod _{i\in S}p=p^{|S|}$. Setting $S=\mathrm{\varnothing }$ gives the pmf summing to one; that is why the design’s sum_one field is a one-liner citing this lemma.

The Lean is that argument with the two sides of Finset.prod_add identified separately — the left as hL, the right termwise as hR — and then joined.

1. classical

   The family g below is defined by an if on i ∈ S, and Finset.prod_add wants a DecidableEq instance. classical supplies decidability for everything in scope so that neither has to be threaded through by hand. It changes no mathematics; it changes which instances Lean will accept.

2. set g : Fin n → ℝ := fun i => if i ∈ S then 0 else (1 - p) with hg

   The $b_i$ of the informal derivation: $0$ on $S$, $1-p$ off it. set introduces the abbreviation and replaces any matching expression in the goal by the new name, and with hg keeps the defining equation around, which the sub-proofs below need in order to unfold g again.

3. have key := Finset.prod_add (fun _ : Fin n => p) g univ

   “Multiply out the brackets”, instantiated at $a_i=p$ and $b_i=g_i$: ∏ i ∈ s, (f i + g i) = ∑ t ∈ s.powerset, (∏ i ∈ t, f i) * ∏ i ∈ s \ t, g i. Nothing is proved by this step; it puts the identity into the context, where the two rewrites at the end can use it.

4. rw [Finset.powerset_univ] at key

   prod_add sums over univ.powerset, the set of subsets of the universe; Finset.powerset_univ : univ.powerset = univ says that this is the same as the universe of the type Finset (Fin n) — the set of all assignments. This is the step where “a sum over subsets of the population is a sum over assignments” becomes a rewrite rather than a remark.

   Goal

   g : Fin n → ℝ := fun i => if i ∈ S then 0 else 1 - p
   key : ∏ i, (p + g i) = ∑ t, (∏ i ∈ t, p) * ∏ i ∈ univ \ t, g i
   ⊢ (∑ z, p ^ #z * (1 - p) ^ (n - #z) * if S ⊆ z then 1 else 0) = p ^ #S

5. have hpt : ∀ i : Fin n, p + g i = if i ∈ S then p else 1

   The left-hand side of prod_add, factor by factor. Inside the have, simp only [hg] unfolds g and by_cases h : i ∈ S <;> simp [h] checks the two cases: on $S$ the factor is $p+0=p$, off it $p+(1-p)=1$.

6. rw [Finset.prod_congr rfl fun i _ => hpt i, Finset.prod_ite_mem, Finset.univ_inter, Finset.prod_const]

   Four rewrites finishing hL, and the goal shown is the state after the first. Finset.prod_congr rfl (…) replaces every factor using hpt (the rfl says the index set is unchanged). Finset.prod_ite_mem : ∏ i ∈ s, (if i ∈ t then f i else 1) = ∏ i ∈ s ∩ t, f i drops the factors equal to 1, leaving a product over univ ∩ S; Finset.univ_inter simplifies that to S; and Finset.prod_const : ∏ _x ∈ s, b = b ^ #s gives p ^ #S. That is the informal “the left-hand side collapses to $p^{|S|}$”.

   Goal

   hpt : ∀ (i : Fin n), p + g i = if i ∈ S then p else 1
   ⊢ (∏ i, if i ∈ S then p else 1) = p ^ #S

7. have hcard : #(univ \ t) = n - #t

   Now hR, the right-hand side, one assignment t at a time; the goal shown is its statement with hcard already available. The little cardinality fact is proved by rw [← Finset.compl_eq_univ_sdiff, Finset.card_compl, Fintype.card_fin]: “everything except t” is the complement of t, whose size is n - #t. It is needed because the exponent on the right-hand side is written n - #t while the product on the left ranges over univ \ t.

   Goal

   t : Assignment n
   hcard : #(univ \ t) = n - #t
   ⊢ (∏ _i ∈ t, p) * ∏ i ∈ univ \ t, g i = p ^ #t * (1 - p) ^ (n - #t) * if S ⊆ t then 1 else 0

8. rw [Finset.prod_const]

   The first factor ∏ _i ∈ t, p is a constant product, so it is p ^ #t — which is already the first factor on the right. by_cases h : S ⊆ t then splits into the two branches shown, in which the indicator has been evaluated by ite_eq_left h and mul_one, respectively ite_eq_right h and mul_zero.

   Goal

   case pos
   h : S ⊆ t
   ⊢ p ^ #t * ∏ i ∈ univ \ t, g i = p ^ #t * (1 - p) ^ (n - #t)
   
   case neg
   h : ¬S ⊆ t
   ⊢ p ^ #t * ∏ i ∈ univ \ t, g i = 0

9. rw [Finset.prod_congr rfl this, Finset.prod_const, hcard]

   The branch where $S\subseteq t$. Then no $i$ outside $t$ lies in $S$ — that is what the little have above this line proves, using Finset.mem_sdiff to get i ∉ t and h to rule out i ∈ S — so every factor of ∏ i ∈ univ \ t, g i is 1 - p. prod_const makes that (1 - p) ^ #(univ \ t) and hcard rewrites the exponent to n - #t.

10. refine Finset.prod_eq_zero (Finset.mem_sdiff.mpr ⟨Finset.mem_univ i, hit⟩) ?_

    The branch where $S⊈t$. The right-hand side is 0, and mul_eq_zero_of_right reduces the goal to showing that the second factor, the product over univ \ t, vanishes. Finset.not_subset.mp h produces a witness: some i with hiS : i ∈ S and hit : i ∉ t. That i lies in univ \ t, and g i = 0 there because i ∈ S — the remaining ?_, closed by ite_eq_left hiS. Finset.prod_eq_zero then kills the whole product: this is the informal “the factor $0$ kills exactly the assignments that fail to treat all of $S$”.

11. rw [← hL, key]

    Back in the main goal. ← hL rewrites the right-hand side p ^ #S backwards into the product ∏ i, (p + g i), and key then turns that product into the sum over all assignments. Both sides are now sums over the same index set.

    Goal

    ⊢ (∑ z, p ^ #z * (1 - p) ^ (n - #z) * if S ⊆ z then 1 else 0) =
        ∑ t, (∏ i ∈ t, p) * ∏ i ∈ univ \ t, g i

12. exact (Finset.sum_congr rfl fun t _ => hR t).symm

    Two sums over the same set whose summands agree are equal, which is Finset.sum_congr; hR is the termwise agreement. It proves the equation in the opposite direction to the goal, hence .symm.

Everything else about the Bernoulli design is a corollary.

Lemma. Under the Bernoulli design with parameter $p$, every unit has propensity ${\pi }_i=p$.

theorem bernoulli_propensity (i : Fin n) : (bernoulliDesign n p hp0 hp1).propensity i = p := by
  rw [Design.propensity_eq_inclusion, bernoulli_inclusion, Finset.card_singleton, pow_one]

1. rw [Design.propensity_eq_inclusion, bernoulli_inclusion, Finset.card_singleton, pow_one]

   Four rewrites, in the order of the informal proof. propensity_eq_inclusion turns ${\pi }_i$ into the inclusion probability of {i}. bernoulli_inclusion, which is the master computation transported to this design by unfolding inclusion and expect and applying prod_Z, evaluates that to p ^ #{i}. Finset.card_singleton computes #{i} = 1, and pow_one finishes. Every later Bernoulli fact — the joint propensity, the expected number treated — is this same chain with a different set in place of {i}.

theorem bernoulli_jointPropensity {i j : Fin n} (hij : i ≠ j) :
    (bernoulliDesign n p hp0 hp1).jointPropensity i j = p ^ 2 := by
  sorry

theorem bernoulli_jointPropensity {i j : Fin n} (hij : i ≠ j) :
    (bernoulliDesign n p hp0 hp1).jointPropensity i j = p ^ 2 := by
  rw [Design.jointPropensity_eq_inclusion _ hij, bernoulli_inclusion,
    Finset.card_insert_of_notMem (by simpa using hij), Finset.card_singleton]

theorem ex_bernoulli_expected_treated :
    (bernoulliDesign n p hp0 hp1).expect (fun z => (#z : ℝ)) = n * p := by
  sorry

theorem ex_bernoulli_expected_treated :
    (bernoulliDesign n p hp0 hp1).expect (fun z => (#z : ℝ)) = n * p := by
  rw [← Design.sum_propensity]
  simp only [bernoulli_propensity, Finset.sum_const, Finset.card_univ, Fintype.card_fin,
    nsmul_eq_mul]

Complete randomization

Fix $n_1$, and draw the treated set uniformly from the $(\genfrac{}{}{0ex}{}{n}{n_1})$ sets of that size.

noncomputable def completeRandomization (n n₁ : ℕ) (h : n₁ ≤ n) : Design n where
  prob z := if #z = n₁ then ((n.choose n₁ : ℕ) : ℝ)⁻¹ else 0
  nonneg z := by
    by_cases hz : #z = n₁
    · simp only [hz, ite_eq_left]
      exact inv_nonneg.mpr (Nat.cast_nonneg (n.choose n₁))
    · simp [hz]
  sum_one := by
    have hpos : 0 < n.choose n₁ := Nat.choose_pos h
    have hne : ((n.choose n₁ : ℕ) : ℝ) ≠ 0 := Nat.cast_ne_zero.mpr hpos.ne'
    calc ∑ z : Assignment n, (if #z = n₁ then ((n.choose n₁ : ℕ) : ℝ)⁻¹ else 0)
        = ∑ z ∈ (univ : Finset (Fin n)).powersetCard n₁, ((n.choose n₁ : ℕ) : ℝ)⁻¹ := by
          rw [← Finset.sum_filter]
          congr 1
          ext z
          simp [Finset.mem_powersetCard]
      _ = ((n.choose n₁ : ℕ) : ℝ) * ((n.choose n₁ : ℕ) : ℝ)⁻¹ := by
          rw [Finset.sum_const, nsmul_eq_mul, Finset.card_powersetCard, Finset.card_univ,
            Fintype.card_fin]
      _ = 1 := mul_inv_cancel₀ hne

The pmf is $1/(\genfrac{}{}{0ex}{}{n}{n_1})$ on the sets of size $n_1$ and $0$ elsewhere.

Lemma. Suppose $n_1\le n$. Then $\sum _{z}P(z)=1$ for the pmf above; that is, complete randomization is a design.

The hypothesis is doing that one job and no other.

1. by_cases hz : #z = n₁

   The nonneg field. In the branch where #z = n₁, simp only [hz, ite_eq_left] selects the then branch and inv_nonneg.mpr (Nat.cast_nonneg _) observes that the inverse of a nonnegative real is nonnegative — a cast of a natural number is nonnegative whatever it is, so nothing about $(\genfrac{}{}{0ex}{}{n}{n_1})$ is needed here. Otherwise simp [hz] selects 0.

2. have hne : ((n.choose n₁ : ℕ) : ℝ) ≠ 0 := Nat.cast_ne_zero.mpr hpos.ne'

   The sum_one field opens by recording the one fact the hypothesis buys: Nat.choose_pos h : 0 < n.choose n₁ (this is where h : n₁ ≤ n is used), and then Nat.cast_ne_zero moves that from ℕ to ℝ. Without it the final cancellation is false, not merely unprovable.

3. rw [← Finset.sum_filter]

   The calc chain begins. The sum has if #z = n₁ then _ else 0 as its summand, and Finset.sum_filter, used backwards, drops the zero terms: the sum is now over the assignments satisfying the condition. What has to be shown is that this filtered set is powersetCard n₁ univ.

   Goal

   ⊢ ∑ a with #a = n₁, (↑(n.choose n₁))⁻¹ = ∑ z ∈ powersetCard n₁ univ, (↑(n.choose n₁))⁻¹

4. congr 1

   The two sides are sums of the same constant over two index sets, so congr 1 peels off the summand and leaves the index sets to be compared; ext z then reduces set equality to membership, one assignment at a time, and simp [Finset.mem_powersetCard] finishes — z ∈ powersetCard n₁ univ unfolds to z ⊆ univ ∧ #z = n₁, whose first half is free.

   Goal

   z : Assignment n
   ⊢ z ∈ {a | #a = n₁} ↔ z ∈ powersetCard n₁ univ

5. rw [Finset.sum_const, nsmul_eq_mul, Finset.card_powersetCard, Finset.card_univ, Fintype.card_fin]

   The counting step. Finset.sum_const : ∑ _x ∈ s, b = #s • b replaces the sum of a constant by a scalar multiple, nsmul_eq_mul turns that ℕ-action into ordinary multiplication by a cast, and Finset.card_powersetCard : #(powersetCard n s) = (#s).choose n supplies the count. card_univ and Fintype.card_fin then rewrite #(univ : Finset (Fin n)) as n, giving $(\genfrac{}{}{0ex}{}{n}{n_1})\cdot {(\genfrac{}{}{0ex}{}{n}{n_1})}^{-1}$.

6. _ = 1 := mul_inv_cancel₀ hne

   mul_inv_cancel₀ : a ≠ 0 → a * a⁻¹ = 1, with hne as the side condition. The _ is calc notation for the previous right-hand side.

The inclusion probability is again a counting problem, and Mathlib already knows the count.

Lemma. Let $S$ be a set of units and let $k\ge |S|$. The number of $k$-element subsets of $\{1,\dots ,n\}$ that contain $S$ is $(\genfrac{}{}{0ex}{}{n-|S|}{k-|S|})$.

theorem card_filter_subset_powersetCard (S : Finset (Fin n)) (k : ℕ) (hS : #S ≤ k) :
    #(((univ : Finset (Fin n)).powersetCard k).filter (S ⊆ ·)) = (n - #S).choose (k - #S) := by
  have h := Finset.card_filter_powersetCard_subset S (univ : Finset (Fin n)) k
    (Finset.subset_univ S) hS
  rwa [Finset.card_univ, Fintype.card_fin] at h

1. have h := Finset.card_filter_powersetCard_subset S (univ : Finset (Fin n)) k (Finset.subset_univ S) hS

   Mathlib states the count for subsets of an arbitrary finite set t: #{x ∈ powersetCard n t | s ⊆ x} = (#t - #s).choose (n - #s), given s ⊆ t and #s ≤ n. Instantiating t := univ needs the first side condition (Finset.subset_univ, free) and the second (hS, the hypothesis).

2. rwa [Finset.card_univ, Fintype.card_fin] at h

   h now says the count is (#univ - #S).choose (k - #S) and the goal says (n - #S).choose (k - #S). The two #univ rewrites — Finset.card_univ to Fintype.card (Fin n), then Fintype.card_fin to n — close the gap, and the a of rwa supplies the resulting hypothesis as the proof. The whole lemma is bookkeeping over the fact that the population is univ.

Having committed to treating all of $S$, one still chooses $n_1-|S|$ units from the remaining $n-|S|$. So

which for $S=\{i\}$ is $(\genfrac{}{}{0ex}{}{n-1}{n_1-1})/(\genfrac{}{}{0ex}{}{n}{n_1})=n_1/n$ and for $S=\{i,j\}$ is $(\genfrac{}{}{0ex}{}{n-2}{n_1-2})/(\genfrac{}{}{0ex}{}{n}{n_1})=n_1(n_1-1)/(n(n-1))$. Those two ratios are pure Nat.choose algebra:

theorem choose_ratio_one (h1 : 1 ≤ n₁) (hn : n₁ ≤ n) :
    (((n - 1).choose (n₁ - 1) : ℕ) : ℝ) / ((n.choose n₁ : ℕ) : ℝ) = (n₁ : ℝ) / (n : ℝ) := by
  obtain ⟨j, rfl⟩ : ∃ j, n₁ = j + 1 := ⟨n₁ - 1, by omega⟩
  obtain ⟨m, rfl⟩ : ∃ m, n = m + 1 := ⟨n - 1, by omega⟩
  have hjm : j ≤ m := by omega
  have hpos : 0 < (m + 1).choose (j + 1) := Nat.choose_pos (by omega)
  have hne : ((m + 1).choose (j + 1) : ℝ) ≠ 0 := Nat.cast_ne_zero.mpr hpos.ne'
  have hm : ((m : ℝ) + 1) ≠ 0 := by positivity
  have key : ((m + 1) * m.choose j : ℕ) = ((m + 1).choose (j + 1) * (j + 1) : ℕ) :=
    Nat.add_one_mul_choose_eq m j
  have key' : ((m : ℝ) + 1) * (m.choose j : ℝ) = ((m + 1).choose (j + 1) : ℝ) * ((j : ℝ) + 1) := by
    exact_mod_cast key
  simp only [Nat.add_sub_cancel]
  push_cast
  rw [div_eq_div_iff hne hm]
  linear_combination key'

On paper the ratio is one cancellation. In Lean the obstacle is that n - 1 and n₁ - 1 are natural subtractions, and the proof’s whole strategy is to make sure they never have to be reasoned about.

1. obtain ⟨j, rfl⟩ : ∃ j, n₁ = j + 1 := ⟨n₁ - 1, by omega⟩

   The trick. Rather than reason about n₁ - 1, produce a j with n₁ = j + 1 and let the rfl pattern substitute it away: after this line the variable n₁ no longer exists, and every occurrence has become j + 1. The witness is n₁ - 1 and the equation is arithmetic in ℕ, which omega decides using h1 : 1 ≤ n₁. The second obtain does the same for n, and the goal shown is the state after both.

   Goal

   j m : ℕ
   h1 : 1 ≤ j + 1
   hn : j + 1 ≤ m + 1
   ⊢ ↑((m + 1 - 1).choose (j + 1 - 1)) / ↑((m + 1).choose (j + 1)) = ↑(j + 1) / ↑(m + 1)

2. have hpos : 0 < (m + 1).choose (j + 1) := Nat.choose_pos (by omega)

   Side conditions, gathered before the algebra. The line before this one, have hjm : j ≤ m := by omega, is hn restated without the + 1s, for omega’s convenience. Then Nat.choose_pos needs j + 1 ≤ m + 1, which omega supplies; hne casts that positivity into ℝ as a non-vanishing denominator; and hm : (m : ℝ) + 1 ≠ 0 is the other denominator, by positivity. Both are needed by name in the rewrite below — clearing a denominator is only legitimate when it is nonzero, and Lean will not take that on trust.

3. have key' : ((m : ℝ) + 1) * (m.choose j : ℝ) = ((m + 1).choose (j + 1) : ℝ) * ((j : ℝ) + 1) := by exact_mod_cast key

   The one real fact about binomial coefficients: Nat.add_one_mul_choose_eq : (n + 1) * C(n, k) = C(n + 1, k + 1) * (k + 1), which is key, stated in ℕ. exact_mod_cast restates it in ℝ, moving every cast to where the goal wants it — the sort of step that is invisible on paper and unavoidable here.

4. simp only [Nat.add_sub_cancel]

   Nat.add_sub_cancel : n + m - m = n. Now that n₁ is j + 1 and n is m + 1, the truncated subtractions in the statement are literally j + 1 - 1 and m + 1 - 1, and this rewrites them to j and m. Truncation never gets a chance to bite.

5. push_cast

   The casts are still arranged as ↑(j + 1) — a cast of a sum. push_cast pushes every cast towards the leaves, giving ↑j + 1, which is the shape key' and hm are stated in.

6. rw [div_eq_div_iff hne hm]

   div_eq_div_iff : b ≠ 0 → d ≠ 0 → (a / b = c / d ↔ a * d = c * b). This is where the two non-vanishing facts are spent, and what is left is a polynomial identity in the two binomial coefficients and ↑m, ↑j.

   Goal

   ⊢ ↑(m.choose j) * (↑m + 1) = (↑j + 1) * ↑((m + 1).choose (j + 1))

7. linear_combination key'

   linear_combination e proves a goal lhs = rhs by checking that lhs - rhs - (e.lhs - e.rhs) = 0 as a polynomial identity — here with coefficient 1 on key'. The goal and key' are the same equation with the factors written in a different order, so the difference is zero and ring inside linear_combination sees it. No division and no cancellation is performed anywhere in ℕ.

The two-element version, choose_ratio_two, is the same proof with + 2 for + 1 and two applications of Nat.add_one_mul_choose_eq chained by a calc instead of one.

And the payoff, twice.

Theorem. Suppose $n_1\le n$. Under complete randomization every unit has propensity ${\pi }_i=n_1/n$, and any two distinct units $i\ne j$ have joint propensity ${\pi }_{ij}=n_1(n_1-1)/(n(n-1))$.

theorem completeRandomization_propensity (i : Fin n) :
    (completeRandomization n n₁ h).propensity i = (n₁ : ℝ) / (n : ℝ) := by
  have hn : 0 < n := Nat.lt_of_le_of_lt (Nat.zero_le (i : ℕ)) i.isLt
  have hcardS : #({i} : Finset (Fin n)) = 1 := Finset.card_singleton i
  rw [Design.propensity_eq_inclusion, completeRandomization_inclusion]
  rcases Nat.lt_or_ge n₁ 1 with hlt | hge
  ·
    have hz : n₁ = 0 := by omega
    rw [card_filter_subset_powersetCard_of_lt _ _ (by omega)]
    simp [hz]
  · rw [card_filter_subset_powersetCard _ _ (by omega), hcardS, ← div_eq_mul_inv]
    exact choose_ratio_one hge h

1. have hn : 0 < n := Nat.lt_of_le_of_lt (Nat.zero_le (i : ℕ)) i.isLt

   There is a unit, so there is at least one unit: i.isLt is the proof that i’s underlying natural number is < n. This is the Lean way of saying “the statement quantifies over a unit, so the population is non-empty” — a fact a textbook would not stop for and a proof about $n_1/n$ needs. The next line, have hcardS : #({i} : Finset (Fin n)) = 1 := Finset.card_singleton i, records the other triviality the counting lemma will hand back: a singleton has one element.

2. rw [Design.propensity_eq_inclusion, completeRandomization_inclusion]

   The same two moves as under the Bernoulli design: propensity to inclusion probability of {i}, then this design’s inclusion lemma. That lemma is the counting statement in probability form — the number of $n_1$-sets containing {i}, times ${(\genfrac{}{}{0ex}{}{n}{n_1})}^{-1}$ — and its own proof is the pmf’s if pushed through prod_Z and Finset.sum_ite_mem.

   Goal

   i : Fin n
   hn : 0 < n
   hcardS : #{i} = 1
   ⊢ ↑(#({x ∈ powersetCard n₁ univ | {i} ⊆ x})) * (↑(n.choose n₁))⁻¹ = ↑n₁ / ↑n

3. rcases Nat.lt_or_ge n₁ 1 with hlt | hge

   The case distinction the counting lemma forces: it needs #S ≤ n₁, which for S = {i} means 1 ≤ n₁. rcases on Nat.lt_or_ge gives two branches, and what is shown is each branch after its opening rewrites: on the left card_filter_subset_powersetCard_of_lt has returned a count of 0, on the right the three rewrites of step 5 have turned the count into a ratio of binomial coefficients.

   Goal

   case inl
   hlt : n₁ < 1
   hz : n₁ = 0
   ⊢ ↑0 * (↑(n.choose n₁))⁻¹ = ↑n₁ / ↑n
   
   case inr
   hge : n₁ ≥ 1
   ⊢ ↑((n - 1).choose (n₁ - 1)) / ↑(n.choose n₁) = ↑n₁ / ↑n

4. simp [hz]

   The degenerate branch, where hz : n₁ = 0. The left side is 0 because the count is, and the right side is 0 / n, which is 0. Nothing here needs hn or the value of $(\genfrac{}{}{0ex}{}{n}{n_1})$.

5. rw [card_filter_subset_powersetCard _ _ (by omega), hcardS, ← div_eq_mul_inv]

   The real branch. The count is (n - #{i}).choose (n₁ - #{i}), hcardS rewrites #{i} to 1, and ← div_eq_mul_inv turns a * b⁻¹ back into a / b, which is the shape choose_ratio_one is stated in. Rewriting into a quotient rather than out of one is the usual direction here: Mathlib prefers a * b⁻¹ for computation and a / b for reading.

6. exact choose_ratio_one hge h

   The binomial algebra proved above, applied with hge : 1 ≤ n₁ and h : n₁ ≤ n. Both hypotheses of choose_ratio_one are exactly the two facts this branch has.

The joint propensity is the same proof at $S=\{i,j\}$, with two places where the arithmetic is fussier.

theorem completeRandomization_jointPropensity {i j : Fin n} (hij : i ≠ j) :
    (completeRandomization n n₁ h).jointPropensity i j
      = ((n₁ : ℝ) * ((n₁ : ℝ) - 1)) / ((n : ℝ) * ((n : ℝ) - 1)) := by
  have hcardS : #({i, j} : Finset (Fin n)) = 2 := by
    rw [Finset.card_insert_of_notMem (by simpa using hij), Finset.card_singleton]
  have hn2 : 2 ≤ n := by
    have := Finset.card_le_card (Finset.subset_univ ({i, j} : Finset (Fin n)))
    rwa [hcardS, Finset.card_univ, Fintype.card_fin] at this
  rw [Design.jointPropensity_eq_inclusion _ hij, completeRandomization_inclusion]
  rcases Nat.lt_or_ge n₁ 2 with hlt | hge
  ·
    rw [card_filter_subset_powersetCard_of_lt _ _ (by omega)]
    interval_cases n₁ <;> norm_num
  · rw [card_filter_subset_powersetCard _ _ (by omega), hcardS, ← div_eq_mul_inv]
    exact choose_ratio_two hge h

1. have hcardS : #({i, j} : Finset (Fin n)) = 2

   {i, j} is insert i {j}, so Finset.card_insert_of_notMem gives #{j} + 1 provided i ∉ {j} — which is hij : i ≠ j after simpa translates membership in a singleton into equality — and Finset.card_singleton finishes. This is where the hypothesis i ≠ j is spent, and it is why the theorem is false without it: at i = j the set has one element, not two.

2. have hn2 : 2 ≤ n

   There are two distinct units, so there are at least two units. The proof is the same shape as hn in the previous theorem but has to go through the set: Finset.card_le_card (Finset.subset_univ {i, j}) says #{i, j} ≤ #univ, and rewriting with hcardS, Finset.card_univ and Fintype.card_fin turns that into 2 ≤ n.

3. rw [Design.jointPropensity_eq_inclusion _ hij, completeRandomization_inclusion]

   As before, with the pair version of the bridge lemma — jointPropensity_eq_inclusion needs i ≠ j because ∏ k ∈ {i, j}, Z z k has two factors only when the two units are distinct.

4. rcases Nat.lt_or_ge n₁ 2 with hlt | hge

   The counting lemma now needs 2 ≤ n₁, so the split is at 2. Again the goals shown are the two branches after their opening rewrites, the count being 0 on one side and $(\genfrac{}{}{0ex}{}{n-2}{n_1-2})/(\genfrac{}{}{0ex}{}{n}{n_1})$ on the other.

   Goal

   case inl
   hlt : n₁ < 2
   ⊢ ↑0 * (↑(n.choose n₁))⁻¹ = ↑n₁ * (↑n₁ - 1) / (↑n * (↑n - 1))
   
   case inr
   hge : n₁ ≥ 2
   ⊢ ↑((n - 2).choose (n₁ - 2)) / ↑(n.choose n₁) = ↑n₁ * (↑n₁ - 1) / (↑n * (↑n - 1))

5. interval_cases n₁ <;> norm_num

   The degenerate branch, and the reason it is not a one-line simp this time: the right-hand side $n_1(n_1-1)/(n(n-1))$ is zero for two different reasons. hlt : n₁ < 2 bounds n₁ between 0 and 1, so interval_cases enumerates those two values, and norm_num checks each: at n₁ = 0 the factor n₁ is zero, at n₁ = 1 the factor n₁ - 1 is.

6. exact choose_ratio_two hge h

   The real branch, closed by the two-element ratio lemma exactly as the propensity was closed by the one-element one.

The contrast between the two designs is central to design-based variance. They can be made to agree on every propensity — set $p=n_1/n$ — and they never agree on the joint propensities:

Fixing the number of treated units makes units negatively dependent: if $i$ is treated, one of the $n_1$ slots is gone, so $j$ is slightly less likely to be. That inequality is the finite-population correction, and it is the last exercise below.

theorem card_eq_of_prob_ne_zero {z : Assignment n}
    (hz : (completeRandomization n n₁ h).prob z ≠ 0) : #z = n₁ := by
  sorry

theorem card_eq_of_prob_ne_zero {z : Assignment n}
    (hz : (completeRandomization n n₁ h).prob z ≠ 0) : #z = n₁ := by
  by_contra hc
  exact hz (by simp [completeRandomization_prob, hc])

theorem card_filter_subset_powersetCard_of_lt (S : Finset (Fin n)) (k : ℕ) (hk : k < #S) :
    #(((univ : Finset (Fin n)).powersetCard k).filter (S ⊆ ·)) = 0 := by
  sorry

theorem card_filter_subset_powersetCard_of_lt (S : Finset (Fin n)) (k : ℕ) (hk : k < #S) :
    #(((univ : Finset (Fin n)).powersetCard k).filter (S ⊆ ·)) = 0 := by
  rw [Finset.card_eq_zero, Finset.filter_eq_empty_iff]
  intro z hz hSz
  have hcard : #z = k := (Finset.mem_powersetCard.mp hz).2
  exact absurd (hcard ▸ Finset.card_le_card hSz) (by omega)

theorem ex_negative_dependence {n₁ : ℕ} (hle : n₁ ≤ n) (hpos : 0 < n₁) (hlt : n₁ < n)
    {i j : Fin n} (hij : i ≠ j) :
    (completeRandomization n n₁ hle).jointPropensity i j
      < (completeRandomization n n₁ hle).propensity i
        * (completeRandomization n n₁ hle).propensity j := by
  sorry

theorem ex_negative_dependence {n₁ : ℕ} (hle : n₁ ≤ n) (hpos : 0 < n₁) (hlt : n₁ < n)
    {i j : Fin n} (hij : i ≠ j) :
    (completeRandomization n n₁ hle).jointPropensity i j
      < (completeRandomization n n₁ hle).propensity i
        * (completeRandomization n n₁ hle).propensity j := by
  have hn1 : (1 : ℝ) ≤ (n₁ : ℝ) := by exact_mod_cast hpos
  have hn1n : (n₁ : ℝ) < (n : ℝ) := by exact_mod_cast hlt
  have hn2 : (2 : ℝ) ≤ (n : ℝ) := by exact_mod_cast (by omega : 2 ≤ n)
  have hkey : 0 < (n₁ : ℝ) * (n : ℝ) * ((n : ℝ) - (n₁ : ℝ)) :=
    mul_pos (mul_pos (by linarith) (by linarith)) (by linarith)
  rw [completeRandomization_jointPropensity hij, completeRandomization_propensity,
    completeRandomization_propensity, div_mul_div_comm,
    div_lt_div_iff₀ (by nlinarith) (by nlinarith)]
  nlinarith [hkey]

Two designs that exist to break things

A design is any pmf on assignments, and the two axioms rule out very little. It is worth seeing what gets through, because chapter 8’s hypotheses are exactly the fence.

First, a design with impeccable propensities and no usable variation. Flip one fair coin for the whole population; treat everybody, or nobody.

noncomputable def allOrNothing (n : ℕ) (hn : 0 < n) : Design n where
  prob z := if z = ∅ then 1 / 2 else if z = univ then 1 / 2 else 0
  nonneg z := by split_ifs <;> norm_num
  sum_one := by
    have hne : (∅ : Assignment n) ≠ univ := by
      intro h
      have hmem : (⟨0, hn⟩ : Fin n) ∈ (∅ : Finset (Fin n)) := h ▸ Finset.mem_univ _
      simp at hmem
    have key : ∀ z : Assignment n,
        (if z = ∅ then (1 : ℝ) / 2 else if z = univ then 1 / 2 else 0)
          = (if z = ∅ then (1 : ℝ) / 2 else 0) + (if z = univ then (1 : ℝ) / 2 else 0) := by
      intro z
      rcases eq_or_ne z ∅ with h1 | h1
      · have h2 : z ≠ univ := fun hc => hne (h1.symm.trans hc)
        simp only [ite_eq_left h1, ite_eq_right h2, add_zero]
      · rcases eq_or_ne z univ with h2 | h2
        · simp only [ite_eq_right h1, ite_eq_left h2, zero_add]
        · simp only [ite_eq_right h1, ite_eq_right h2, zero_add]
    rw [Finset.sum_congr rfl fun z _ => key z, Finset.sum_add_distrib,
      Finset.sum_ite_eq' univ (∅ : Assignment n) fun _ => (1 : ℝ) / 2,
      Finset.sum_ite_eq' univ (univ : Assignment n) fun _ => (1 : ℝ) / 2]
    norm_num

@[simp] theorem allOrNothing_prob {n : ℕ} {hn : 0 < n} (z : Assignment n) :
    (allOrNothing n hn).prob z = if z = ∅ then 1 / 2 else if z = univ then 1 / 2 else 0 := rfl

1. nonneg z := by split_ifs <;> norm_num

   split_ifs splits on every condition in the goal at once — there are two nested ifs here — leaving three goals, 0 ≤ 1 / 2 twice and 0 ≤ 0, all numeric.

2. have hne : (∅ : Assignment n) ≠ univ

   The sum_one field, and the first thing it needs: the two assignments carrying probability are different assignments. This is where hn : 0 < n earns its place in the definition — with no units, ∅ and univ are the same set, the two halves of the pmf would land on the same assignment and the probabilities would sum to 1 / 2. The proof takes the unit ⟨0, hn⟩, which lies in univ, and derives a contradiction from its lying in ∅.

3. have key : ∀ z, (if z = ∅ then 1 / 2 else if z = univ then 1 / 2 else 0) = (if z = ∅ then 1 / 2 else 0) + (if z = univ then 1 / 2 else 0)

   Nested ifs do not split under a sum; a sum of two single ifs does. So the pmf is rewritten as one term per atom. The proof is three cases — z = ∅, z = univ, neither — and in the first hne is what rules out the second if, which is exactly the step that would fail if the two assignments coincided.

4. rw [Finset.sum_congr rfl fun z _ => key z, Finset.sum_add_distrib, Finset.sum_ite_eq' univ (∅ : Assignment n) fun _ => (1 : ℝ) / 2, Finset.sum_ite_eq' univ (univ : Assignment n) fun _ => (1 : ℝ) / 2]

   sum_congr applies key under the sum, sum_add_distrib splits it into two, and then Finset.sum_ite_eq' s a b : (∑ x ∈ s, if x = a then b x else 0) = if a ∈ s then b a else 0 collapses each: a sum with one surviving term is that term. Both conditions ∅ ∈ univ and univ ∈ univ are true, so norm_num finishes with 1 / 2 + 1 / 2 = 1.

Only two assignments carry probability, so every expectation collapses to an average of two numbers.

Lemma. Under the all-or-nothing design, every statistic $f$ has $\mathbb{E}[f]=\frac{1}{2}f(\mathrm{\varnothing })+\frac{1}{2}f(\{1,\dots ,n\})$. Consequently ${\pi }_i=1/2$ for every unit, ${\pi }_{ij}=1/2$ for every pair, and no two units are independent.

theorem allOrNothing_expect (f : Assignment n → ℝ) :
    (allOrNothing n hn).expect f = f ∅ / 2 + f univ / 2 := by
  have hne : (∅ : Assignment n) ≠ univ := by
    intro h
    have hmem : (⟨0, hn⟩ : Fin n) ∈ (∅ : Finset (Fin n)) := h ▸ Finset.mem_univ _
    simp at hmem
  have key : ∀ z : Assignment n,
      (allOrNothing n hn).prob z * f z
        = (if z = ∅ then f z / 2 else 0) + (if z = univ then f z / 2 else 0) := by
    intro z
    rw [allOrNothing_prob]
    rcases eq_or_ne z ∅ with h1 | h1
    · have h2 : z ≠ univ := fun hc => hne (h1.symm.trans hc)
      simp only [ite_eq_left h1, ite_eq_right h2, add_zero]
      ring
    · rcases eq_or_ne z univ with h2 | h2
      · simp only [ite_eq_right h1, ite_eq_left h2, zero_add]
        ring
      · simp only [ite_eq_right h1, ite_eq_right h2, zero_add, zero_mul]
  rw [Design.expect, Finset.sum_congr rfl fun z _ => key z, Finset.sum_add_distrib,
    Finset.sum_ite_eq' univ (∅ : Assignment n) fun z => f z / 2,
    Finset.sum_ite_eq' univ (univ : Assignment n) fun z => f z / 2]
  norm_num

1. have key : ∀ z, (allOrNothing n hn).prob z * f z = (if z = ∅ then f z / 2 else 0) + (if z = univ then f z / 2 else 0)

   The same decomposition as in sum_one, now with f z carried along: three cases, hne again ruling out the second if in the first of them, and a ring in each of the two live branches to turn 1 / 2 * f z into f z / 2. The goal shown is the main goal after rw [Design.expect] and sum_congr have applied it under the sum.

   Goal

   ⊢ ∑ z, ((if z = ∅ then f z / 2 else 0) + if z = univ then f z / 2 else 0) = f ∅ / 2 + f univ / 2

2. rw [Finset.sum_add_distrib]

   One sum of two ifs becomes two sums of one if each — the shape Finset.sum_ite_eq' recognizes.

   Goal

   ⊢ ((∑ x, if x = ∅ then f x / 2 else 0) + ∑ x, if x = univ then f x / 2 else 0) = f ∅ / 2 + f univ / 2

3. rw [Finset.sum_ite_eq' univ (∅ : Assignment n) fun z => f z / 2, Finset.sum_ite_eq' univ (univ : Assignment n) fun z => f z / 2]

   Each sum collapses to its one surviving term, guarded by a membership condition that is trivially true.

   Goal

   ⊢ ((if ∅ ∈ univ then f ∅ / 2 else 0) + if univ ∈ univ then f univ / 2 else 0) = f ∅ / 2 + f univ / 2

4. norm_num

   ∅ ∈ univ and univ ∈ univ are both true, so both ifs take their first branch and the goal is an identity.

theorem allOrNothing_propensity (i : Fin n) : (allOrNothing n hn).propensity i = 1 / 2 := by
  rw [Design.propensity, allOrNothing_expect,
    Z_of_not_mem (by simp : i ∉ (∅ : Assignment n)), Z_of_mem (Finset.mem_univ i)]
  norm_num

theorem allOrNothing_jointPropensity (i j : Fin n) :
    (allOrNothing n hn).jointPropensity i j = 1 / 2 := by
  rw [Design.jointPropensity, allOrNothing_expect,
    Z_of_not_mem (by simp : i ∉ (∅ : Assignment n)), Z_of_mem (Finset.mem_univ i),
    Z_of_mem (Finset.mem_univ j)]
  norm_num

theorem allOrNothing_not_independent (i j : Fin n) :
    (allOrNothing n hn).jointPropensity i j
      ≠ (allOrNothing n hn).propensity i * (allOrNothing n hn).propensity j := by
  rw [allOrNothing_jointPropensity, allOrNothing_propensity, allOrNothing_propensity]
  norm_num

1. rw [Design.propensity, allOrNothing_expect, Z_of_not_mem (by simp : i ∉ (∅ : Assignment n)), Z_of_mem (Finset.mem_univ i)]

   Unfold the propensity to an expectation, apply the lemma just proved to get Z ∅ i / 2 + Z univ i / 2, and then evaluate the two indicators with the two lemmas from chapter 6: nobody is in ∅, everybody is in univ. norm_num then computes 0 / 2 + 1 / 2.

2. rw [Design.jointPropensity, allOrNothing_expect, Z_of_not_mem (by simp : i ∉ (∅ : Assignment n)), Z_of_mem (Finset.mem_univ i), Z_of_mem (Finset.mem_univ j)]

   The joint propensity is the same chain with one more indicator. The statistic is Z z i * Z z j, so the term at ∅ is Z ∅ i * Z ∅ j; only the first factor is rewritten, to 0, and the second is never evaluated at all — norm_num discards the product because of the zero in front of it. At univ both factors become 1, and 1 * 1 / 2 = 1 / 2.

3. rw [allOrNothing_jointPropensity, allOrNothing_propensity, allOrNothing_propensity]

   Non-independence, and it is a calculation rather than an argument: rewrite all three quantities to their values and let norm_num observe that 1 / 2 ≠ 1 / 2 * (1 / 2). The theorem holds for every pair of units, including i = j — where it is not a failure of independence but the ordinary fact that $\mathbb{E}[Z_i^2]\ne \mathbb{E}[Z_i{]}^2$ unless the variance is zero.

Every unit has ${\pi }_i=1/2$, exactly as under a fair Bernoulli design, so any result that depends only on the propensities applies verbatim. Chapter 8’s theorem is one of them, and the corollary is two lines — the hypotheses of HT_unbiased are and , and $1/2$ satisfies both:

theorem allOrNothing_HT_unbiased (P : Population n) :
    (allOrNothing n hn).expect (HT (allOrNothing n hn) P) = P.tau :=
  HT_unbiased _ P
    (fun i => by rw [allOrNothing_propensity]; norm_num)
    (fun i => by rw [allOrNothing_propensity]; norm_num)

1. HT_unbiased _ P (fun i => …) (fun i => …)

   A term-mode proof: the general theorem, applied. Its first _ is the design, which Lean reads off the goal; then the population; then its two hypotheses, each a function from a unit to a proof about that unit’s propensity.

2. fun i => by rw [allOrNothing_propensity]; norm_num

   Each hypothesis is discharged by the same two-tactic script inside a by block: rewrite (allOrNothing n hn).propensity i to 1 / 2, then let norm_num check 0 < 1 / 2 in the first case and 1 / 2 < 1 in the second. Reading a term-mode proof means reading its arguments as the obligations they discharge; here the obligations are the two sides of “the propensity is strictly between zero and one”.

And ${\pi }_{ij}=1/2\ne 1/4={\pi }_i{\pi }_j$ for every pair, so no two units are independent, and the estimator is unbiased for a quantity it can never estimate: the treated mean and the control mean are never both observed. Unbiasedness is a statement about ${\pi }_i$; usefulness is a statement about ${\pi }_{ij}$.

Second, the degenerate case. A point mass is a design:

noncomputable def pointMass (z₀ : Assignment n) : Design n where
  prob z := if z = z₀ then 1 else 0
  nonneg z := by split <;> norm_num
  sum_one := by simp

theorem pointMass_expect (z₀ : Assignment n) (f : Assignment n → ℝ) :
    (pointMass z₀).expect f = f z₀ := by
  show ∑ z, (if z = z₀ then (1 : ℝ) else 0) * f z = f z₀
  simp

theorem pointMass_propensity (z₀ : Assignment n) (i : Fin n) :
    (pointMass z₀).propensity i = Z z₀ i :=
  pointMass_expect z₀ _

theorem pointMass_propensity_not_interior (z₀ : Assignment n) (i : Fin n) :
    ¬(0 < (pointMass z₀).propensity i ∧ (pointMass z₀).propensity i < 1) := by
  rw [pointMass_propensity]
  unfold Z
  split <;> rintro ⟨h0, h1⟩ <;> linarith

1. nonneg z := by split <;> norm_num

   The first axiom, as for oneOfThree: split the if, and both branches are numeric.

2. sum_one := by simp

   The pmf is if z = z₀ then 1 else 0, and a sum of that shape over univ has one surviving term, which is 1. simp knows the lemma (Finset.sum_ite_eq' again) without being told.

3. show ∑ z, (if z = z₀ then (1 : ℝ) else 0) * f z = f z₀

   pointMass_expect. show restates the goal in a definitionally equal form — here unfolding expect and prob by hand — because simp needs to see the if before it can pick out the surviving term.

4. simp

   And it does: one term of the sum survives, with the factor 1 in front of f z₀.

5. pointMass_expect z₀ _

   pointMass_propensity is that lemma at the statistic fun z => Z z i, and the proof is the term itself: the propensity of i under a point mass at z₀ is Z z₀ i, one if i ∈ z₀ and zero otherwise. No design is more informative about its own propensities and none is less useful.

6. rw [pointMass_propensity]; unfold Z; split <;> rintro ⟨h0, h1⟩ <;> linarith

   pointMass_propensity_not_interior, the theorem that the hypotheses of HT_unbiased fail. After rewriting and unfolding, split gives the two values of the indicator; rintro ⟨h0, h1⟩ then assumes the conjunction 0 < π ∧ π < 1 and destructures it in one move — the goal is a negation, so assuming the statement and deriving a contradiction is the proof — and linarith finds the contradiction: 0 < 0 in one branch, 1 < 1 in the other.

Every propensity is $0$ or $1$, so for every unit one of the two hypotheses of HT_unbiased fails — and the failure is not a technicality. Take the point mass at “everybody treated”, where all of them are $1$, and compute.

Proposition. Let $D$ be the point mass at “everybody treated”. Then $\mathbb{E}[{\hat{\tau }}_{\mathrm{HT}}]={\overline{y}}_1$, so the bias $\mathbb{E}[{\hat{\tau }}_{\mathrm{HT}}]-\tau$ equals ${\overline{y}}_0$.

theorem pointMass_univ_propensity (i : Fin n) :
    (pointMass (univ : Assignment n)).propensity i = 1 := by
  rw [pointMass_propensity, Z_of_mem (Finset.mem_univ i)]

theorem pointMass_univ_expect_HT (P : Population n) :
    (pointMass (univ : Assignment n)).expect (HT (pointMass univ) P) = mean P.y1 := by
  rw [pointMass_expect]
  have hZ : ∀ i : Fin n, Z (univ : Assignment n) i = 1 := fun i => Z_of_mem (Finset.mem_univ i)
  simp only [HT, pointMass_univ_propensity, hZ, one_mul, div_one, sub_self, zero_mul,
    zero_div, sub_zero]
  rw [mean]
  ring

theorem pointMass_univ_HT_bias (P : Population n) :
    (pointMass (univ : Assignment n)).expect (HT (pointMass univ) P) - P.tau = mean P.y0 := by
  rw [pointMass_univ_expect_HT, P.tau_eq_mean_sub_mean]
  ring

1. rw [pointMass_propensity, Z_of_mem (Finset.mem_univ i)]

   pointMass_univ_propensity: the propensity is Z univ i, and everybody belongs to univ, so it is 1. This is the hypothesis of HT_unbiased failing, stated as an equation.

2. rw [pointMass_expect]

   The expectation of anything under a point mass is its value at the one assignment. From here on there is no probability left in the problem — what remains is arithmetic in the definition of the estimator.

   Goal

   P : Population n
   ⊢ HT (pointMass univ) P univ = mean P.y1

3. have hZ : ∀ i : Fin n, Z (univ : Assignment n) i = 1 := fun i => Z_of_mem (Finset.mem_univ i)

   Everybody is treated, recorded as a fact about every unit so that the simp only below can use it under the binder of the estimator’s sum. Z_of_mem is the chapter-6 lemma; Finset.mem_univ i is the proof that unit i belongs to the full set.

4. simp only [HT, pointMass_univ_propensity, hZ, one_mul, div_one, sub_self, zero_mul, zero_div, sub_zero]

   The estimator, evaluated. HT unfolds it; pointMass_univ_propensity and hZ (the local fact that every indicator is 1) put 1 everywhere; then one_mul and div_one clear the treated term down to P.y1 i, and the control term dies in three steps — sub_self makes both 1 - Z z i and 1 - π i into 0, zero_mul and zero_div make the term 0, and sub_zero drops it. The zero_div is the division convention doing its work in plain sight: it is what turns $0/0$ into $0$ rather than into an error.

   Goal

   hZ : ∀ (i : Fin n), Z univ i = 1
   ⊢ 1 / ↑n * ∑ x, P.y1 x = mean P.y1

5. rw [mean]

   The goal is 1 / ↑n * ∑ x, P.y1 x = mean P.y1. Unfolding mean on the right puts both sides in the same vocabulary, 1 / n times a sum against a sum over n.

6. ring

   (1/n) · s = s/n is an identity in a field, with the sum as an atom.

7. rw [pointMass_univ_expect_HT, P.tau_eq_mean_sub_mean]

   The bias theorem. The expectation just computed is ${\overline{y}}_1$, and $\tau ={\overline{y}}_1-{\overline{y}}_0$ is the identity from chapter 6, so the goal becomes mean P.y1 - (mean P.y1 - mean P.y0) = mean P.y0, and the closing ring does the cancellation.

Note what the division convention did: it turned a statement that should be undefined into a false-but-well-typed one. That is the right trade — every theorem stays an unconditional equation — but it means a hypothesis such as is necessary, and its absence is easy to miss.

theorem ex_expect_affine (D : Design n) (f : Assignment n → ℝ) (a b : ℝ) :
    D.expect (fun z => a * f z + b) = a * D.expect f + b := by
  sorry

theorem ex_expect_affine (D : Design n) (f : Assignment n → ℝ) (a b : ℝ) :
    D.expect (fun z => a * f z + b) = a * D.expect f + b := by
  rw [D.expect_add (fun z => a * f z) (fun _ => b), D.expect_const_mul f a, D.expect_const b]

How a probabilist would set this up

They would begin with a probability space $(\mathrm{\Omega },\mathcal{F},\mathbb{P})$, define $Z:\mathrm{\Omega }\to \{0,1{\}}^n$ as a measurable map, and define $\mathbb{E}[f]$ as a Lebesgue integral. In Mathlib that is PMF (Assignment n) with values in ℝ≥0∞, its induced Measure, and ∫ z, f z ∂μ, with Integrable and Measurable side conditions on every lemma.

For a finite population, all of that machinery is doing nothing:

• the σ-algebra is $\mathcal{P}(\mathrm{\Omega })$, so every function is measurable;
• $\mathrm{\Omega }$ is finite, so every function is integrable;
• the integral is the finite sum ∑ z, D.prob z * f z.