Designs — MrCLean

The population is fixed. The potential outcomes are fixed. The one thing that is random is which units get treated, and the experimenter chose the mechanism that makes it random. That mechanism is the design, and in a finite population it is nothing more exotic than a probability mass function on the $2^{n}$ possible treated sets.

This chapter is where probability enters the tutorial, and it enters as a finite sum. Everything after it — unbiasedness, variance, blocking, randomization tests — is algebra performed against the object defined here.

A design is its pmf

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

DefinitionDesign

A design for $n$ units is a function $P$ from treated sets to $[0, 1]$ with $\sum_{z} P (z) = 1$ .

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

/-- A **design** (randomisation scheme) for a population of `n` units: a probability mass
function on the finite set of possible assignments.

* `prob z` is the probability that the experimenter assigns exactly the set `z` to treatment;
* `nonneg` and `sum_one` are the two Kolmogorov axioms, which in a finite setting is all
  there is to say. -/
structure Design (n : ℕ) where
  /-- The probability of drawing the assignment `z`. -/
  prob : Assignment n → ℝ
  /-- Probabilities are nonnegative. -/
  nonneg : ∀ z, 0 ≤ prob z
  /-- Probabilities sum to one over all `2 ^ n` assignments. -/
  sum_one : ∑ z, prob z = 1

Three things are worth saying about that block, because they are the shape of every definition in 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.¹

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

ExpectationNew tabOpens with 262 lines of library preamble — the code above is at the bottom of the editor.

/-- The **expectation** of a statistic `f` of the assignment, `E[f] = ∑ z, P(z) · f z`.

`f` is a function of the assignment alone.  Everything we will ever take an expectation of --
an estimator, an indicator, a squared deviation -- is such a function, because the assignment
is the only random object in sight. -/
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:

Linearity of expectationNew tabOpens with 287 lines of library preamble — the code above is at the bottom of the editor.

/-- **Linearity of expectation**, additive part: `E[f + g] = E[f] + 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 earns its keep in exactly one lemma:

The expectation of a constantNew tabOpens with 318 lines of library preamble — the code above is at the bottom of the editor.

/-- The expectation of a constant is that constant.  This is exactly where `sum_one` is used. -/
theorem expect_const (c : ℝ) : D.expect (fun _ => c) = c := by
  simp only [expect, ← Finset.sum_mul, D.sum_one, one_mul]

Read the tactic list right to left: ← 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 it is used backwards), 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 $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:

Expectation commutes with a finite sumNew tabOpens with 322 lines of library preamble — the code above is at the bottom of the editor.

/-- Expectation commutes with a finite sum of statistics:
`E[∑ i ∈ s, f i] = ∑ i ∈ s, E[f i]`.

This is the form of linearity that unbiasedness proofs actually use, since every estimator we
study is a sum over units. -/
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

The proof is Finset.sum_comm: exchanging the order of a double finite sum. That exchange is the only “theorem” in the whole of our probability theory, and in a finite population it is free.

ExerciseScalar linearity

Warm-up. Pull a constant out of an expectation. The proof is the same shape as expect_add: unfold expect, then use Finset.mul_sum to move c out of the sum. Finishing with Finset.sum_congr rfl fun z _ => by ring is the honest way to reassociate D.prob z * (c * f z) into c * (D.prob z * f z) inside the binder.

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

/-- **Linearity of expectation**, scalar part: `E[c · f] = c · E[f]`. -/
theorem expect_const_mul (c : ℝ) : D.expect (fun z => c * f z) = c * D.expect f := by
  sorry

Show solution

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

/-- **Linearity of expectation**, scalar part: `E[c · f] = c · E[f]`. -/
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

ExerciseStatistics that agree on the support

Core. Two statistics that agree wherever the design puts positive probability have the same expectation. This is the lemma that lets chapter 8 replace difference-in-means by Horvitz–Thompson on the assignments complete randomization can actually produce. Reduce to a termwise statement with Finset.sum_congr rfl, then by_cases hz : D.prob z = 0: in the zero branch both sides are 0 * _.

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

/-- Two statistics that agree on the **support** of the design have the same expectation.

This is how we transfer a computation from "all assignments" to "the assignments the design
can actually produce" -- for instance, difference-in-means only behaves well on assignments
with the right number of treated units, and under complete randomisation those are exactly
the assignments of positive probability. -/
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

Show solution

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

/-- Two statistics that agree on the **support** of the design have the same expectation.

This is how we transfer a computation from "all assignments" to "the assignments the design
can actually produce" -- for instance, difference-in-means only behaves well on assignments
with the right number of treated units, and under complete randomisation those are exactly
the assignments of positive probability. -/
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:

The probability of an eventNew tabOpens with 342 lines of library preamble — the code above is at the bottom of the editor.

/-- The **probability of an event**, i.e. of a set of assignments singled out by a predicate. -/
noncomputable def probOf (p : Assignment n → Prop) [DecidablePred p] : ℝ :=
  ∑ z ∈ univ.filter p, D.prob z

Expectation of an indicatorNew tabOpens with 346 lines of library preamble — the code above is at the bottom of the editor.

/-- **The expectation of an indicator is a probability**, `E[1{p}] = P(p)`.

This is the bridge between the algebraic world (indicators are numbers, so we may add and
multiply them) and the probabilistic world (events have probabilities). -/
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. Finset.sum_filter is the lemma that states this for arbitrary sums, and the proof is one rewrite in the other direction.

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.

The propensity scoreNew tabOpens with 365 lines of library preamble — the code above is at the bottom of the editor.

/-- The **propensity score** of unit `i`: `π i = E[Z i] = P(unit i is treated)`.

Under a design-based view the propensity score is not a nuisance function to be estimated; it
is a known number determined by the experimenter's randomisation device. -/
noncomputable def propensity (i : Fin n) : ℝ := D.expect (fun z => Z z i)

The joint propensityNew tabOpens with 371 lines of library preamble — the code above is at the bottom of the editor.

/-- The **joint propensity** of units `i` and `j`: `π i j = P(both i and j are treated)`.

Unlike the individual propensities, this quantity distinguishes designs that randomise units
independently (Bernoulli) from designs that constrain the number of treated units
(complete randomisation), and it is what drives the variance of every estimator. -/
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 $π_{i} = E [Z_{i}] = P (i treated)$ is not a nuisance function to be estimated from covariates. It is a number the experimenter chose. The joint propensity $π_{i j} = P (i and j both treated)$ is what distinguishes designs that agree on every $π_{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:

Inclusion probabilityNew tabOpens with 357 lines of library preamble — the code above is at the bottom of the editor.

/-- The **inclusion probability** of a set `S` of units: the probability that *every* unit in
`S` is treated, written as the expectation of a product of indicators.

Specialising `S` to a singleton gives the propensity score, and to a pair gives the joint
propensity; keeping the general version means each design only needs one computation. -/
noncomputable def inclusion (S : Finset (Fin n)) : ℝ :=
  D.expect (fun z => ∏ i ∈ S, Z z i)

$inclusion (S) = 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:

π i is the inclusion probability of {i}New tabOpens with 378 lines of library preamble — the code above is at the bottom of the editor.

/-- The propensity score is the inclusion probability of the singleton `{i}`. -/
theorem propensity_eq_inclusion (i : Fin n) : D.propensity i = D.inclusion {i} := by
  simp only [propensity, inclusion, Finset.prod_singleton]

The propensity is also a probability in the ordinary sense, which is a one-line consequence of expect_indicator:

π i = P(i ∈ z)New tabOpens with 389 lines of library preamble — the code above is at the bottom of the editor.

/-- The propensity score is the probability of the event "unit `i` is treated". -/
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)

ExerciseThe control indicator

Warm-up. The expected value of the control indicator is $1 - π_{i}$ . Subtract using expect_sub with the constant statistic fun _ => 1, then rewrite the two pieces with expect_const and expect_Z. simpa only [...] using h applies a hypothesis after simplifying it, which is what the one-line proof does.

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

/-- The expectation of the control indicator is `1 - π i`. -/
@[simp] theorem expect_one_sub_Z (i : Fin n) :
    D.expect (fun z => 1 - Z z i) = 1 - D.propensity i := by
  sorry

Show solution

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

/-- The expectation of the control indicator is `1 - π i`. -/
@[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

ExercisePropensities sum to the expected sample size

Core. $\sum_{i} π_{i} = E [# z]$ : the propensities always sum to the expected number of treated units, whatever the design. Use expect_sum backwards to turn the sum of expectations into the expectation of a sum, then expect_congr with sum_Z : ∑ i, Z z i = #z to identify the summand.

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

/-- The propensity scores sum to the expected number of treated units. -/
theorem sum_propensity : ∑ i, D.propensity i = D.expect (fun z => (#z : ℝ)) := by
  sorry

Show solution

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

/-- The propensity scores sum to the expected number of treated units. -/
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

Variance under the designNew tabOpens with 426 lines of library preamble — the code above is at the bottom of the editor.

/-- The **variance** of a statistic under the design, `Var f = E[(f - E f)²]`. -/
noncomputable def var : ℝ := D.expect (fun z => (f z - D.expect f) ^ 2)

The shortcut formulaNew tabOpens with 433 lines of library preamble — the code above is at the bottom of the editor.

/-- **Covariance shortcut formula**: `Cov(f, g) = E[f·g] - E[f]·E[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

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

The shortcut formula is proved by expanding the product termwise with ring inside a Finset.sum_congr, splitting the sum, pulling the constants $E [f]$ and $E [g]$ out, and using sum_one once more. 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.

ExerciseVariance of a treatment indicator

Core. $Var (Z_{i}) = π_{i} (1 - π_{i})$ , for every design. Start from Design.var_eq_expect_sq_sub_sq; the squared indicator is the indicator (Z_sq), so expect_congr turns the first term into a propensity, and ring finishes.

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

/-- Exercise (core).  **The variance of a treatment indicator is `π (1 - π)`**, for every
design whatsoever.

Start from `Design.var_eq_expect_sq_sub_sq`.  The squared indicator is the indicator
(`Z_sq`), so `Design.expect_congr` turns the first expectation into a propensity
(`Design.expect_Z`); then `ring`. -/
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

Show solution

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

/-- Exercise (core).  **The variance of a treatment indicator is `π (1 - π)`**, for every
design whatsoever.

Start from `Design.var_eq_expect_sq_sub_sq`.  The squared indicator is the indicator
(`Z_sq`), so `Design.expect_congr` turns the first expectation into a propensity
(`Design.expect_Z`); then `ring`. -/
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

ExerciseCovariance of two treatment indicators

Core. $Cov (Z_{i}, Z_{j}) = π_{i j} - π_{i} π_{j}$ . This is the identity that makes the joint propensity the interesting summary: it is zero exactly when the two units are uncorrelated. Use Design.cov_eq_expect_mul_sub_mul and Design.expect_Z; the remaining expectation is Design.jointPropensity by definition, so rfl closes it.

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

/-- Exercise (core).  **The covariance of two treatment indicators** is the joint propensity
minus the product of the propensities.  Under Bernoulli it is `p² - p² = 0`; under complete
randomisation it is negative, which is the finite-population correction.

`Design.cov_eq_expect_mul_sub_mul` is the shortcut formula, `Design.expect_Z` handles the two
marginal terms, and the surviving expectation is `Design.jointPropensity` by definition, so
`rfl` closes the goal. -/
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

Show solution

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

/-- Exercise (core).  **The covariance of two treatment indicators** is the joint propensity
minus the product of the propensities.  Under Bernoulli it is `p² - p² = 0`; under complete
randomisation it is negative, which is the finite-population correction.

`Design.cov_eq_expect_mul_sub_mul` is the shortcut formula, `Design.expect_Z` handles the two
marginal terms, and the surviving expectation is `Design.jointPropensity` by definition, so
`rfl` closes the goal. -/
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.

All eight assignments of a three-unit populationNew tabOpens with 1,028 lines of library preamble — the code above is at the bottom of the editor.

/- With `n = 3` there are `2 ^ 3 = 8` assignments, and Lean will list them for you.
`Assignment 3` is `Finset (Fin 3)`, a finite type, so `Finset.univ` is the set of *all*
treated sets and `decide` settles any decidable question about it. -/
#eval (univ : Finset (Assignment 3))
-- {∅, {0}, {1}, {0, 1}, {2}, {0, 2}, {1, 2}, {0, 1, 2}}

/-- A three-unit population admits eight assignments. -/
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:

Every expectation on three units is an eight-term sumNew tabOpens with 1,037 lines of library preamble — the code above is at the bottom of the editor.

/-- **Every expectation over a three-unit design is an eight-term sum.**

This is the lemma that turns "take an expectation" into arithmetic: rewrite `univ` as the
explicit list of the eight treated sets (`decide` checks that the two `Finset`s are equal),
peel the sum apart with `Finset.sum_insert` — whose side condition "the head is not in the
tail" is again decidable — and finish with `ring`, which only has to reassociate. -/
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

The show … from by decide is a proof of a Finset equality supplied inline: Finset.univ and the explicit list denote the same set, and decide checks it. Finset.sum_insert then peels the head off the sum, once per element, with the side condition “the head is not in the tail” — again decidable, hence (by decide). Only the final ring is real arithmetic, and all it does is reassociate.

Now a real design, written out as a table:

Complete randomization on three units, by handNew tabOpens with 1,052 lines of library preamble — the code above is at the bottom of the editor.

/-- **Complete randomisation on three units with one treated**, written out as a pmf: three
assignments of probability `1/3`, five of probability `0`.

The two proof fields are the whole content of "this is a probability distribution":
`nonneg` is a case split on the `if`, and `sum_one` is the eight-term sum, evaluated once the
cardinalities `#{0, 2} = 2` and friends have been computed. -/
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. nonneg is a case split (split on the if), and sum_one is the eight-term sum, evaluated by norm_num after the combinatorial facts (#{0,2} = 1 is false, and so on) have been settled. Those cardinalities need two lemmas in the norm_num list: Finset.card_insert_of_notMem, to peel an element off a literal set, and Fin.ext_iff, which reduces (0 : Fin 3) ≠ 2 to 0 ≠ 2 in ℕ — without it norm_num cannot discharge the side condition and the sum stops half-evaluated. Then the propensity is a computation rather than a theorem:

π = 1/3 by brute force, and by the general theoremNew tabOpens with 1,068 lines of library preamble — the code above is at the bottom of the editor.

/-- The propensity of every unit is `1/3`, computed by brute force: expand the expectation
into eight terms, split into the three units, and evaluate twenty-four indicators. -/
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]

/-- The same design, obtained from the library rather than by hand: the pmf of
`completeRandomization 3 1` is `1 / C(3,1) = 1/3` on the singletons, so
`completeRandomization_propensity` gives the same answer with no enumeration at all. -/
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

fin_cases i splits into the three units, and each branch is eight indicator evaluations. That 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 second lemma checks the tiny design against the library’s: completeRandomization 3 1 has pmf $1 / (\binom{3}{1}) = 1 / 3$ on the singletons, and the general propensity theorem then says $π_{i} = 1 / 3$ without any enumeration at all.

The Bernoulli design

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

The Bernoulli designNew tabOpens with 562 lines of library preamble — the code above is at the bottom of the editor.

/-- The **Bernoulli design** with treatment probability `p`: each unit is independently
assigned to treatment with probability `p`.

The number of treated units is itself random, which is the defining difference from complete
randomisation. -/
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

The sum_one field is the binomial theorem, and it is the $S = \emptyset$ case of the one lemma this design needs:

The master computationNew tabOpens with 510 lines of library preamble — the code above is at the bottom of the editor.

/-- The key combinatorial computation behind the Bernoulli design.

Summing the Bernoulli weight of `z` against the indicator that `S` is entirely treated gives
`p ^ #S`.  Taking `S = ∅` this says the weights sum to one; taking `S = {i}` it gives the
propensity; taking `S = {i, j}` it gives the joint propensity.  Everything about Bernoulli
first- and second-order inclusion probabilities is this one lemma.

The proof chooses `g i = if i ∈ S then 0 else (1 - p)` and applies `Finset.prod_add` to
`∏ i, (p + g i)`.  The factor `0` kills exactly the assignments that fail to treat all of
`S`, so the surviving terms are the ones we want to sum. -/
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
  -- The left-hand side of `prod_add` collapses to `p ^ #S`.
  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]
  -- Each summand on the right-hand side of `prod_add` is a Bernoulli weight times an indicator.
  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
    · -- all of `S` is treated: no zero factor appears
      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]
    · -- some unit of `S` is untreated: it contributes the factor `0`
      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:

\prod_{i} (a_{i} + b_{i}) = \sum_{t \subseteq {1, \dots, n}} (\prod_{i \in t} a_{i}) (\prod_{i \notin t} b_{i}) .

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

b_{i} = {\begin{cases} 0 & i \in S \\ 1 - p & i \notin S, \end{cases}

and the right-hand side becomes $\sum_{z} p^{# z} (1 - p)^{n - # z} 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 = \emptyset$ gives the pmf summing to one; that is why the design’s sum_one field is a one-liner citing this lemma.

Everything else about the Bernoulli design is a corollary:

π i = pNew tabOpens with 586 lines of library preamble — the code above is at the bottom of the editor.

/-- **Every unit has propensity `p`** under the Bernoulli design. -/
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]

ExerciseIndependence: π ij = p²

Core. Distinct units are independent. Rewrite with Design.jointPropensity_eq_inclusion (which needs i ≠ j) and bernoulli_inclusion, and then compute $# {i, j} = 2$ with Finset.card_insert_of_notMem, whose side condition is i ∉ {j} — simpa using hij discharges it.

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

/-- **Distinct units are independent** under the Bernoulli design: the probability that both
`i` and `j` are treated is `p²`, the product of their propensities. -/
theorem bernoulli_jointPropensity {i j : Fin n} (hij : i ≠ j) :
    (bernoulliDesign n p hp0 hp1).jointPropensity i j = p ^ 2 := by
  sorry

Show solution

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

/-- **Distinct units are independent** under the Bernoulli design: the probability that both
`i` and `j` are treated is `p²`, the product of their propensities. -/
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]

ExerciseExpected number of treated units

Core. $E [# z] = n p$ under the Bernoulli design — the number treated is itself random, which is the design’s defining feature and, in chapter 8, its defining problem. Use Design.sum_propensity backwards, rewrite every propensity with bernoulli_propensity, and count the constant summands with Finset.sum_const, Finset.card_univ, Fintype.card_fin and nsmul_eq_mul.

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

/-- Exercise (core).  **The expected number of treated units under a Bernoulli design is
`n p`.**  The count is random, which is exactly what complete randomisation refuses to allow.

`Design.sum_propensity` says `∑ i, π i = E[#z]`; use it backwards (`←`), rewrite every
propensity with `bernoulli_propensity`, and count the constant summands with
`Finset.sum_const`, `Finset.card_univ`, `Fintype.card_fin` and `nsmul_eq_mul`. -/
theorem ex_bernoulli_expected_treated :
    (bernoulliDesign n p hp0 hp1).expect (fun z => (#z : ℝ)) = n * p := by
  sorry

Show solution

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

/-- Exercise (core).  **The expected number of treated units under a Bernoulli design is
`n p`.**  The count is random, which is exactly what complete randomisation refuses to allow.

`Design.sum_propensity` says `∑ i, π i = E[#z]`; use it backwards (`←`), rewrite every
propensity with `bernoulli_propensity`, and count the constant summands with
`Finset.sum_const`, `Finset.card_univ`, `Fintype.card_fin` and `nsmul_eq_mul`. -/
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 $(\binom{n}{n_{1}})$ sets of that size.

Complete randomizationNew tabOpens with 587 lines of library preamble — the code above is at the bottom of the editor.

/-- The **completely randomised design** with `n₁` treated units: the uniform distribution on
the `C(n, n₁)` assignments that treat exactly `n₁` units.

`prob z` is `1 / C(n, n₁)` when `#z = n₁` and `0` otherwise.  The hypothesis `n₁ ≤ n` is what
makes `C(n, n₁)` nonzero, hence the pmf well defined. -/
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 / (\binom{n}{n_{1}})$ on the sets of size $n_{1}$ and $0$ elsewhere, and sum_one is a counting argument: there are $(\binom{n}{n_{1}})$ such sets (Finset.card_powersetCard), each of that probability. The hypothesis $n_{1} \leq n$ is doing exactly one job — making $(\binom{n}{n_{1}}) \neq 0$ , so the pmf is well defined.

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

How many n₁-subsets contain a given setNew tabOpens with 506 lines of library preamble — the code above is at the bottom of the editor.

/-- **How many `k`-subsets of the population contain a given set `S`?**  Answer:
`C(n - #S, k - #S)`, because after committing to include `S` one still chooses `k - #S`
units freely from the remaining `n - #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

Having committed to treating all of $S$ , one still chooses $n_{1} - # S$ units from the remaining $n - # S$ . So

π (S) = \frac{(\binom{n - # S}{n_{1} - # S})}{(\binom{n}{n_{1}})},

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

C(n-1, n₁-1) / C(n, n₁) = n₁ / nNew tabOpens with 527 lines of library preamble — the code above is at the bottom of the editor.

/-- `C(n-1, n₁-1) / C(n, n₁) = n₁ / n`: the fraction of `n₁`-subsets containing a fixed unit.

This is the counting fact behind "the propensity score of complete randomisation is `n₁/n`". -/
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'

The proof is a trick worth learning. Rather than manipulate n - 1 and n₁ - 1 — truncated natural subtraction, which is a minefield — it destructs the hypotheses: obtain ⟨j, rfl⟩ : ∃ j, n₁ = j + 1 rewrites n₁ to j + 1 everywhere, and similarly for n. Now the subtractions are literal (Nat.add_sub_cancel), and Mathlib’s Nat.add_one_mul_choose_eq : (n+1) * C(n,k) = C(n+1,k+1) * (k+1) chains into the result by linear_combination, which proves a linear-arithmetic goal from a given identity by subtracting it. No cancellation, no division in ℕ.

And the payoff, twice:

TheoremPropensity under complete randomization

Under complete randomization with $n_{1}$ treated units, every unit has propensity $π_{i} = n_{1} / n$ .

Note what is not assumed: nothing beyond $n_{1} \leq n$ . The degenerate case $n_{1} = 0$ is covered by the branch that counts zero subsets, and its right-hand side $0 / n$ is $0$ . This is why the theorem can be stated as an unconditional equation.

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

/-- **The propensity score of complete randomisation is `n₁ / n`** for every unit.

Note this matches the Bernoulli design with `p = n₁/n`: the two designs are
indistinguishable at the level of individual units, and differ only in their joint behaviour. -/
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
  · -- degenerate: nobody is treated, so the propensity is `0 = 0 / n`
    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

Joint propensity under complete randomizationNew tabOpens with 662 lines of library preamble — the code above is at the bottom of the editor.

/-- **The joint propensity of complete randomisation** for two distinct units:
`P(i and j both treated) = n₁(n₁-1) / (n(n-1))`.

Compare with the Bernoulli value `p² = (n₁/n)²`: the completely randomised value is strictly
smaller whenever `0 < n₁ < n`, which is the *negative dependence* induced by fixing the number
of treated units. -/
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
  · -- degenerate: `n₁ ∈ {0, 1}`, so both sides are zero
    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

The contrast between the two designs is the whole story of 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:

\underset{complete randomization}{\frac{n_{1} (n_{1} - 1)}{n (n - 1)}} < \underset{Bernoulli}{(\frac{n_{1}}{n})^{2}} whenever 0 < n_{1} < n .

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.

ExerciseThe support of complete randomization

Warm-up. Only assignments with exactly $n_{1}$ treated units have positive probability. Argue by contradiction (by_contra): if #z ≠ n₁ then the if in the pmf takes its else branch and the probability is 0, contradicting the hypothesis. Chapter 8 uses this lemma with expect_congr_of_support to replace difference-in-means by Horvitz–Thompson.

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

/-- Only assignments treating exactly `n₁` units can occur: this describes the **support**
of the design, and is what lets us replace an estimator by an equal one on the support. -/
theorem card_eq_of_prob_ne_zero {z : Assignment n}
    (hz : (completeRandomization n n₁ h).prob z ≠ 0) : #z = n₁ := by
  sorry

Show solution

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

/-- Only assignments treating exactly `n₁` units can occur: this describes the **support**
of the design, and is what lets us replace an estimator by an equal one on the support. -/
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])

ExerciseThe degenerate count

Core. When fewer units are treated than $S$ has members, no assignment contains $S$ . Finset.card_eq_zero and Finset.filter_eq_empty_iff reduce the goal to “no n₁-subset contains S”; Finset.card_le_card on S ⊆ z then contradicts #z = n₁ < #S, and omega closes the arithmetic.

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

/-- If we ask for fewer treated units than `S` has members, no assignment can treat all of
`S`, so the count is zero.  This degenerate case is what makes the joint-propensity formula
below hold with no side condition on `n₁`. -/
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

Show solution

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

/-- If we ask for fewer treated units than `S` has members, no assignment can treat all of
`S`, so the count is zero.  This degenerate case is what makes the joint-propensity formula
below hold with no side condition on `n₁`. -/
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)

ExerciseNegative dependence

Stretch. Prove the inequality above. Rewrite both sides with the two propensity theorems, put the product over a common denominator with div_mul_div_comm, clear the denominators with div_lt_div_iff₀ (which wants both to be positive), and hand the polynomial inequality to nlinarith. It will not find the witness unaided: after clearing denominators the two sides differ by exactly $n_{1} n (n - n_{1})$ , so prove that product positive first and pass it to nlinarith as a hint. The casts (1 : ℝ) ≤ n₁, (n₁ : ℝ) < n and (2 : ℝ) ≤ n come from exact_mod_cast.

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

/-- Exercise (stretch).  **Fixing the number of treated units makes units negatively
dependent.**  For `0 < n₁ < n` the joint propensity of complete randomisation is strictly
below the product of the two propensities — unlike Bernoulli, where the two are equal.

Rewrite both sides with `completeRandomization_jointPropensity` and
`completeRandomization_propensity`, put the right-hand side over one denominator with
`div_mul_div_comm`, clear the denominators with `div_lt_div_iff₀` (both must be shown
positive), and hand the polynomial inequality to `nlinarith`.  It will not find the
witness on its own: the gap between the two sides is exactly `n₁ · n · (n - n₁)`, so
supply that product as a hint.  The three casts it needs, `(1 : ℝ) ≤ n₁`, `(n₁ : ℝ) < n`
and `(2 : ℝ) ≤ n`, all come from `exact_mod_cast`. -/
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

Show solution

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

/-- Exercise (stretch).  **Fixing the number of treated units makes units negatively
dependent.**  For `0 < n₁ < n` the joint propensity of complete randomisation is strictly
below the product of the two propensities — unlike Bernoulli, where the two are equal.

Rewrite both sides with `completeRandomization_jointPropensity` and
`completeRandomization_propensity`, put the right-hand side over one denominator with
`div_mul_div_comm`, clear the denominators with `div_lt_div_iff₀` (both must be shown
positive), and hand the polynomial inequality to `nlinarith`.  It will not find the
witness on its own: the gap between the two sides is exactly `n₁ · n · (n - n₁)`, so
supply that product as a hint.  The three casts it needs, `(1 : ℝ) ≤ n₁`, `(n₁ : ℝ) < n`
and `(2 : ℝ) ≤ n`, all come from `exact_mod_cast`. -/
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.

The all-or-nothing designNew tabOpens with 1,086 lines of library preamble — the code above is at the bottom of the editor.

/-- The **all-or-nothing design**: flip one fair coin for the whole population, and either
treat everybody or nobody.

It satisfies both `Design` axioms and gives every unit propensity `1/2`, and it is useless,
because the two arms are never populated at the same time.  The defect is invisible in the
propensity scores and glaring in the joint propensities. -/
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

Only two assignments carry probability, so every expectation collapses to an average of two numbers, and every propensity below is a two-line corollary of that one lemma:

Expectation under the all-or-nothing designNew tabOpens with 1,120 lines of library preamble — the code above is at the bottom of the editor.

/-- Every expectation under the all-or-nothing design is an average of two numbers. -/
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/2 for every unit, and π ij = 1/2 for every pairNew tabOpens with 1,145 lines of library preamble — the code above is at the bottom of the editor.

/-- Every unit has propensity `1/2`, exactly as under a fair Bernoulli design. -/
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

/-- But *every* pair of units is treated together with probability `1/2`, not `1/4`: the
all-or-nothing design is as far from independent as a design can be. -/
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

/-- Independence fails for every pair. -/
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

Every unit has $π_{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 0 < \pi_i and \pi_i < 1, and 1/2 satisfies both:

Horvitz–Thompson is unbiased under the all-or-nothing designNew tabOpens with 1,167 lines of library preamble — the code above is at the bottom of the editor.

/-- **Horvitz--Thompson is unbiased under the all-or-nothing design.**  `HT_unbiased` asks only
for `0 < π i < 1`, and `allOrNothing_propensity` supplies `π i = 1/2`; nothing about the joint
behaviour of the units is needed, or available.

The corollary is worth stating precisely because the design is useless.  Unbiasedness is a
statement about the marginals `π i`, and this design has the marginals of a fair coin flip per
unit while never producing a treated unit and a control unit at the same time. -/
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)

And $π_{i j} = 1 / 2 \neq 1 / 4 = π_{i} π_{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 $π_{i}$ ; usefulness is a statement about $π_{i j}$ .

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

A deterministic designNew tabOpens with 1,182 lines of library preamble — the code above is at the bottom of the editor.

/-- The **deterministic design** that always produces the assignment `z₀`.  It satisfies both
`Design` axioms and carries no randomness at all. -/
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

/-- Expectation under a point mass is evaluation. -/
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

/-- Its propensities are the indicators of `z₀`: every unit has propensity `0` or `1`, so the
hypotheses `0 < π i` and `π i < 1` of `HT_unbiased` both fail somewhere. -/
theorem pointMass_propensity (z₀ : Assignment n) (i : Fin n) :
    (pointMass z₀).propensity i = Z z₀ i :=
  pointMass_expect z₀ _

/-- No unit of a deterministic design has a propensity strictly between `0` and `1`. -/
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

Its propensities are $0$ or $1$ , so it violates both hypotheses of HT_unbiased — and the violation is not a technicality. Take the point mass at “everybody treated” and compute:

The bias of Horvitz–Thompson when the hypotheses failNew tabOpens with 1,208 lines of library preamble — the code above is at the bottom of the editor.

/-- Treat everybody, with probability one: every propensity is `1`. -/
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)]

/-- **Horvitz–Thompson under a design that violates its hypotheses.**

Under the point mass at `univ` no unit is ever a control, the control term of `HT` is
`0 * y0 / 0 = 0` by Lean's division convention, and the estimator collapses to the mean of
the treated potential outcomes. -/
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

/-- The bias is exactly the control mean, which is not zero for a generic population.  This
is what the hypothesis `π i < 1` in `HT_unbiased` buys. -/
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

The control term of the estimator is $(1 - Z_{i}) y_{0 i} / (1 - π_{i}) = 0 \cdot y_{0 i} / 0$ , which Lean evaluates to $0$ by the division-by-zero convention. The estimator therefore returns $\overset{―}{y_{1}}$ , and the bias is exactly $\overset{―}{y_{0}}$ . Note what the 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 like $0 < π_{i} < 1$ is load-bearing and its absence is silent.

ExerciseAffine statistics

Warm-up. $E [a f + b] = a E [f] + b$ . Three rewrites: expect_add, then expect_const_mul on the first piece and expect_const on the second. The only difficulty is writing the two statistics as explicit lambdas so that rw can match them.

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

/-- Exercise (warm-up).  **Linearity in the form you actually use it**: expectation is
affine, `E[a f + b] = a E[f] + b`.

`Design.expect_add` splits the sum, `Design.expect_const_mul` pulls the constant out of the
first piece, and `Design.expect_const` evaluates the second.  Each is stated about an
explicit `fun z => …`, so the whole proof is one `rw` chain. -/
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

Show solution

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

/-- Exercise (warm-up).  **Linearity in the form you actually use it**: expectation is
affine, `E[a f + b] = a E[f] + b`.

`Design.expect_add` splits the sum, `Design.expect_const_mul` pulls the constant out of the
first piece, and `Design.expect_const` evaluates the second.  Each is stated about an
explicit `fun z => …`, so the whole proof is one `rw` chain. -/
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 $(Ω, F, P)$ , define $Z : Ω \to {0, 1}^{n}$ as a measurable map, and define $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 $P (Ω)$ , so every function is measurable;
$Ω$ is finite, so every function is integrable;
the integral is the finite sum ∑ z, D.prob z * f z.

So the elementary development in this chapter is not a simplification that will have to be redone later; it is the same objects with the bureaucracy removed. Chapter 13 makes that precise — Design.toPMF, the theorem ∫ z, f z ∂μ = D.expect f, and HT_unbiased restated in MeasureTheory vocabulary — and proves the Bernoulli design’s units are iIndepFun in Mathlib’s sense, which is bernoulli_inclusion and nothing more. You need to be able to read that statement; you do not need it to prove anything in between.

Reading autoformalized Lean

Two ways a “design” can be formalized into something that is not a design.

A design given by its propensities. Here is a plausible-looking structure:

-- WRONG: this is not a design
structure Design' (n : ℕ) where
  propensity : Fin n → ℝ
  nonneg : ∀ i, 0 ≤ propensity i
  le_one : ∀ i, propensity i ≤ 1

Every field is true of a real design. It even supports a definition of Horvitz–Thompson, and it is what a paper’s notation looks like on the page. But it is a strictly weaker object: it cannot express “exactly $n_{1}$ units are treated”, it has no joint propensities, and therefore no variance formula in chapters 9 to 12 can even be stated over it. When a formalization replaces a distribution by a vector of summaries, the tell is that the interesting theorems about it become unstateable rather than unprovable.

An expectation with the wrong weights. This one type-checks, is provable, and is wrong:

-- WRONG: uniform, not `D`
noncomputable def expect' (D : Design n) (f : Assignment n → ℝ) : ℝ :=
  (∑ z, f z) / (2 ^ n : ℝ)

D appears in the signature and is never used, so every theorem proved about expect' is a theorem about the uniform design dressed up as a theorem about an arbitrary one. It would even satisfy linearity, and $E [c] = c$ , and expect_mono — all the lemmas of this chapter go through. The check is not to read the theorems; it is to read the definition and ask whether each argument is used. #print expect' shows you the body.

Next chapter: with a design in hand and its propensities computed, the payoff. Horvitz–Thompson is unbiased for any design with $0 < π_{i} < 1$ , and difference-in-means is unbiased under complete randomization — because, on the support of that design, it is Horvitz–Thompson.

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

A design is its pmf

Expectation is a weighted sum, and that is all

Indicators, events, probabilities

Propensities

Variance and covariance

Computing with a design you can see

The Bernoulli design

Complete randomization

Two designs that exist to break things

How a probabilist would set this up

Reading autoformalized Lean

Footnotes