Potential outcomes

Part I was about Lean. From here on Lean is in service of something.

This chapter builds the fixed half of the design-based model: a finite population whose potential outcomes are constants, and an assignment that decides which of them you get to see. Nothing here is probabilistic — the word expectation does not appear until the next chapter — and yet the central impossibility result of the field is already provable, which is a reasonable advertisement for doing it this way round.

The population is two vectors

There are $n$ units. Unit $i$ carries two numbers: $y_1(i)$, the outcome it would show under treatment, and $y_0(i)$, the outcome it would show under control. Both numbers exist for every unit. At most one of them is ever seen.

/-- A **finite population** of `n` units, described by its potential outcomes.

`y1 i` is the outcome unit `i` would show if treated, `y0 i` the outcome it would show if
left in control.  Both are fixed real numbers: in design-based inference the population is
*not* a random sample from anything, and these numbers are not random variables.

Because `y1` and `y0` are functions of a single unit `i`, no unit's potential outcomes can
depend on anybody else's treatment.  That is SUTVA / no interference, enforced by the type. -/
structure Population (n : ℕ) where
  /-- The outcome unit `i` would show under treatment. -/
  y1 : Fin n → ℝ
  /-- The outcome unit `i` would show under control. -/
  y0 : Fin n → ℝ

structure … where declares a record; y1 and y0 are its fields. Each is a function Fin n → ℝ, and Fin n is the type with exactly n elements, so a Population 500 is literally a pair of length-500 real vectors. To build one you supply both fields; to take one apart you write P.y1 and P.y0.

Read what is not there. There is no distribution, no error term, no superpopulation the units were drawn from, and no random variable of any kind. P.y1 i is a real number in the same sense that $\pi$ is: fixed, and mostly unknown to you.¹This is Neyman’s 1923 model, and the reason the tutorial is arranged around it is that its randomness has exactly one source. A Lean development of the superpopulation version would need measure theory on page one; this one needs finite sums.

The average treatment effect is the mean of the individual effects:

/-- The **individual treatment effect** of unit `i`: `y1 i - y0 i`.

This is the quantity we would like to know and can never observe for any single unit. -/
def effect (i : Fin n) : ℝ := P.y1 i - P.y0 i

/-- The **average treatment effect** (ATE) of the population,
`τ = (1/n) ∑ i, (y1 i - y0 i)`.

This is a fixed number attached to the population, not a parameter of a distribution. -/
noncomputable def tau : ℝ := mean P.effect

with mean the obvious thing:

/-- The **finite-population mean** of a vector `v : Fin n → ℝ`, i.e. `(∑ i, v i) / n`.

Note Lean's convention that `x / 0 = 0`, so `mean` on an empty population is `0` rather than
undefined.  This convention is a genuine convenience: it lets every statement below be a
plain equation with no side conditions about `n ≠ 0` hiding in the definition. -/
noncomputable def mean (v : Fin n → ℝ) : ℝ := (∑ i, v i) / n

tau is a noncomputable def because real division has no algorithm. effect is not, because subtraction does. Neither fact affects a single proof; Lean just insists you say so.

The identity every applied paper opens with — the ATE is the difference of the two means — is one line, because mean is linear:

/-- The ATE is the difference of the two potential-outcome means:
`τ = mean y1 - mean y0`.

Obvious on paper; worth having as a named lemma because it is used constantly. -/
theorem tau_eq_mean_sub_mean : P.tau = mean P.y1 - mean P.y0 :=
  mean_sub P.y1 P.y0

Interrogating a definition

Three commands do almost all the work of reading someone else’s Lean. #check e prints the type of e. #check @f prints it with the implicit arguments made visible — @ switches off Lean’s habit of filling them in, and it is usually what you want when you are auditing a statement. #print f prints the body of f, which is where the content of a definition lives. And example states a proposition with no name, to be checked and thrown away.

-- `#check` reports a type.  `@` shows the arguments Lean usually infers.
#check @Population              -- Population : ℕ → Type
#check @Population.y1           -- {n : ℕ} → Population n → Fin n → ℝ
#check @Population.tau          -- {n : ℕ} → Population n → ℝ
#check @Z                       -- {n : ℕ} → Assignment n → Fin n → ℝ
#check @Yobs                    -- {n : ℕ} → Population n → Assignment n → Fin n → ℝ

-- `#print` shows a definition's body, which is where the content actually is.
#print Yobs
#print Population.effect

-- `Assignment n` is an `abbrev`, so it is *the same type* as `Finset (Fin n)`,
-- not a copy of it: `rfl` proves the equation.
example : Assignment 3 = Finset (Fin 3) := rfl

-- And `Yobs` really is the algebraic form; `rfl` means "true by unfolding".
example (P : Population 3) (z : Assignment 3) (i : Fin 3) :
    Yobs P z i = Z z i * P.y1 i + (1 - Z z i) * P.y0 i := rfl

rfl proves an equation that holds by unfolding definitions alone. Both examples above are proved by rfl, and that is a meaningful check: the first says Assignment 3 is the same type as Finset (Fin 3) rather than a copy of it, the second says Yobs is by definition the algebraic expression and not merely equal to it by some argument.

Nothing here is abstract. A three-unit population can be written down and its $\tau$ computed:

/-- Three units, written out.  `![a, b, c]` is Mathlib's notation for the
function `Fin 3 → ℝ` taking those three values. -/
def toy : Population 3 where
  y1 := ![3, 5, 4]
  y0 := ![1, 2, 3]

/-- Its average treatment effect is `(2 + 3 + 1) / 3 = 2`.  Nothing here is
observable; this is the number a study is trying to estimate. -/
example : toy.tau = 2 := by
  norm_num [Population.tau, mean, Population.effect, toy, Fin.sum_univ_three,
    Matrix.cons_val_two, Matrix.tail_cons, Matrix.head_cons]

/-- Treat units `0` and `2`.  The observed vector reveals `y1 0`, `y0 1`, `y1 2`
and says nothing whatever about `y0 0`, `y1 1`, `y0 2`. -/
example : Yobs toy {0, 2} = ![3, 2, 4] := by
  funext i
  fin_cases i <;> simp [Yobs_eq_ite, toy]

The second example is the whole chapter in miniature. Treating units 0 and 2 reveals y1 0, y0 1, y1 2 — and the three numbers y0 0, y1 1, y0 2, which $\tau$ depends on just as much, are not in the observed vector anywhere.

An assignment is the treated set

An assignment says who got treated. The obvious encoding is a vector of zeros and ones, Fin n → Bool. The library uses the set of treated units instead.

/-- An **assignment** of treatment to a population of `n` units: the (finite) set of units
that receive treatment.

We represent an assignment by the *treated set* `Finset (Fin n)` rather than by an indicator
vector `Fin n → Bool`.  See `SPEC.md` for the full rationale; the short version is that
"exactly `n₁` units are treated" is then literally `Finset.powersetCard n₁ univ`, and
Mathlib's combinatorics of subsets (`Finset.card_powersetCard`,
`Finset.card_filter_powersetCard_subset`, `Finset.prod_add`) applies directly. -/
abbrev Assignment (n : ℕ) := Finset (Fin n)

This is a design decision worth a paragraph, because it is the one that decides how hard the rest of the library is to prove.

• “Exactly $n_1$ units are treated” becomes Finset.powersetCard n₁ univ, an object Mathlib already knows how to count. Complete randomization in chapter 7 is then a Mathlib lemma away rather than a bespoke combinatorial argument.
• Bernoulli assignment sums become Finset.prod_add, again already in Mathlib.
• The control group is free: it is zᶜ, and every lemma proved about the treated side transfers to the control side by complementation.
• Finset (Fin n) is a Fintype, so ∑ z, … over all $2^n$ assignments is a finite sum that needs no extra bookkeeping. That sum is what an expectation over a design will be.

abbrev rather than def matters: an abbrev is reducible, so Lean will unfold Assignment n to Finset (Fin n) on its own, and every Finset lemma and instance applies to an assignment without a word from you.

Nothing downstream pattern-matches on the set. The algebra is written in terms of the indicator, which is real-valued:

/-- The **treatment indicator** `Z z i ∈ {0, 1}`, as a *real number*: `1` if unit `i` is
treated under assignment `z`, and `0` otherwise.

Keeping the indicator in `ℝ` (rather than `Bool` or `Prop`) is what lets us write the usual
algebraic identities of design-based inference, e.g. `Yobs = Z * y1 + (1 - Z) * y0`. -/
def Z (z : Assignment n) (i : Fin n) : ℝ := if i ∈ z then 1 else 0

@[simp] theorem Z_of_mem {z : Assignment n} {i : Fin n} (h : i ∈ z) : Z z i = 1 := ite_eq_left h

@[simp] theorem Z_of_not_mem {z : Assignment n} {i : Fin n} (h : i ∉ z) : Z z i = 0 := ite_eq_right h

Z z i is 1 or 0 — the number, not the Boolean and not the proposition. Design-based inference is algebra with indicators. In Bool you would be writing a coercion at every occurrence of Z; in Prop you would be reasoning about propositions rather than computing with them. In ℝ the usual identities are ordinary equations in a field, so ring, field_simp and linarith apply:

/-- The indicator is idempotent: `Z² = Z`, because `0² = 0` and `1² = 1`.

This tiny lemma is the workhorse behind every variance computation for indicators. -/
theorem Z_mul_self (z : Assignment n) (i : Fin n) : Z z i * Z z i = Z z i := by
  unfold Z; split <;> norm_num

/-- `Z z i ^ 2 = Z z i`. -/
theorem Z_sq (z : Assignment n) (i : Fin n) : Z z i ^ 2 = Z z i := by
  rw [sq, Z_mul_self]

/-- The **control indicator** `1 - Z z i` is the treatment indicator of the complementary
assignment.  This lets every fact proved about treated units be reused for control units. -/
theorem one_sub_Z (z : Assignment n) (i : Fin n) : 1 - Z z i = Z zᶜ i := by
  unfold Z
  by_cases h : i ∈ z <;> simp [h]

The second is the complementation trick made precise: 1 - Z z i is not merely like an indicator, it is the indicator of zᶜ.

Counting the two groups

The treated count is $n_1=\mathrm{\#}z$ and the control count is $n_0=n-n_1$. In Lean the treated count arrives as a sum of indicators:

/-- The number of treated units is the sum of the indicators: `∑ i, Z z i = #z`. -/
theorem sum_Z (z : Assignment n) : ∑ i, Z z i = (#z : ℝ) := by
  sorry

/-- The number of treated units is the sum of the indicators: `∑ i, Z z i = #z`. -/
theorem sum_Z (z : Assignment n) : ∑ i, Z z i = (#z : ℝ) := by
  simp only [Z, Finset.sum_ite_mem, Finset.univ_inter, Finset.sum_const, nsmul_eq_mul, mul_one]

and the control count follows by complementation, with no new work:

/-- The control indicators sum to the number of control units.

The proof is one rewrite: `1 - Z z i` *is* the treatment indicator of the
complementary assignment (`one_sub_Z`), so this is `sum_Z` applied to `zᶜ`. -/
theorem sum_one_sub_Z (z : Assignment n) : ∑ i, (1 - Z z i) = (#zᶜ : ℝ) := by
  sorry

/-- The control indicators sum to the number of control units.

The proof is one rewrite: `1 - Z z i` *is* the treatment indicator of the
complementary assignment (`one_sub_Z`), so this is `sum_Z` applied to `zᶜ`. -/
theorem sum_one_sub_Z (z : Assignment n) : ∑ i, (1 - Z z i) = (#zᶜ : ℝ) := by
  simp only [one_sub_Z]
  exact sum_Z zᶜ

The cast is where it gets annoying. card_compl says #zᶜ = n - #z in ℕ, where subtraction is truncated, and getting that into ℝ requires knowing that the subtraction does not truncate:

/-- The control count as a real number: `#zᶜ = n - #z`, with `-` the *real*
subtraction.

`card_compl` says the same thing in `ℕ`, where `-` is truncated, so the cast
across the subtraction needs `Nat.cast_sub` and therefore a proof that
`#z ≤ n`.  This is the commonest cast obstacle in the whole development. -/
theorem card_compl_cast (z : Assignment n) : ((#zᶜ : ℕ) : ℝ) = (n : ℝ) - #z := by
  sorry

/-- The control count as a real number: `#zᶜ = n - #z`, with `-` the *real*
subtraction.

`card_compl` says the same thing in `ℕ`, where `-` is truncated, so the cast
across the subtraction needs `Nat.cast_sub` and therefore a proof that
`#z ≤ n`.  This is the commonest cast obstacle in the whole development. -/
theorem card_compl_cast (z : Assignment n) : ((#zᶜ : ℕ) : ℝ) = (n : ℝ) - #z := by
  have h : #z ≤ n := by simpa using Finset.card_le_univ z
  rw [MrCLean.card_compl, Nat.cast_sub h]

The observed outcome

You see $y_1(i)$ for treated units and $y_0(i)$ for control units, and that is all you see.

/-- The **observed outcome** of unit `i` under assignment `z`:
`Yobs = Z * y1 + (1 - Z) * y0`.

Exactly one of the two potential outcomes is revealed, and which one is revealed is the only
thing the assignment controls.  The *algebraic* form given here (rather than an `if`) is the
one that plays well with linearity of expectation; `Yobs_eq_ite` says it agrees with the
`if`-form you would write informally. -/
def Yobs (P : Population n) (z : Assignment n) (i : Fin n) : ℝ :=
  Z z i * P.y1 i + (1 - Z z i) * P.y0 i

The library takes the algebraic form as the definition and proves the if form as a lemma, rather than the other way round. That ordering is deliberate: the algebraic form is linear in Z, which is what makes every expectation computation in chapters 8 and 9 a rewrite rather than a case split.

/-- The algebraic form of `Yobs` agrees with the obvious case split: the observed outcome is
`y1 i` for treated units and `y0 i` for control units. -/
theorem Yobs_eq_ite (P : Population n) (z : Assignment n) (i : Fin n) :
    Yobs P z i = if i ∈ z then P.y1 i else P.y0 i := by
  sorry

/-- The algebraic form of `Yobs` agrees with the obvious case split: the observed outcome is
`y1 i` for treated units and `y0 i` for control units. -/
theorem Yobs_eq_ite (P : Population n) (z : Assignment n) (i : Fin n) :
    Yobs P z i = if i ∈ z then P.y1 i else P.y0 i := by
  unfold Yobs Z
  by_cases h : i ∈ z <;> simp [h]

Once you have that, the two branches are immediate, and they are the lemmas you will actually reach for:

/-- A treated unit reveals its treated potential outcome. -/
theorem Yobs_of_mem {P : Population n} {z : Assignment n} {i : Fin n} (h : i ∈ z) :
    Yobs P z i = P.y1 i := by rw [Yobs_eq_ite, ite_eq_left h]

/-- A control unit reveals its control potential outcome. -/
theorem Yobs_of_not_mem {P : Population n} {z : Assignment n} {i : Fin n} (h : i ∉ z) :
    Yobs P z i = P.y0 i := by rw [Yobs_eq_ite, ite_eq_right h]

An exercise in reading Yobs as a regression writes it:

/-- **Core.** The observed outcome written the way a regression writes it:
control outcome plus the individual effect, switched on by the indicator. -/
theorem Yobs_eq_y0_add_Z_mul_effect (P : Population n) (z : Assignment n) (i : Fin n) :
    Yobs P z i = P.y0 i + Z z i * P.effect i := by
  sorry

/-- **Core.** The observed outcome written the way a regression writes it:
control outcome plus the individual effect, switched on by the indicator. -/
theorem Yobs_eq_y0_add_Z_mul_effect (P : Population n) (z : Assignment n) (i : Fin n) :
    Yobs P z i = P.y0 i + Z z i * P.effect i := by
  simp only [Yobs, Population.effect]
  ring

No interference is a type, not an assumption

In a paper, SUTVA is a numbered assumption you state and then hope nobody checks. Here there is nothing to state. y1 : Fin n → ℝ takes a unit and returns a number; it has no argument in which anybody else’s treatment could appear. A model with interference would need

-- Not the model in this library: a different type, and everything downstream
-- would have to change with it.
structure InterferingPopulation (n : ℕ) where
  y : Assignment n → Fin n → ℝ

which is a different structure, so the assumption is not something you can forget: you would have to rewrite the model to violate it.

The consequence is provable rather than assumed:

/-- **No interference / SUTVA.**  If two assignments agree about unit `i`, then
unit `i` shows the same observed outcome under both — however differently the
other `n - 1` units are treated.

There is no hypothesis to state: this is provable because `y1` and `y0` are
functions of a unit alone, so the type of `Population` already forbids
interference. -/
theorem Yobs_congr (P : Population n) (z z' : Assignment n) (i : Fin n)
    (h : i ∈ z ↔ i ∈ z') : Yobs P z i = Yobs P z' i := by
  rw [Yobs_eq_ite, Yobs_eq_ite]
  by_cases hi : i ∈ z
  · rw [ite_eq_left hi, ite_eq_left (h.mp hi)]
  · rw [ite_eq_right hi, ite_eq_right fun hc => hi (h.mpr hc)]

Read the statement: if two assignments agree about unit i then unit i shows the same observed outcome under both, no matter how differently the other $n-1$ units are treated. The proof rewrites with Yobs_eq_ite, splits on i ∈ z, and uses ite_eq_left / ite_eq_right — the lemmas that collapse an if once you know which branch you are in.²If you have read older Lean 4 code you will know these as if_pos and if_neg. They are deprecated as of v4.34; Mathlib is mid-migration on several such names, which is one reason the toolchain here is pinned.

The fundamental problem of causal inference

Informally: you observe one potential outcome per unit, never both, so $\tau$ is not a function of the data. Stated that baldly it is not yet a theorem — “function of the data” needs saying in a way Lean can check. The sharp version is that the observed vector constrains $\tau$ not at all.

/-- **The fundamental problem of causal inference**, in its sharpest form: fix an
assignment `z` and a population `P`, and pick any number `t` you like.  There is
a population that produces *exactly* the same observed outcome vector as `P`
under `z` and whose average treatment effect is `t`.

The data alone therefore restrict `τ` not at all.  Everything later in this
tutorial buys identification back by making `z` random, not by learning more
about any single unit. -/
theorem exists_population_tau_eq (hn : 0 < n) (P : Population n) (z : Assignment n) (t : ℝ) :
    ∃ P' : Population n, Yobs P' z = Yobs P z ∧ P'.tau = t := by
  refine ⟨shift P z (t - P.tau), Yobs_shift P z _, ?_⟩
  rw [tau_shift hn]
  ring

The proof is a construction. Move every unobserved potential outcome by c and leave every observed one alone:

/-- Move every *unobserved* potential outcome of `P` by `c`, leaving every
observed one alone: a treated unit keeps `y1` and has its (unseen) `y0` lowered
by `c`; a control unit keeps `y0` and has its (unseen) `y1` raised by `c`.

Either way the individual effect `y1 i - y0 i` goes up by exactly `c`, and no
observed number changes. -/
def shift (P : Population n) (z : Assignment n) (c : ℝ) : Population n where
  y1 := fun i => if i ∈ z then P.y1 i else P.y1 i + c
  y0 := fun i => if i ∈ z then P.y0 i - c else P.y0 i

A treated unit keeps y1 and has its unseen y0 lowered by c; a control unit keeps y0 and has its unseen y1 raised by c. Both branches raise the individual effect y1 i - y0 i by exactly c, which is why the two ifs lean opposite ways. Nothing observable moves:

/-- The shifted population is observationally identical to `P` under `z`. -/
theorem Yobs_shift (P : Population n) (z : Assignment n) (c : ℝ) :
    Yobs (shift P z c) z = Yobs P z := by
  funext i
  by_cases h : i ∈ z <;> simp [Yobs_eq_ite, shift, h]

funext i turns the equation between two functions into an equation at an arbitrary unit i — two functions are equal when they agree everywhere, and funext is how you say so in Lean. Then by_cases h : i ∈ z splits on whether i was treated, and in each branch simp unfolds shift, resolves the if with h, and finds both sides identical.

The effect moves uniformly, so the average moves with it:

/-- …and its average treatment effect is `c` larger. -/
theorem tau_shift (hn : 0 < n) (P : Population n) (z : Assignment n) (c : ℝ) :
    (shift P z c).tau = P.tau + c := by
  have heff : (shift P z c).effect = fun i => P.effect i + c := by
    funext i
    by_cases h : i ∈ z <;> simp [Population.effect, shift, h] <;> ring
  calc (shift P z c).tau
      = mean (fun i => P.effect i + c) := by simp only [Population.tau, heff]
    _ = mean P.effect + mean (fun _ : Fin n => c) := mean_add P.effect fun _ => c
    _ = P.tau + c := by simp only [mean_const hn, Population.tau]

The have establishes the pointwise fact — after the case split, both branches are P.effect i + c, so the shifted effect vector is fun i => P.effect i + c — and the calc chain then does the averaging: mean_add splits the mean of a sum, and mean_const hn evaluates the mean of the constant c. That hn : 0 < n is the only hypothesis in the theorem, and it is there because of the division convention: with no units at all, every population has $\tau =0/0=0$ and nothing can be shifted.

The main theorem is now three lines: take c = t - P.tau, hand Yobs_shift back as the first component, and let ring finish. refine ⟨_, _, ?_⟩ supplies the two easy components of the existential and leaves the third as a goal.

The form one usually says out loud follows immediately:

/-- The same fact in the form one usually says out loud: for any assignment of
any non-empty population there are two populations with the same observed
outcomes and different average treatment effects. -/
theorem fundamental_problem (hn : 0 < n) (z : Assignment n) :
    ∃ P P' : Population n, Yobs P z = Yobs P' z ∧ P.tau ≠ P'.tau := by
  obtain ⟨P', hY, ht⟩ :=
    exists_population_tau_eq hn (⟨fun _ => 0, fun _ => 0⟩ : Population n) z 1
  refine ⟨⟨fun _ => 0, fun _ => 0⟩, P', hY.symm, ?_⟩
  have h0 : (⟨fun _ => 0, fun _ => 0⟩ : Population n).tau = 0 := by
    simp [Population.tau, Population.effect, mean]
  rw [h0, ht]
  norm_num

What the data does pin down is exactly the treated units’ y1 and the control units’ y0:

/-- What the observed vector *does* pin down: the treated units' `y1` and the
control units' `y0`, and nothing else.  Read the right-hand side as the exact
inventory of what a single experiment can ever tell you. -/
theorem Yobs_eq_iff (P P' : Population n) (z : Assignment n) :
    Yobs P z = Yobs P' z ↔ (∀ i ∈ z, P.y1 i = P'.y1 i) ∧ (∀ i ∉ z, P.y0 i = P'.y0 i) := by
  sorry

/-- What the observed vector *does* pin down: the treated units' `y1` and the
control units' `y0`, and nothing else.  Read the right-hand side as the exact
inventory of what a single experiment can ever tell you. -/
theorem Yobs_eq_iff (P P' : Population n) (z : Assignment n) :
    Yobs P z = Yobs P' z ↔ (∀ i ∈ z, P.y1 i = P'.y1 i) ∧ (∀ i ∉ z, P.y0 i = P'.y0 i) := by
  constructor
  · intro h
    refine ⟨fun i hi => ?_, fun i hi => ?_⟩
    · have hc := congrFun h i
      rwa [Yobs_of_mem hi, Yobs_of_mem hi] at hc
    · have hc := congrFun h i
      rwa [Yobs_of_not_mem hi, Yobs_of_not_mem hi] at hc
  · rintro ⟨h1, h0⟩
    funext i
    by_cases hi : i ∈ z
    · rw [Yobs_of_mem hi, Yobs_of_mem hi, h1 i hi]
    · rw [Yobs_of_not_mem hi, Yobs_of_not_mem hi, h0 i hi]

None of this is fixed by collecting more units, and none of it is fixed by measuring more carefully. It is fixed — in the weaker sense of “made estimable on average” — by making $z$ random, which is the subject of the next chapter. The estimators of chapter 8 are not functions that recover $\tau$ from the data; they are functions whose average over the design equals $\tau$. Chapter 9 then asks how far from $\tau$ any single realization is likely to be.

Exercises

The exercises above are the load-bearing ones. Three more, in increasing order of effort.

/-- The mean is additive. -/
theorem mean_add (u v : Fin n → ℝ) : mean (fun i => u i + v i) = mean u + mean v := by
  sorry

/-- The mean is additive. -/
theorem mean_add (u v : Fin n → ℝ) : mean (fun i => u i + v i) = mean u + mean v := by
  simp only [mean, Finset.sum_add_distrib, add_div]

/-- **Warm-up.** If everybody is treated, the observed vector is `y1`. -/
theorem Yobs_univ (P : Population n) : Yobs P Finset.univ = P.y1 := by
  sorry

/-- **Warm-up.** If everybody is treated, the observed vector is `y1`. -/
theorem Yobs_univ (P : Population n) : Yobs P Finset.univ = P.y1 := by
  funext i
  exact Yobs_of_mem (Finset.mem_univ i)

/-- **Stretch.** The mean of the observed vector is *not* an estimate of
anything in particular: it is the treated units' `y1` and the control units'
`y0`, averaged over the whole population.  Prove the identity, then notice that
`z` appears in it and `τ` does not. -/
theorem mean_Yobs (P : Population n) (z : Assignment n) :
    mean (Yobs P z) = ((∑ i ∈ z, P.y1 i) + ∑ i ∈ zᶜ, P.y0 i) / n := by
  sorry

/-- **Stretch.** The mean of the observed vector is *not* an estimate of
anything in particular: it is the treated units' `y1` and the control units'
`y0`, averaged over the whole population.  Prove the identity, then notice that
`z` appears in it and `τ` does not. -/
theorem mean_Yobs (P : Population n) (z : Assignment n) :
    mean (Yobs P z) = ((∑ i ∈ z, P.y1 i) + ∑ i ∈ zᶜ, P.y0 i) / n := by
  rw [sum_mem_eq_sum_Z_mul, sum_compl_eq_sum_one_sub_Z_mul, ← Finset.sum_add_distrib]
  rfl

Reading autoformalized Lean

Suppose you asked a model to formalize “the fundamental problem of causal inference” and got one of the following. Each type-checks. None of them says it.

The quantifiers swapped

-- WRONG
theorem fundamental_problem (hn : 0 < n) (z : Assignment n) (P P' : Population n) :
    Yobs P z = Yobs P' z → P.tau ≠ P'.tau

The hypothesis nobody can satisfy

-- WRONG
theorem fundamental_problem' (hn : 0 < n) (z : Assignment n) (y : Fin n → ℝ)
    (h : ∀ P : Population n, Yobs P z = y) :
    ∃ P P' : Population n, Yobs P z = Yobs P' z ∧ P.tau ≠ P'.tau

The definition doing the theorem’s work

-- WRONG, and the theorem statement is unchanged
def Yobs (P : Population n) (z : Assignment n) (i : Fin n) : ℝ := P.y1 i