Regression as linear algebra

In a finite population, potential outcomes are fixed and only the assignment is random. For fixed $X$ and $y$ , an OLS coefficient is therefore an algebraic function of the observed data: a solution to the normal equations. No stochastic error model is needed to define it.

The identities in this chapter hold for each assignment separately. They show that the treatment coefficient in a simple regression equals difference in means, that group indicators recover group means, and that a fully interacted regression separates into within-group fits. Statistical properties follow by combining these identities with results about the corresponding estimators.

Matrices in Mathlib

A design matrix is Matrix (Fin n) (Fin k) ℝ: $n$ rows, one per unit, and $k$ columns, one per regressor. Underneath it is nothing but a function Fin n → Fin k → ℝ; the type Matrix exists so that * can mean matrix multiplication rather than pointwise multiplication. Matrix.of is the relabelling that takes you from the function to the matrix, and it is why simpleMatrix below is written the way it is.

Three notations do all the work. Xᵀ is the transpose. X * Y is matrix multiplication, so Xᵀ * X is the Gram matrix — the sums of squares and cross products, $X^{⊤} X$ . And X *ᵥ v is Matrix.mulVec, the matrix–vector product, whose $i$ -th entry is the fitted value $\sum_{j} X_{i j} v_{j}$ .

Gram entries, associativity, and what IsOLS unfolds toNew tabOpens with 1,188 lines of library preamble — the code above is at the bottom of the editor.

theorem gram_apply {k : ℕ} (X : Matrix (Fin n) (Fin k) ℝ) (j l : Fin k) :
    (Xᵀ * X) j l = ∑ i, X i j * X i l := by
  simp [Matrix.mul_apply, Matrix.transpose_apply]

theorem mulVec_assoc {k : ℕ} (X : Matrix (Fin n) (Fin k) ℝ) (β : Fin k → ℝ) :
    (Xᵀ * X) *ᵥ β = Xᵀ *ᵥ (X *ᵥ β) :=
  (Matrix.mulVec_mulVec β Xᵀ X).symm

theorem isOLS_unfold {k : ℕ} (X : Matrix (Fin n) (Fin k) ℝ) (y : Fin n → ℝ) (β : Fin k → ℝ) :
    IsOLS X y β ↔ (Xᵀ * X) *ᵥ β = Xᵀ *ᵥ y := Iff.rfl

All three are definitional, and the walkthrough is mostly about what “definitional” means in each case.

```
simp [Matrix.mul_apply, Matrix.transpose_apply]
```
gram_apply. Matrix.mul_apply unfolds an entry of a product, $(A B)_{j l} = \sum_{i} A_{j i} B_{i l}$ , and Matrix.transpose_apply says $(X^{⊤})_{j i} = X_{i j}$ . Together they turn $(X^{⊤} X)_{j l}$ into $\sum_{i} X_{i j} X_{i l}$ — the cross-product of columns $j$ and $l$ , which is what the Gram matrix collects.
```
(Matrix.mulVec_mulVec β Xᵀ X).symm
```
mulVec_assoc, a term-mode proof rather than a tactic block. Matrix.mulVec_mulVec states $A \cdot (B v) = (A B) v$ ; .symm turns the equation round, giving $(X^{⊤} X) β = X^{⊤} (X β)$ . Read right to left it is the one move every proof in this chapter starts with: the normal equations are about the residual $y - X β$ , and getting there means peeling $X^{⊤}$ off the front.
```
Iff.rfl
```
isOLS_unfold, and worth a second look: Iff.rfl says the two sides are the same proposition up to unfolding definitions, so no proof is needed. It is Lean’s way of saying “that is literally the definition” — IsOLS X y β is notation for the normal equations and nothing more.

The precedence trap

*ᵥ binds more tightly than *. The Gram matrix applied to a coefficient vector must therefore be parenthesized:

(Xᵀ * X) *ᵥ β

Without the parentheses, Lean parses the expression as

Xᵀ * (X *ᵥ β)

The second expression attempts to multiply a matrix by a vector using matrix multiplication, so it does not type-check. This particular mistake is benign because Lean rejects it. Chaining is benign too: *ᵥ is right-associative, so Xᵀ *ᵥ X *ᵥ β means Xᵀ *ᵥ (X *ᵥ β), which mulVec_assoc above says is the same vector as (Xᵀ * X) *ᵥ β. What does not announce itself is a transpose in the wrong place: with $k = n$ both X *ᵥ β and Xᵀ *ᵥ β type-check and they are different vectors. Read the parenthesization and the transposes of every *ᵥ you meet.

Small matrices, and one determinant

Literal vectors are ![a, b, …] and literal matrices are !![a, b; c, d], given row by row. A $2 \times 2$ determinant is Matrix.det_fin_two_of, and it is the only determinant this chapter needs.

Vector and matrix literalsNew tabOpens with 1,200 lines of library preamble — the code above is at the bottom of the editor.

theorem literal_mulVec : (!![(1 : ℝ), 2; 3, 4]) *ᵥ ![1, 1] = ![3, 7] := by
  funext i
  fin_cases i <;> simp [Matrix.mulVec, dotProduct, Fin.sum_univ_two] <;> norm_num

theorem literal_det : (!![(1 : ℝ), 2; 3, 4]).det = -2 := by
  norm_num [Matrix.det_fin_two_of]

Neither is a theorem about anything; they are here so that the tactics that make literals compute have been seen once.

```
funext i
```
literal_mulVec is an equality of two functions Fin 2 → ℝ, and funext reduces that to an equality at each index i. Vectors in Mathlib are functions, so this is how nearly every vector identity begins.
```
fin_cases i <;> simp [Matrix.mulVec, dotProduct, Fin.sum_univ_two] <;> norm_num
```
fin_cases i splits Fin 2 into its two elements, giving one goal per entry. In each, simp unfolds Matrix.mulVec and dotProduct to a sum and Fin.sum_univ_two writes that two-term sum out explicitly; norm_num then does the arithmetic, $1 \cdot 1 + 2 \cdot 1 = 3$ and $3 \cdot 1 + 4 \cdot 1 = 7$ . The <;> combinator applies what follows it to every goal the previous tactic left.
```
norm_num [Matrix.det_fin_two_of]
```
literal_det. Matrix.det_fin_two_of is the formula $a d - b c$ for the determinant of a literal $2 \times 2$ matrix with rows $(a, b)$ and $(c, d)$ , and norm_num evaluates $1 \cdot 4 - 2 \cdot 3 = - 2$ .

The inverse is total, and that is a trap

Mathlib’s Matrix.inv is a total function. Every square matrix has an inverse in the sense that A⁻¹ always elaborates; when A.det is not a unit, A⁻¹ is defined to be the zero matrix. This is the same design decision as $x / 0 = 0$ for reals, and it has the same consequence: a statement about $(X^{⊤} X)^{- 1}$ is never ill-formed, it is just sometimes about zero.

ExerciseEveryone treated: the estimate is zero

Take z to be all of Finset.univ. Then the treatment column equals the intercept column, the Gram matrix is singular, and ols is the zero vector. This is a core exercise: it is short, and the fact it establishes is the reason every later theorem carries a hypothesis.

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

theorem ols_simpleMatrix_of_all_treated (y : Fin n → ℝ) :
    ols (simpleMatrix (Finset.univ : Assignment n)) y = 0 := by
  sorry

Hint

simpleMatrix_gram_det gives you the determinant and simp shows it is not a unit; then unfold ols, Matrix.nonsing_inv_apply_not_isUnit converts ¬ IsUnit A.det into A⁻¹ = 0, and Matrix.zero_mulVec finishes.

Show solution

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

theorem ols_simpleMatrix_of_all_treated (y : Fin n → ℝ) :
    ols (simpleMatrix (Finset.univ : Assignment n)) y = 0 := by
  have hdet : ¬ IsUnit ((simpleMatrix (Finset.univ : Assignment n))ᵀ
      * simpleMatrix (Finset.univ : Assignment n)).det := by
    rw [simpleMatrix_gram_det]
    simp
  rw [ols, Matrix.nonsing_inv_apply_not_isUnit _ hdet, Matrix.zero_mulVec]

Least squares without a model

The formalization does not define OLS by minimization. It defines it by the normal equations.

The definition of an OLS fitNew tabOpens with 833 lines of library preamble — the code above is at the bottom of the editor.

def IsOLS {k : ℕ} (X : Matrix (Fin n) (Fin k) ℝ) (y : Fin n → ℝ) (β : Fin k → ℝ) : Prop :=
  (Xᵀ * X) *ᵥ β = Xᵀ *ᵥ y

IsOLS X y β is a Prop, not a number: it is the plain linear equation $X^{⊤} X β = X^{⊤} y$ . It may have one solution, or — when the columns are collinear — infinitely many. Working with the relation rather than with a formula is the single most useful decision in the formalization, because the interesting theorems (“the coefficient is the difference in means”) hold for every solution, whether or not the design matrix has full rank. No inverse, no rank condition, no minimization.

The equivalent geometric statement is the one most statisticians actually carry in their heads: the residual is orthogonal to every column of $X$ .

ExerciseNormal equations = orthogonality (core)

Column $j$ of $X$ is fun i => X i j, so orthogonality to it is $\sum_{i} X_{i j} (y_{i} - (X β)_{i}) = 0$ . Show that $β$ solves the normal equations exactly when the residual is orthogonal to every column.

This is the lemma every later proof goes through: each column of the design matrix contributes one scalar equation about sums over units, and we get to choose columns that isolate the treated and control groups.

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

theorem isOLS_iff_orthogonal {k : ℕ} (X : Matrix (Fin n) (Fin k) ℝ) (y : Fin n → ℝ)
    (β : Fin k → ℝ) :
    IsOLS X y β ↔ ∀ j, ∑ i, X i j * (y i - (X *ᵥ β) i) = 0 := by
  sorry

Hint

After unfold IsOLS, Matrix.mulVec_mulVec used right-to-left rewrites (Xᵀ * X) *ᵥ β as Xᵀ *ᵥ (X *ᵥ β), and funext_iff turns the equality of vectors into a statement about each entry. Both directions are then the same simp only [Matrix.mulVec, dotProduct, Matrix.transpose_apply, mul_sub, Finset.sum_sub_distrib] at the hypothesis and the goal, followed by linarith.

Show solution

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

theorem isOLS_iff_orthogonal {k : ℕ} (X : Matrix (Fin n) (Fin k) ℝ) (y : Fin n → ℝ)
    (β : Fin k → ℝ) :
    IsOLS X y β ↔ ∀ j, ∑ i, X i j * (y i - (X *ᵥ β) i) = 0 := by
  unfold IsOLS
  rw [← Matrix.mulVec_mulVec, funext_iff]
  constructor
  · intro h j
    have hj := h j
    simp only [Matrix.mulVec, dotProduct, Matrix.transpose_apply, mul_sub,
      Finset.sum_sub_distrib] at hj ⊢
    linarith
  · intro h j
    have hj := h j
    simp only [Matrix.mulVec, dotProduct, Matrix.transpose_apply, mul_sub,
      Finset.sum_sub_distrib] at hj ⊢
    linarith

The textbook formula is a definition here, not a theorem, and it is kept firmly separate from IsOLS.

The closed formNew tabOpens with 853 lines of library preamble — the code above is at the bottom of the editor.

noncomputable def ols {k : ℕ} (X : Matrix (Fin n) (Fin k) ℝ) (y : Fin n → ℝ) : Fin k → ℝ :=
  (Xᵀ * X)⁻¹ *ᵥ (Xᵀ *ᵥ y)

Two facts connect them, both short enough to be warm-ups. Existence: when the Gram matrix is invertible the formula does solve the normal equations.

ExerciseThe formula solves the equations (warm-up)

When the Gram matrix is invertible, the textbook formula does solve the normal equations.

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

theorem ols_isOLS {k : ℕ} (X : Matrix (Fin n) (Fin k) ℝ) (y : Fin n → ℝ)
    (h : IsUnit (Xᵀ * X).det) : IsOLS X y (ols X y) := by
  sorry

Hint

Unfold both definitions, then three rewrites: Matrix.mulVec_mulVec collapses the two matrix–vector products into one, Matrix.mul_nonsing_inv cancels the inverse, and Matrix.one_mulVec clears up.

Show solution

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

theorem ols_isOLS {k : ℕ} (X : Matrix (Fin n) (Fin k) ℝ) (y : Fin n → ℝ)
    (h : IsUnit (Xᵀ * X).det) : IsOLS X y (ols X y) := by
  unfold IsOLS ols
  rw [Matrix.mulVec_mulVec, Matrix.mul_nonsing_inv _ h, Matrix.one_mulVec]

And uniqueness, which is the no-perfect-collinearity condition doing its job.

ExerciseA nonsingular Gram matrix pins β down (warm-up)

Two solutions of the normal equations coincide when the Gram matrix is nonsingular. Statistically this is the no-perfect-collinearity condition doing its job.

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

theorem isOLS_unique {k : ℕ} {X : Matrix (Fin n) (Fin k) ℝ} {y : Fin n → ℝ} {β β' : Fin k → ℝ}
    (h : (Xᵀ * X).det ≠ 0) (hβ : IsOLS X y β) (hβ' : IsOLS X y β') : β = β' := by
  sorry

Hint

One exact. Matrix.mulVec_injective_of_det_ne_zero says that A *ᵥ · is injective when A.det ≠ 0, and hβ.trans hβ'.symm gives you two vectors with the same image.

Show solution

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

theorem isOLS_unique {k : ℕ} {X : Matrix (Fin n) (Fin k) ℝ} {y : Fin n → ℝ} {β β' : Fin k → ℝ}
    (h : (Xᵀ * X).det ≠ 0) (hβ : IsOLS X y β) (hβ' : IsOLS X y β') : β = β' := by
  exact Matrix.mulVec_injective_of_det_ne_zero h (hβ.trans hβ'.symm)

Put together they give isOLS_iff_eq_ols: under IsUnit (Xᵀ * X).det, “solves the normal equations” and “equals $(X^{⊤} X)^{- 1} X^{⊤} y$ ” are the same statement. Notice which direction the dependency runs. The formula is a convenience for computing; the relation is what the theorems are about.

Simple regression is the difference in means

Regress the observed outcome on a constant and the treatment indicator. Two columns: ones, and $Z$ .

The design matrixNew tabOpens with 906 lines of library preamble — the code above is at the bottom of the editor.

noncomputable def simpleMatrix (z : Assignment n) : Matrix (Fin n) (Fin 2) ℝ :=
  Matrix.of fun i => ![1, Z z i]

Its Gram matrix is the two-by-two array with entries $n, n_{1}$ on the first row and $n_{1}, n_{1}$ on the second, where $n_{1} = | z |$ is the number of treated units. The only two facts that go into it are $\sum_{i} Z_{i} = n_{1}$ and $Z_{i}^{2} = Z_{i}$ : a dummy is its own square, which is why the $(1, 1)$ entry is $n_{1}$ and not something else.

ExerciseThe Gram matrix (core)

Compute the Gram matrix of the two-column design: it is !![n, n₁; n₁, n₁] with n₁ = #z.

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

theorem simpleMatrix_gram (z : Assignment n) :
    (simpleMatrix z)ᵀ * (simpleMatrix z) = !![(n : ℝ), (#z : ℝ); (#z : ℝ), (#z : ℝ)] := by
  sorry

Hint

ext j k reduces the matrix equation to its four entries and fin_cases j <;> fin_cases k splits them out; then one simp with Matrix.mul_apply (which unfolds (A * B) j k into ∑ i, A j i * B i k), simpleMatrix, Matrix.transpose_apply, sum_Z and Z_mul_self closes all four — the last two are the whole content.

Show solution

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

theorem simpleMatrix_gram (z : Assignment n) :
    (simpleMatrix z)ᵀ * (simpleMatrix z) = !![(n : ℝ), (#z : ℝ); (#z : ℝ), (#z : ℝ)] := by
  ext j k
  fin_cases j <;> fin_cases k <;>
    simp [Matrix.mul_apply, simpleMatrix, Matrix.transpose_apply, sum_Z, Z_mul_self]

Its determinant is $n_{1} (n - n_{1})$ (simpleMatrix_gram_det), so the regression is non-degenerate exactly when both groups are non-empty — the same condition under which the difference in means is defined at all. That is not a coincidence, and the next theorem says why.

Unpacking the normal equations is a genuine ↔; no rank condition is needed to read the system, only to solve it.

Proposition. Let $z$ be an assignment with $n_{1} = | z |$ treated units, and let $y$ be any response vector. Then $β = (β_{0}, β_{1})$ solves the normal equations of the regression of $y$ on $(1, Z)$ if and only if

\sum_{i = 1}^{n} y_{i} = n β_{0} + n_{1} β_{1} and \sum_{i \in z} y_{i} = n_{1} β_{0} + n_{1} β_{1} .

Proof.

By the orthogonality form of the normal equations, $β$ is a least-squares fit exactly when the residual $r_{i} = y_{i} - (β_{0} + Z_{i} β_{1})$ is orthogonal to each of the two columns. Orthogonality to the column of ones is $\sum_{i} r_{i} = 0$ , and since $\sum_{i} Z_{i} = n_{1}$ this is the first displayed equation. Orthogonality to the treatment column is $\sum_{i} Z_{i} r_{i} = 0$ ; here $Z_{i}^{2} = Z_{i}$ — a dummy is its own square — so $Z_{i} r_{i} = Z_{i} y_{i} - Z_{i} β_{0} - Z_{i} β_{1}$ , and summing gives the second equation, because $\sum_{i} Z_{i} y_{i} = \sum_{i \in z} y_{i}$ .∎

The two normal equations, written outNew tabOpens with 938 lines of library preamble — the code above is at the bottom of the editor.

theorem isOLS_simpleMatrix_iff (z : Assignment n) (y : Fin n → ℝ) (β : Fin 2 → ℝ) :
    IsOLS (simpleMatrix z) y β ↔
      (∑ i, y i = (n : ℝ) * β 0 + (#z : ℝ) * β 1)
        ∧ (∑ i ∈ z, y i = (#z : ℝ) * β 0 + (#z : ℝ) * β 1) := by
  rw [isOLS_iff_orthogonal, Fin.forall_fin_two]
  have expand0 : ∑ i, (y i - (β 0 + Z z i * β 1))
      = (∑ i, y i) - ((n : ℝ) * β 0 + (#z : ℝ) * β 1) := by
    rw [Finset.sum_sub_distrib, Finset.sum_add_distrib, Finset.sum_const, Finset.card_univ,
      Fintype.card_fin, nsmul_eq_mul, ← Finset.sum_mul, sum_Z]
  have step1 : ∀ i, Z z i * (y i - (β 0 + Z z i * β 1))
      = Z z i * y i - (Z z i * β 0 + Z z i * β 1) := by
    intro i; linear_combination (-(β 1)) * Z_mul_self z i
  have expand1 : ∑ i, Z z i * (y i - (β 0 + Z z i * β 1))
      = (∑ i ∈ z, y i) - ((#z : ℝ) * β 0 + (#z : ℝ) * β 1) := by
    rw [Finset.sum_congr rfl fun i _ => step1 i, Finset.sum_sub_distrib, Finset.sum_add_distrib,
      ← Finset.sum_mul, ← Finset.sum_mul, sum_Z, ← sum_mem_eq_sum_Z_mul]
  simp only [simpleMatrix_apply_zero, simpleMatrix_apply_one, simpleMatrix_mulVec,
    one_mul, expand0, expand1]
  constructor
  · rintro ⟨h0, h1⟩; exact ⟨by linarith, by linarith⟩
  · rintro ⟨h0, h1⟩; exact ⟨by linarith, by linarith⟩

The Lean does exactly that, with the two sums expanded in named haves so that the final step is a pair of linear rearrangements.

rw [isOLS_iff_orthogonal, …]

Trade the matrix equation for the orthogonality condition — the exercise a few sections up. The left-hand side is now “for every column $j$ , the residual is orthogonal to column $j$ ”.

Goal

z : Assignment n
y : Fin n → ℝ
β : Fin 2 → ℝ
⊢ (∀ (j : Fin 2), ∑ i, simpleMatrix z i j * (y i - (simpleMatrix z *ᵥ β) i) = 0) ↔
    ∑ i, y i = ↑n * β 0 + ↑(#z) * β 1 ∧ ∑ i ∈ z, y i = ↑(#z) * β 0 + ↑(#z) * β 1

rw […, Fin.forall_fin_two]

Fin.forall_fin_two turns a statement quantified over Fin 2 into the conjunction of its two instances, so the two normal equations are now visible separately rather than hidden behind a $\forall$ .

Goal

⊢ ∑ i, simpleMatrix z i 0 * (y i - (simpleMatrix z *ᵥ β) i) = 0 ∧
      ∑ i, simpleMatrix z i 1 * (y i - (simpleMatrix z *ᵥ β) i) = 0 ↔
    ∑ i, y i = ↑n * β 0 + ↑(#z) * β 1 ∧ ∑ i ∈ z, y i = ↑(#z) * β 0 + ↑(#z) * β 1

```
have expand0 : ∑ i, (y i - (β 0 + Z z i * β 1)) = (∑ i, y i) - ((n : ℝ) * β 0 + (#z : ℝ) * β 1) := by rw [Finset.sum_sub_distrib, Finset.sum_add_distrib, Finset.sum_const, Finset.card_univ, Fintype.card_fin, nsmul_eq_mul, ← Finset.sum_mul, sum_Z]
```
The intercept equation, expanded. Finset.sum_sub_distrib and Finset.sum_add_distrib split the sum of a difference and of a sum; Finset.sum_const with Finset.card_univ, Fintype.card_fin and nsmul_eq_mul evaluates $\sum_{i} β_{0}$ as $n β_{0}$ ; and Finset.sum_mul, reversed, pulls $β_{1}$ out of $\sum_{i} Z_{i} β_{1}$ so that sum_Z can replace $\sum_{i} Z_{i}$ by $n_{1}$ .
```
have step1 : ∀ i, Z z i * (y i - (β 0 + Z z i * β 1)) = Z z i * y i - (Z z i * β 0 + Z z i * β 1) := by intro i; linear_combination (-(β 1)) * Z_mul_self z i
```
The one place the dummy’s idempotence is used. Expanding the left-hand side produces a $Z_{i}^{2} β_{1}$ that the right-hand side writes as $Z_{i} β_{1}$ , so the two differ by $β_{1} (Z_{i}^{2} - Z_{i})$ . linear_combination closes a goal of the form “these two sides differ by the given multiple of a known equation”, and the known equation is Z_mul_self z i, $Z_{i} Z_{i} = Z_{i}$ .

have expand1 : ∑ i, Z z i * (y i - (β 0 + Z z i * β 1)) = (∑ i ∈ z, y i) - ((#z : ℝ) * β 0 + (#z : ℝ) * β 1) := by rw [Finset.sum_congr rfl fun i _ => step1 i, Finset.sum_sub_distrib, Finset.sum_add_distrib, ← Finset.sum_mul, ← Finset.sum_mul, sum_Z, ← sum_mem_eq_sum_Z_mul]

The slope equation, expanded. Rewrite each term with step1, split the sums as before, pull $β_{0}$ and then $β_{1}$ out with the two reversed Finset.sum_muls so that sum_Z can turn both $\sum_{i} Z_{i}$ into $n_{1}$ , and finally sum_mem_eq_sum_Z_mul reversed turns $\sum_{i} Z_{i} y_{i}$ into $\sum_{i \in z} y_{i}$ — the treated-group total.

```
simp only [simpleMatrix_apply_zero, simpleMatrix_apply_one, simpleMatrix_mulVec, one_mul, expand0, expand1]
```
Substitute. The first two @[simp] lemmas say the columns of the design matrix are $1$ and $Z$ ; simpleMatrix_mulVec says the fitted value is $β_{0} + Z_{i} β_{1}$ ; one_mul clears the multiplication by the constant column; and expand0/expand1 replace the two sums by their closed forms. Both sides of the iff are now the same two equations, one written as “difference $= 0$ ” and the other as “left $=$ right”.
Goal
```
⊢ ∑ i, y i - (↑n * β 0 + ↑(#z) * β 1) = 0
      ∧ ∑ i ∈ z, y i - (↑(#z) * β 0 + ↑(#z) * β 1) = 0 ↔
    ∑ i, y i = ↑n * β 0 + ↑(#z) * β 1 ∧ ∑ i ∈ z, y i = ↑(#z) * β 0 + ↑(#z) * β 1
```
```
constructor
· rintro ⟨h0, h1⟩; exact ⟨by linarith, by linarith⟩
· rintro ⟨h0, h1⟩; exact ⟨by linarith, by linarith⟩
```
Both directions are the same move twice: destructure the conjunction and move a term across the equals sign, which linarith does. There is no mathematics left at this point.

In words: orthogonality to the column of ones says the residuals sum to zero over everybody; orthogonality to the treatment column says they sum to zero over the treated. Subtract the second from the first and you get $\sum_{i \notin z} y_{i} = (n - n_{1}) β_{0}$ , so $β_{0}$ is the control mean; the second equation then hands you $β_{1}$ .

TheoremThe slope on a treatment dummy is the difference in means

If $β$ solves the normal equations of the regression of the observed outcome on $(1, Z)$ , and both groups are non-empty, then $β_{0}$ is the control-group mean of $y (0)$ and

β_{1} = \frac{1}{n_{1}} \sum_{i \in z} Y_{i} - \frac{1}{n_{0}} \sum_{i \notin z} Y_{i} = {\hat{τ}}_{DiM} .

Read the hypotheses. There is no design, no expectation, no assumption about the potential outcomes: only 0 < #z, #z < n, and IsOLS. The conclusion holds for every assignment $z$ with both groups non-empty, including the ones the design would never produce. That is what “algebraic identity” means.

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

theorem simpleMatrix_ols_coeff (P : Population n) (z : Assignment n) (β : Fin 2 → ℝ)
    (hpos : 0 < #z) (hlt : #z < n) (h : IsOLS (simpleMatrix z) (Yobs P z) β) :
    β 0 = meanOn zᶜ P.y0 ∧ β 1 = DiM P z := by
  obtain ⟨e0, e1⟩ := (isOLS_simpleMatrix_iff z (Yobs P z) β).mp h
  have hcard : (#zᶜ : ℝ) = (n : ℝ) - (#z : ℝ) := by rw [card_compl, Nat.cast_sub hlt.le]
  have hz0 : (0 : ℝ) < (#z : ℝ) := by exact_mod_cast hpos
  have hzc0 : (0 : ℝ) < (n : ℝ) - (#z : ℝ) := by
    have : (#z : ℝ) < (n : ℝ) := by exact_mod_cast hlt
    linarith
  have hsplit : (∑ i ∈ z, P.y1 i) + (∑ i ∈ zᶜ, P.y0 i) = ∑ i, Yobs P z i := by
    rw [Finset.sum_congr rfl fun i hi => (Yobs_of_mem hi).symm,
      Finset.sum_congr rfl fun i hi => (Yobs_of_not_mem (by simpa using hi)).symm]
    exact Finset.sum_add_sum_compl z _
  rw [Finset.sum_congr rfl fun i hi => Yobs_of_mem hi] at e1
  rw [← hsplit] at e0
  have hb0 : β 0 = (∑ i ∈ zᶜ, P.y0 i) / ((n : ℝ) - (#z : ℝ)) := by
    field_simp
    linarith
  have hb1 : β 1 = (∑ i ∈ z, P.y1 i) / (#z : ℝ) - β 0 := by
    field_simp
    linarith
  exact ⟨by rw [meanOn, hcard, hb0], by rw [DiM, hb1, hb0, hcard]⟩

Proof.

Write $n_{0} = n - n_{1}$ , $\bar{Z} = n_{1} / n$ , $\bar{Y} = \frac{1}{n} \sum_{i} Y_{i}$ , and let ${\bar{Y}}_{1} = \frac{1}{n_{1}} \sum_{i \in z} Y_{i}$ and ${\bar{Y}}_{0} = \frac{1}{n_{0}} \sum_{i \notin z} Y_{i}$ be the two group means, so that $n \bar{Y} = n_{1} {\bar{Y}}_{1} + n_{0} {\bar{Y}}_{0}$ . For a regression on a constant and one regressor the slope and intercept are

β_{1} = \frac{\sum_{i} (Z_{i} - \bar{Z}) (Y_{i} - \bar{Y})}{\sum_{i} (Z_{i} - \bar{Z})^{2}}, β_{0} = \bar{Y} - β_{1} \bar{Z},

and both collapse because the regressor is binary. Since $Z_{i}^{2} = Z_{i}$ ,

\sum_{i} (Z_{i} - \bar{Z})^{2} = \sum_{i} Z_{i}^{2} - n {\bar{Z}}^{2} = n_{1} - \frac{n_{1}^{2}}{n} = \frac{n_{1} n_{0}}{n},

which is nonzero exactly when both groups are non-empty. For the numerator, $\sum_{i} Z_{i} Y_{i} = \sum_{i \in z} Y_{i} = n_{1} {\bar{Y}}_{1}$ , so

\sum_{i} (Z_{i} - \bar{Z}) (Y_{i} - \bar{Y}) = \sum_{i} Z_{i} Y_{i} - n \bar{Z} \bar{Y} = n_{1} {\bar{Y}}_{1} - \frac{n_{1}}{n} (n_{1} {\bar{Y}}_{1} + n_{0} {\bar{Y}}_{0}) = \frac{n_{1} n_{0}}{n} ({\bar{Y}}_{1} - {\bar{Y}}_{0}) .

Dividing, $β_{1} = {\bar{Y}}_{1} - {\bar{Y}}_{0}$ ; and then

β_{0} = \bar{Y} - β_{1} \bar{Z} = \frac{n_{1} {\bar{Y}}_{1} + n_{0} {\bar{Y}}_{0}}{n} - \frac{n_{1}}{n} ({\bar{Y}}_{1} - {\bar{Y}}_{0}) = {\bar{Y}}_{0} .

Finally the observed outcome is $y_{i} (1)$ on the treated units and $y_{i} (0)$ on the controls, so ${\bar{Y}}_{1}$ is the treated mean of $y (1)$ , ${\bar{Y}}_{0}$ is the control mean of $y (0)$ , and $β_{1} = {\hat{τ}}_{DiM}$ .∎

The Lean never forms that ratio. It treats the two normal equations as a linear system in $(β_{0}, β_{1})$ : subtracting the treated equation from the total one eliminates $β_{1}$ and isolates $β_{0}$ , and the treated equation then gives $β_{1}$ . That is the same content in the other order — the covariance ratio above is what you get by eliminating $β_{0}$ first — and the two field_simp/linarith pairs are the two divisions.

obtain ⟨e0, e1⟩ := (isOLS_simpleMatrix_iff z (Yobs P z) β).mp h

Replace the hypothesis IsOLS by the two scalar equations of the previous proposition, named e0 (over everybody) and e1 (over the treated). From here on the proof is arithmetic with two equations in two unknowns.

Goal

e0 : ∑ i, Yobs P z i = ↑n * β 0 + ↑(#z) * β 1
e1 : ∑ i ∈ z, Yobs P z i = ↑(#z) * β 0 + ↑(#z) * β 1
⊢ β 0 = meanOn zᶜ P.y0 ∧ β 1 = DiM P z

```
have hcard : (#zᶜ : ℝ) = (n : ℝ) - (#z : ℝ) := by rw [card_compl, Nat.cast_sub hlt.le]
```
The control count as a real number. card_compl gives $n_{0} = n - n_{1}$ in $N$ , and Nat.cast_sub hlt.le is what allows that truncated subtraction to be cast to $R$ as an honest subtraction. This is the first of the two places hlt is needed.

have hz0 : (0 : ℝ) < (#z : ℝ) := by exact_mod_cast hpos

$n_{1} > 0$ in $R$ , so that the treated group’s denominator is nonzero.

```
have hzc0 : (0 : ℝ) < (n : ℝ) - (#z : ℝ) := by
  have : (#z : ℝ) < (n : ℝ) := by exact_mod_cast hlt
  linarith
```
And $n_{0} > 0$ . These two are the hypotheses field_simp will consult below before it is willing to clear a denominator; without them Lean’s $x / 0 = 0$ would make the statement false, as the pitfall at the end of the chapter shows.
```
have hsplit : (∑ i ∈ z, P.y1 i) + (∑ i ∈ zᶜ, P.y0 i) = ∑ i, Yobs P z i := by rw [Finset.sum_congr rfl fun i hi => (Yobs_of_mem hi).symm, Finset.sum_congr rfl fun i hi => (Yobs_of_not_mem (by simpa using hi)).symm]; exact Finset.sum_add_sum_compl z _
```
The bridge between potential and observed outcomes. Each Finset.sum_congr rewrites one of the two sums termwise: Yobs_of_mem says the observed outcome of a treated unit is $y_{i} (1)$ and Yobs_of_not_mem that a control’s is $y_{i} (0)$ , both used backwards so that both sums become sums of Yobs. Then Finset.sum_add_sum_compl glues a sum over $z$ and a sum over $z^{c}$ into a sum over everybody.
```
rw [Finset.sum_congr rfl fun i hi => Yobs_of_mem hi] at e1
```
Apply the same rewriting to the treated equation, so e1 now reads $\sum_{i \in z} y_{i} (1) = n_{1} β_{0} + n_{1} β_{1}$ .

rw [← hsplit] at e0

And to the total equation, in the other direction: e0 now reads $\sum_{i \in z} y_{i} (1) + \sum_{i \notin z} y_{i} (0) = n β_{0} + n_{1} β_{1}$ . The two equations are now stated in the same vocabulary as the goal.

Goal

e0 : ∑ i ∈ z, P.y1 i + ∑ i ∈ zᶜ, P.y0 i = ↑n * β 0 + ↑(#z) * β 1
e1 : ∑ i ∈ z, P.y1 i = ↑(#z) * β 0 + ↑(#z) * β 1
hcard : ↑(#zᶜ) = ↑n - ↑(#z)
hz0 : 0 < ↑(#z)
hzc0 : 0 < ↑n - ↑(#z)
⊢ β 0 = meanOn zᶜ P.y0 ∧ β 1 = DiM P z

```
have hb0 : β 0 = (∑ i ∈ zᶜ, P.y0 i) / ((n : ℝ) - (#z : ℝ)) := by field_simp; linarith
```
The intercept. field_simp, using hzc0, multiplies out the division and leaves $β_{0} (n - n_{1}) = \sum_{i \notin z} y_{i} (0)$ — which is exactly e0 minus e1, so linarith closes it. Eliminating $β_{1}$ between the two normal equations is what isolates the control mean.
Goal
```
⊢ β 0 * (↑n - ↑(#z)) = ∑ i ∈ zᶜ, P.y0 i
```
```
have hb1 : β 1 = (∑ i ∈ z, P.y1 i) / (#z : ℝ) - β 0 := by field_simp; linarith
```
The slope. field_simp, using hz0 this time, leaves $β_{1} n_{1} = \sum_{i \in z} y_{i} (1) - n_{1} β_{0}$ , which is e1 rearranged; linarith again. Note that $β_{1}$ comes out as (treated mean) $- β_{0}$ , and $β_{0}$ is the control mean, so this is the difference in means — but Lean still has to be shown that.
Goal
```
hb0 : β 0 = (∑ i ∈ zᶜ, P.y0 i) / (↑n - ↑(#z))
⊢ β 1 * ↑(#z) = ∑ i ∈ z, P.y1 i - ↑(#z) * β 0
```
```
exact ⟨by rw [meanOn, hcard, hb0], by rw [DiM, hb1, hb0, hcard]⟩
```
The two halves of the conclusion. For the first, unfold meanOn and use hcard to make its denominator $n - n_{1}$ syntactically the one in hb0. For the second, unfold DiM and substitute hb1, then hb0 for the $β_{0}$ it still contains, then hcard once more to line the denominators up.

The same holds for the closed form (ols_simpleMatrix_eq_DiM), because with both groups non-empty the determinant $n_{1} (n - n_{1})$ is nonzero.

Inheriting a sampling property

Now, and only now, does a design appear.

The one probabilistic theorem in the fileNew tabOpens with 998 lines of library preamble — the code above is at the bottom of the editor.

theorem ols_simpleMatrix_unbiased_completeRandomization {n₁ : ℕ} (hle : n₁ ≤ n)
    (P : Population n) (hpos : 0 < n₁) (hlt : n₁ < n) :
    (completeRandomization n n₁ hle).expect
      (fun z => ols (simpleMatrix z) (Yobs P z) 1) = P.tau := by
  have hsupp : ∀ z, (completeRandomization n n₁ hle).prob z ≠ 0 →
      ols (simpleMatrix z) (Yobs P z) 1 = DiM P z := by
    intro z hz
    have hcard : #z = n₁ := card_eq_of_prob_ne_zero hz
    exact ols_simpleMatrix_eq_DiM P (by omega) (by omega)
  rw [(completeRandomization n n₁ hle).expect_congr_of_support hsupp]
  exact DiM_unbiased_completeRandomization hle P hpos hlt

There is no new argument in that proof and there is not supposed to be. On every assignment the design can produce, $| z | = n_{1}$ , so both groups are non-empty, so the coefficient equals the difference in means; and two functions agreeing on the support of a design have the same expectation. Then quote the unbiasedness of difference in means from chapter 8 and stop.

intro z hz; have hcard : #z = n₁ := card_eq_of_prob_ne_zero hz

Inside hsupp: fix an assignment with nonzero probability. card_eq_of_prob_ne_zero is the chapter-7 fact that complete randomization puts all its mass on the assignments of size exactly $n_{1}$ .

Goal

z : Assignment n
hz : (completeRandomization n n₁ hle).prob z ≠ 0
hcard : #z = n₁
⊢ ols (simpleMatrix z) (Yobs P z) 1 = DiM P z

```
exact ols_simpleMatrix_eq_DiM P (by omega) (by omega)
```
The algebraic identity, applied to this assignment. Its two hypotheses are $0 < | z |$ and $| z | < n$ ; with hcard in context they follow from hpos and hlt by pure arithmetic, which is what omega does.
```
rw [(completeRandomization n n₁ hle).expect_congr_of_support hsupp]
```
expect_congr_of_support swaps the regression coefficient for ${\hat{τ}}_{DiM}$ under the expectation, which is legitimate precisely because they agree on the support. Every trace of regression has now left the goal.
Goal
```
⊢ (completeRandomization n n₁ hle).expect (DiM P) = P.tau
```
```
exact DiM_unbiased_completeRandomization hle P hpos hlt
```
And that goal is a theorem from chapter 8, quoted verbatim.

That is the whole message. Running lm(y ~ z) does not invoke a model of the outcomes, does not need homoskedasticity, and does not need normal errors. Under this design it recomputes an estimator whose sampling distribution is governed entirely by the randomization — and whose unbiasedness you proved three chapters ago.

Three units, in numbers

Identities are easier to believe once you have watched one happen. Here is a population of three units with both potential outcomes written down, which we can do because this is a thought experiment rather than data.

A population and a realized assignmentNew tabOpens with 1,216 lines of library preamble — the code above is at the bottom of the editor.

noncomputable def exP : Population 3 := ⟨![4, 6, 9], ![1, 2, 3]⟩

def exZ : Assignment 3 := {0, 1}

theorem exZ_card : #exZ = 2 := by decide

theorem exZ_compl : exZᶜ = {2} := by decide

The two decides attached to it evaluate small Finset facts in the kernel: $| z | = 2$ because $z = {0, 1}$ , and $z^{c} = {2}$ because the third unit is the only one left. Both are decidable statements about three explicit units, so “just compute it” is a proof.

The estimand and the estimate:

τ = 13/3, and the difference in means on this assignment is 2New tabOpens with 1,225 lines of library preamble — the code above is at the bottom of the editor.

theorem exP_tau : exP.tau = 13 / 3 := by
  rw [Population.tau_eq]
  norm_num [exP, Population.effect, Fin.sum_univ_three, Matrix.cons_val_two,
    Matrix.head_cons, Matrix.tail_cons]

theorem exP_DiM : DiM exP exZ = 2 := by
  rw [DiM, exZ_compl, exZ]
  norm_num [exP, Matrix.cons_val_two, Matrix.head_cons, Matrix.tail_cons]

```
rw [Population.tau_eq]
```
exP_tau. Population.tau_eq restates $τ$ as the mean of the unit-level effects, $\frac{1}{n} \sum_{i} (y_{i} (1) - y_{i} (0))$ , which is the form the numbers can be substituted into.
```
norm_num [exP, Population.effect, Fin.sum_univ_three, Matrix.cons_val_two, Matrix.head_cons, Matrix.tail_cons]
```
Substitute and compute. exP unfolds the population to its two literal vectors, Population.effect to $y_{i} (1) - y_{i} (0)$ , and Fin.sum_univ_three expands the sum over Fin 3 into three terms. The three Matrix.cons_* lemmas are what read an entry out of a ![…] literal — a vector literal is a nested cons, so getting at its third entry means peeling two of them. norm_num then does $\frac{(4 - 1) + (6 - 2) + (9 - 3)}{3} = \frac{13}{3}$ .
```
rw [DiM, exZ_compl, exZ]
```
exP_DiM. Unfold the estimator, then replace $z^{c}$ by the literal ${2}$ using the decided fact above and $z$ by the literal ${0, 1}$ , so that both sums are over explicit sets.
```
norm_num [exP, Matrix.cons_val_two, Matrix.head_cons, Matrix.tail_cons]
```
And compute: $\frac{4 + 6}{2} - \frac{3}{1} = 5 - 3 = 2$ . Note which potential outcome each group contributes — $y (1)$ for the treated units, $y (0)$ for the control — which is what makes this the difference in means and not a difference of comparable quantities.

$2 \neq 13 / 3$ . The estimator is unbiased and the estimate is wrong, which is the normal state of affairs and worth seeing spelled out by a proof assistant at least once.

ExerciseCheck a coefficient vector by hand (warm-up)

The observed outcomes are $(4, 6, 3)$ . Claim: $β = (3, 2)$ solves the normal equations. The two equations of isOLS_simpleMatrix_iff are $4 + 6 + 3 = 3 β_{0} + 2 β_{1}$ and $4 + 6 = 2 β_{0} + 2 β_{1}$ , i.e. $13 = 13$ and $10 = 10$ . Note where the numbers came from: $3$ is the control mean of $y (0)$ and $2$ is the difference in means.

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

theorem exP_isOLS : IsOLS (simpleMatrix exZ) (Yobs exP exZ) ![3, 2] := by
  sorry

Hint

Start with rw [isOLS_simpleMatrix_iff] and constructor; then in each goal simp with Yobs_eq_ite (which turns the observed outcome into an if), exZ, exP and Fin.sum_univ_three (which expands the sums over Fin 3), and norm_num does the arithmetic.

Show solution

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

theorem exP_isOLS : IsOLS (simpleMatrix exZ) (Yobs exP exZ) ![3, 2] := by
  rw [isOLS_simpleMatrix_iff]
  constructor <;>
    simp [Yobs_eq_ite, exZ, exP, Fin.sum_univ_three] <;> norm_num

And the closed form agreesNew tabOpens with 1,239 lines of library preamble — the code above is at the bottom of the editor.

theorem exP_ols : ols (simpleMatrix exZ) (Yobs exP exZ) 1 = 2 := by
  rw [ols_simpleMatrix_eq_DiM exP (by rw [exZ_card]; norm_num) (by rw [exZ_card]; norm_num),
    exP_DiM]

```
rw [ols_simpleMatrix_eq_DiM exP (by rw [exZ_card]; norm_num) (by rw [exZ_card]; norm_num), exP_DiM]
```
One rewrite with two side goals. ols_simpleMatrix_eq_DiM is the closed-form version of the headline theorem, so applying it turns the left-hand side into ${\hat{τ}}_{DiM}$ , and exP_DiM — just proved — evaluates that as $2$ . The two by blocks discharge its hypotheses $0 < | z |$ and $| z | < 3$ : exZ_card replaces $| z |$ by $2$ and norm_num checks $0 < 2$ and $2 < 3$ . Nothing is recomputed; the matrix inverse never appears.

Group indicators return group means

Replace the single dummy by one dummy per group and drop the intercept. This is the saturated, or fixed-effects, regression: with blocks crossed with treatment it is the model a blocked experiment is usually written as.

One column per groupNew tabOpens with 1,021 lines of library preamble — the code above is at the bottom of the editor.

def groupMatrix {K : ℕ} (g : Fin n → Fin K) : Matrix (Fin n) (Fin K) ℝ :=
  Matrix.of fun i k => if g i = k then 1 else 0

Every row has exactly one $1$ , so the columns are orthogonal and no matrix algebra is needed at all.

ExerciseThe fitted value is your own group's coefficient (warm-up)

The fitted value of unit i under the group-indicator design is the coefficient of i’s own group, β (g i).

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

theorem groupMatrix_mulVec {K : ℕ} (g : Fin n → Fin K) (β : Fin K → ℝ) (i : Fin n) :
    (groupMatrix g *ᵥ β) i = β (g i) := by
  sorry

Hint

One simp with Matrix.mulVec, dotProduct and groupMatrix in the set: the sum ∑ k, (if g i = k then 1 else 0) * β k collapses by Finset.sum_ite_eq, which simp will find for you.

Show solution

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

theorem groupMatrix_mulVec {K : ℕ} (g : Fin n → Fin K) (β : Fin K → ℝ) (i : Fin n) :
    (groupMatrix g *ᵥ β) i = β (g i) := by
  simp [Matrix.mulVec, dotProduct, groupMatrix]

ExerciseThe coefficients are the cell means (core)

If $β$ solves the normal equations of the no-intercept regression on the group dummies, then for every non-empty group $k$ , $β_{k}$ is the mean of $y$ over that group. The normal equation for column $k$ is $\sum_{i} 1 {g_{i} = k} (y_{i} - β_{g_{i}}) = 0$ ; the indicator restricts the sum to group $k$ , on which the fitted value is the constant $β_{k}$ .

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

theorem groupMatrix_ols_coeff {K : ℕ} (g : Fin n → Fin K) (y : Fin n → ℝ) (β : Fin K → ℝ)
    (h : IsOLS (groupMatrix g) y β) (k : Fin K) (hk : (groupOf g k).Nonempty) :
    β k = meanOn (groupOf g k) y := by
  sorry

Hint

Get the column equation with (isOLS_iff_orthogonal _ _ _).mp h k, prove a hstep rewriting each summand as if g i = k then y i - β k else 0 by groupMatrix_mulVec and by_cases hi : g i = k, then ← Finset.sum_filter turns the sum of an if into a sum over groupOf g k. After Finset.sum_sub_distrib and Finset.sum_const the equation reads $\sum_{i \in s} y_{i} - | s | \cdot β_{k} = 0$ , and eq_meanOn_of_sum_eq (fed Finset.card_pos.mpr hk) plus linarith converts that into $β_{k} = {\bar{y}}_{s}$ .

Show solution

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

theorem groupMatrix_ols_coeff {K : ℕ} (g : Fin n → Fin K) (y : Fin n → ℝ) (β : Fin K → ℝ)
    (h : IsOLS (groupMatrix g) y β) (k : Fin K) (hk : (groupOf g k).Nonempty) :
    β k = meanOn (groupOf g k) y := by
  have hk' := (isOLS_iff_orthogonal _ _ _).mp h k
  have hstep : ∀ i, groupMatrix g i k * (y i - (groupMatrix g *ᵥ β) i)
      = if g i = k then y i - β k else 0 := by
    intro i
    rw [groupMatrix_mulVec]
    by_cases hi : g i = k <;> simp [groupMatrix, hi]
  rw [Finset.sum_congr rfl fun i _ => hstep i, ← Finset.sum_filter,
    show Finset.univ.filter (fun i => g i = k) = groupOf g k from rfl,
    Finset.sum_sub_distrib, Finset.sum_const, nsmul_eq_mul] at hk'
  refine eq_meanOn_of_sum_eq (Finset.card_pos.mpr hk) ?_
  linarith

This is the sense in which a saturated regression is non-parametric: it imposes nothing, it reports the cell means. Combine it with chapter 10 and the story about blocked experiments writes itself — a block-by-treatment saturated regression reports the within-block treated and control means, and the blocked estimator is a particular weighted combination of them.

Lin’s fully interacted regression

Regress the outcome on $(1, Z, x, Z \cdot x)$ for a covariate $x$ . Lin (2013) recommends centering $x$ first; the coefficient on $Z$ is then his covariate-adjusted estimator.

The algebra that makes it tick is a change of basis. First, a name for “the simple regression of $y$ on $(1, x)$ , restricted to the units in $s$ ”:

A within-group simple regressionNew tabOpens with 1,056 lines of library preamble — the code above is at the bottom of the editor.

def IsSimpleOLSOn (s : Finset (Fin n)) (x y : Fin n → ℝ) (a b : ℝ) : Prop :=
  (∑ i ∈ s, (y i - (a + b * x i)) = 0) ∧ (∑ i ∈ s, x i * (y i - (a + b * x i)) = 0)

ExerciseThe fitted line passes through the center of mass (core)

The intercept of a within-group fit is ${\bar{y}}_{s} - b {\bar{x}}_{s}$ . Only the first normal equation is used — the one that says the residuals sum to zero. The slope equation is what would pin $b$ down, and we deliberately never need it, which is why $b$ appears in the conclusion as a free quantity.

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

theorem intercept_eq_of_isSimpleOLSOn {s : Finset (Fin n)} {x y : Fin n → ℝ} {a b : ℝ}
    (h : IsSimpleOLSOn s x y a b) (hs : 0 < #s) :
    a = meanOn s y - b * meanOn s x := by
  sorry

Hint

The first normal equation is h.1. Expand its $\sum_{i \in s} (y_{i} - (a + b x_{i}))$ with Finset.sum_sub_distrib, Finset.sum_add_distrib, Finset.sum_const (plus nsmul_eq_mul) and ← Finset.mul_sum, unfold both meanOns, and let field_simp and linarith divide by $| s |$ — field_simp needs (0 : ℝ) < #s as a have first.

Show solution

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

theorem intercept_eq_of_isSimpleOLSOn {s : Finset (Fin n)} {x y : Fin n → ℝ} {a b : ℝ}
    (h : IsSimpleOLSOn s x y a b) (hs : 0 < #s) :
    a = meanOn s y - b * meanOn s x := by
  have hs' : (0 : ℝ) < (#s : ℝ) := by exact_mod_cast hs
  have hexp : ∑ i ∈ s, (y i - (a + b * x i))
      = (∑ i ∈ s, y i) - ((#s : ℝ) * a + b * ∑ i ∈ s, x i) := by
    rw [Finset.sum_sub_distrib, Finset.sum_add_distrib, Finset.sum_const, nsmul_eq_mul,
      ← Finset.mul_sum]
  have h1 := h.1
  rw [hexp] at h1
  rw [meanOn, meanOn]
  field_simp
  linarith

And the design matrix:

Four columnsNew tabOpens with 1,073 lines of library preamble — the code above is at the bottom of the editor.

noncomputable def interactedMatrix (z : Assignment n) (x : Fin n → ℝ) :
    Matrix (Fin n) (Fin 4) ℝ :=
  Matrix.of fun i => ![1, Z z i, x i, Z z i * x i]

TheoremDecoupling

The fully interacted regression is two separate within-group regressions. $β$ solves the four normal equations if and only if the line $(β_{0} + β_{1}) + (β_{2} + β_{3}) x$ is a least-squares fit among the treated units and the line $β_{0} + β_{2} x$ is a least-squares fit among the controls.

Nothing is assumed about $y$ , about $x$ , or about the design. Both groups may even be empty; the statement is an equivalence of two systems of equations, and non-degeneracy is needed only when one wants to solve them.

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

theorem isOLS_interactedMatrix_iff (z : Assignment n) (x y : Fin n → ℝ) (β : Fin 4 → ℝ) :
    IsOLS (interactedMatrix z x) y β ↔
      IsSimpleOLSOn z x y (β 0 + β 1) (β 2 + β 3) ∧ IsSimpleOLSOn zᶜ x y (β 0) (β 2) := by
  rw [isOLS_iff_orthogonal, forall_fin_four]
  set r : Fin n → ℝ := fun i => y i - (interactedMatrix z x *ᵥ β) i with hr
  have treated : ∀ i, Z z i * r i = Z z i * (y i - ((β 0 + β 1) + (β 2 + β 3) * x i)) := by
    intro i
    simp only [hr, interactedMatrix_mulVec]
    linear_combination (-(β 1 + β 3 * x i)) * Z_mul_self z i
  have treatedx : ∀ i, Z z i * x i * r i
      = Z z i * (x i * (y i - ((β 0 + β 1) + (β 2 + β 3) * x i))) := by
    intro i
    simp only [hr, interactedMatrix_mulVec]
    linear_combination (-(x i * (β 1 + β 3 * x i))) * Z_mul_self z i
  have control : ∀ i, (1 - Z z i) * r i = (1 - Z z i) * (y i - (β 0 + β 2 * x i)) := by
    intro i
    simp only [hr, interactedMatrix_mulVec]
    linear_combination (β 1 + β 3 * x i) * Z_mul_self z i
  have controlx : ∀ i, (1 - Z z i) * (x i * r i)
      = (1 - Z z i) * (x i * (y i - (β 0 + β 2 * x i))) := by
    intro i
    simp only [hr, interactedMatrix_mulVec]
    linear_combination (x i * (β 1 + β 3 * x i)) * Z_mul_self z i
  have bz1 : ∑ i ∈ z, (y i - ((β 0 + β 1) + (β 2 + β 3) * x i)) = ∑ i, Z z i * r i := by
    rw [sum_mem_eq_sum_Z_mul]
    exact (Finset.sum_congr rfl fun i _ => treated i).symm
  have bz2 : ∑ i ∈ z, x i * (y i - ((β 0 + β 1) + (β 2 + β 3) * x i))
      = ∑ i, Z z i * x i * r i := by
    rw [sum_mem_eq_sum_Z_mul]
    exact (Finset.sum_congr rfl fun i _ => treatedx i).symm
  have bc1 : ∑ i ∈ zᶜ, (y i - (β 0 + β 2 * x i)) = (∑ i, r i) - ∑ i, Z z i * r i := by
    rw [sum_compl_eq_sum_one_sub_Z_mul, ← Finset.sum_sub_distrib]
    refine Finset.sum_congr rfl fun i _ => ?_
    rw [← control i]; ring
  have bc2 : ∑ i ∈ zᶜ, x i * (y i - (β 0 + β 2 * x i))
      = (∑ i, x i * r i) - ∑ i, Z z i * x i * r i := by
    rw [sum_compl_eq_sum_one_sub_Z_mul, ← Finset.sum_sub_distrib]
    refine Finset.sum_congr rfl fun i _ => ?_
    rw [← controlx i]; ring
  simp only [interactedMatrix_apply_zero, interactedMatrix_apply_one,
    interactedMatrix_apply_two, interactedMatrix_apply_three, one_mul, ← hr]
  rw [IsSimpleOLSOn, IsSimpleOLSOn, bz1, bz2, bc1, bc2]
  constructor
  · rintro ⟨h0, h1, h2, h3⟩
    exact ⟨⟨h1, h3⟩, ⟨by rw [h0, h1, sub_zero], by rw [h2, h3, sub_zero]⟩⟩
  · rintro ⟨⟨t1, t2⟩, ⟨c1, c2⟩⟩
    exact ⟨by linarith, t1, by linarith, t2⟩

Proof.

Write $r_{i} = y_{i} - (β_{0} + β_{1} Z_{i} + β_{2} x_{i} + β_{3} Z_{i} x_{i})$ for the residual. By the orthogonality form of the normal equations, $β$ is a least-squares fit exactly when

\sum_{i} r_{i} = 0, \sum_{i} Z_{i} r_{i} = 0, \sum_{i} x_{i} r_{i} = 0, \sum_{i} Z_{i} x_{i} r_{i} = 0,

i.e. exactly when $r$ is orthogonal to each of the four columns $1, Z, x, Z x$ . Now change basis. The four vectors $Z, Z x, 1 - Z, (1 - Z) x$ span the same space, since $1 = Z + (1 - Z)$ and $x = Z x + (1 - Z) x$ , and the change of basis is invertible. Orthogonality to a spanning set does not depend on which spanning set, so the four equations above hold if and only if

\sum_{i} Z_{i} r_{i} = 0, \sum_{i} Z_{i} x_{i} r_{i} = 0, \sum_{i} (1 - Z_{i}) r_{i} = 0, \sum_{i} (1 - Z_{i}) x_{i} r_{i} = 0.

The new basis is what makes the system block diagonal: $Z_{i}$ vanishes off the treated group and $1 - Z_{i}$ vanishes off the control group, so the first pair of equations involves only treated units and the second only controls.

It remains to identify each pair. Because $Z_{i}^{2} = Z_{i}$ ,

Z_{i} r_{i} = Z_{i} (y_{i} - [(β_{0} + β_{1}) + (β_{2} + β_{3}) x_{i}]),

so the first two equations say that the residuals of the line $(β_{0} + β_{1}) + (β_{2} + β_{3}) x$ sum to zero over the treated units and are orthogonal to $x$ there — the two normal equations of the simple regression of $y$ on $(1, x)$ restricted to $z$ . Symmetrically $(1 - Z_{i}) r_{i} = (1 - Z_{i}) (y_{i} - [β_{0} + β_{2} x_{i}])$ , so the last two say the same for the line $β_{0} + β_{2} x$ over the controls.∎

The Lean proof is that argument in two halves. Four haves do the collapse of the fitted value — two on the treated side, two on the control side, each one an application of $Z_{i}^{2} = Z_{i}$ — and four more convert between “sum over a group” and “indicator-weighted sum over everybody”, which is what puts the change of basis in a form rw can use. The final constructor is the invertibility of that change of basis, one direction at a time.

rw [isOLS_iff_orthogonal, forall_fin_four]

Trade the matrix equation for orthogonality to each column, then split the Fin 4 quantifier into its four instances. forall_fin_four is a private lemma in the file rather than fin_cases, because fin_cases leaves the index as an anonymous constructor and the @[simp] lemmas about interactedMatrix z x i 0 only fire on the numeral.

Goal

z : Assignment n
x y : Fin n → ℝ
β : Fin 4 → ℝ
⊢ ∑ i, interactedMatrix z x i 0 * (y i - (interactedMatrix z x *ᵥ β) i) = 0 ∧
      ∑ i, interactedMatrix z x i 1 * (y i - (interactedMatrix z x *ᵥ β) i) = 0 ∧
        ∑ i, interactedMatrix z x i 2 * (y i - (interactedMatrix z x *ᵥ β) i) = 0 ∧
          ∑ i, interactedMatrix z x i 3 * (y i - (interactedMatrix z x *ᵥ β) i) = 0 ↔
    IsSimpleOLSOn z x y (β 0 + β 1) (β 2 + β 3) ∧ IsSimpleOLSOn zᶜ x y (β 0) (β 2)

```
set r : Fin n → ℝ := fun i => y i - (interactedMatrix z x *ᵥ β) i with hr
```
Name the residual. This is bookkeeping — set abbreviates the term everywhere it appears and records the abbreviation as hr — but it is what makes the eight sum identities below readable.
```
have treated : ∀ i, Z z i * r i = Z z i * (y i - ((β 0 + β 1) + (β 2 + β 3) * x i)) := by intro i; simp only [hr, interactedMatrix_mulVec]; linear_combination (-(β 1 + β 3 * x i)) * Z_mul_self z i
```
The first of the four idempotence rewrites, and the heart of the theorem: on a treated unit the four-parameter fitted value collapses to the two-parameter line $(β_{0} + β_{1}) + (β_{2} + β_{3}) x_{i}$ . After interactedMatrix_mulVec expands the fitted value, the two sides differ by $(β_{1} + β_{3} x_{i}) (Z_{i}^{2} - Z_{i})$ , and linear_combination closes exactly such a goal when handed that multiple of Z_mul_self z i.
Goal
```
i : Fin n
⊢ Z z i * (y i - (β 0 + Z z i * β 1 + x i * β 2 + Z z i * x i * β 3))
    = Z z i * (y i - (β 0 + β 1 + (β 2 + β 3) * x i))
```

have treatedx : ∀ i, Z z i * x i * r i = Z z i * (x i * (y i - ((β 0 + β 1) + (β 2 + β 3) * x i))) := by … linear_combination (-(x i * (β 1 + β 3 * x i))) * Z_mul_self z i

The same statement with an extra factor of $x_{i}$ , for the second treated normal equation. The scalar in the linear_combination picks up the same $x_{i}$ .

```
have control : ∀ i, (1 - Z z i) * r i = (1 - Z z i) * (y i - (β 0 + β 2 * x i)) := by … linear_combination (β 1 + β 3 * x i) * Z_mul_self z i
```
The control-side mirror: on a control unit the fitted value collapses to $β_{0} + β_{2} x_{i}$ . The scalar is the same as in treated with the sign flipped, because $1 - Z_{i}$ replaces $Z_{i}$ .
```
have controlx : ∀ i, (1 - Z z i) * (x i * r i) = (1 - Z z i) * (x i * (y i - (β 0 + β 2 * x i))) := by … linear_combination (x i * (β 1 + β 3 * x i)) * Z_mul_self z i
```
And the same with the extra $x_{i}$ . Those are the four linear_combination calls; each is “multiply $Z^{2} = Z$ by the right scalar”, and together they are the step in the informal proof where the four-parameter fitted value collapses to a two-parameter line inside each group.
```
have bz1 : ∑ i ∈ z, (y i - ((β 0 + β 1) + (β 2 + β 3) * x i)) = ∑ i, Z z i * r i := by rw [sum_mem_eq_sum_Z_mul]; exact (Finset.sum_congr rfl fun i _ => treated i).symm
```
The first sum conversion. The left-hand side is a normal equation of the within-group regression, a sum over $z$ ; sum_mem_eq_sum_Z_mul rewrites it as an indicator-weighted sum over everybody, and then treated, applied termwise and backwards, turns each summand into $Z_{i} r_{i}$ . This is what puts the two sides of the eventual iff in the same language.
Goal
```
⊢ ∑ i, Z z i * (y i - (β 0 + β 1 + (β 2 + β 3) * x i)) = ∑ i, Z z i * r i
```

have bz2 : ∑ i ∈ z, x i * (y i - ((β 0 + β 1) + (β 2 + β 3) * x i)) = ∑ i, Z z i * x i * r i := by rw [sum_mem_eq_sum_Z_mul]; exact (Finset.sum_congr rfl fun i _ => treatedx i).symm

The same for the second treated equation, using treatedx.

```
have bc1 : ∑ i ∈ zᶜ, (y i - (β 0 + β 2 * x i)) = (∑ i, r i) - ∑ i, Z z i * r i := by rw [sum_compl_eq_sum_one_sub_Z_mul, ← Finset.sum_sub_distrib]; refine Finset.sum_congr rfl fun i _ => ?_; rw [← control i]; ring
```
The control side, and the one place the shape differs. sum_compl_eq_sum_one_sub_Z_mul turns the sum over $z^{c}$ into $\sum_{i} (1 - Z_{i}) (\dots)$ , and Finset.sum_sub_distrib reversed writes the target difference of two sums as one sum of differences. Comparing termwise, control i reversed replaces the left summand by $(1 - Z_{i}) r_{i}$ and ring checks that this is $r_{i} - Z_{i} r_{i}$ . The result is stated as a difference because the four goals coming out of forall_fin_four are about the old columns $1$ and $Z$ , not about $1 - Z$ .
Goal
```
⊢ ∑ i, (1 - Z z i) * (y i - (β 0 + β 2 * x i)) = ∑ i, (r i - Z z i * r i)
```

have bc2 : ∑ i ∈ zᶜ, x i * (y i - (β 0 + β 2 * x i)) = (∑ i, x i * r i) - ∑ i, Z z i * x i * r i := by …

The same for the second control equation, using controlx.

```
simp only [interactedMatrix_apply_zero, interactedMatrix_apply_one, interactedMatrix_apply_two, interactedMatrix_apply_three, one_mul, ← hr]
```
Replace the four columns of the design matrix by $1$ , $Z_{i}$ , $x_{i}$ and $Z_{i} x_{i}$ — those are the four @[simp] lemmas — clear the multiplication by the constant column with one_mul, and fold the residual back up into r with hr reversed.
```
rw [IsSimpleOLSOn, IsSimpleOLSOn, bz1, bz2, bc1, bc2]
```
Unfold both within-group conditions into their two normal equations each, then rewrite all four with the conversions just proved. Both sides of the iff are now statements about the same four sums, and the only difference is the basis: ${1, Z, x, Z x}$ on the left, ${Z, Z x, 1 - Z, (1 - Z) x}$ — written as differences — on the right. (Lean prints the three left-hand residuals in full rather than as r i; they are folded up below to keep the two bases legible.)
Goal
```
⊢ univ.sum r = 0 ∧ ∑ i, Z z i * r i = 0 ∧ ∑ i, x i * r i = 0
      ∧ ∑ i, Z z i * x i * r i = 0 ↔
    (∑ i, Z z i * r i = 0 ∧ ∑ i, Z z i * x i * r i = 0) ∧
      ∑ i, r i - ∑ i, Z z i * r i = 0 ∧ ∑ i, x i * r i - ∑ i, Z z i * x i * r i = 0
```
```
constructor
· rintro ⟨h0, h1, h2, h3⟩
  exact ⟨⟨h1, h3⟩, ⟨by rw [h0, h1, sub_zero], by rw [h2, h3, sub_zero]⟩⟩
```
Left to right. Of the four hypotheses, h1 and h3 (the $Z$ and $Z x$ equations) are literally the treated pair. For the control pair, rw [h0, h1, sub_zero] substitutes $0$ for both sums in $\sum_{i} r_{i} - \sum_{i} Z_{i} r_{i}$ and sub_zero finishes; likewise h2, h3 for the $x$ version.
```
· rintro ⟨⟨t1, t2⟩, ⟨c1, c2⟩⟩
  exact ⟨by linarith, t1, by linarith, t2⟩
```
Right to left. t1 and t2 are again two of the four goals outright; the other two, the $1$ and $x$ equations, follow by adding: $\sum_{i} r_{i} = (\sum_{i} r_{i} - \sum_{i} Z_{i} r_{i}) + \sum_{i} Z_{i} r_{i} = c_{1} + t_{1} = 0$ , which is linarith’s job. That the change of basis is invertible is exactly this addition.

Once you see the decoupling, the coefficient on $Z$ is visibly the gap between two fitted intercepts, and each intercept is a center of mass.

Corollary (Lin’s coefficient). If $β$ solves the interacted normal equations and both groups are non-empty, then

β_{1} = ({\bar{y}}_{1} - (β_{2} + β_{3}) {\bar{x}}_{1}) - ({\bar{y}}_{0} - β_{2} {\bar{x}}_{0}),

where ${\bar{y}}_{1}, {\bar{x}}_{1}$ are treated-group means and ${\bar{y}}_{0}, {\bar{x}}_{0}$ control-group means.

Proof.

By the theorem, $(β_{0} + β_{1}, β_{2} + β_{3})$ is a least-squares fit on the treated units and $(β_{0}, β_{2})$ one on the controls. A fitted line passes through its group’s center of mass, so its intercept is ${\bar{y}}_{s} - b {\bar{x}}_{s}$ ; applying that to each group gives $β_{0} + β_{1} = {\bar{y}}_{1} - (β_{2} + β_{3}) {\bar{x}}_{1}$ and $β_{0} = {\bar{y}}_{0} - β_{2} {\bar{x}}_{0}$ . Subtract.∎

Lin's coefficient is a difference of interceptsNew tabOpens with 1,144 lines of library preamble — the code above is at the bottom of the editor.

theorem lin_coeff (z : Assignment n) (x y : Fin n → ℝ) (β : Fin 4 → ℝ)
    (hpos : 0 < #z) (hlt : #z < n) (h : IsOLS (interactedMatrix z x) y β) :
    β 1 = (meanOn z y - (β 2 + β 3) * meanOn z x) - (meanOn zᶜ y - β 2 * meanOn zᶜ x) := by
  obtain ⟨htreat, hctrl⟩ := (isOLS_interactedMatrix_iff z x y β).mp h
  have hzc : 0 < #zᶜ := by rw [card_compl]; omega
  have h1 := intercept_eq_of_isSimpleOLSOn htreat hpos
  have h0 := intercept_eq_of_isSimpleOLSOn hctrl hzc
  linarith

```
obtain ⟨htreat, hctrl⟩ := (isOLS_interactedMatrix_iff z x y β).mp h
```
Split the interacted fit into the two within-group fits, using the decoupling theorem in the forward direction. Everything after this is about two simple regressions.
```
have hzc : 0 < #zᶜ := by rw [card_compl]; omega
```
The control group is non-empty. card_compl gives $| z^{c} | = n - | z |$ and omega derives $0 < n - | z |$ from hlt; this is needed because the center-of-mass lemma divides by the group size.

have h1 := intercept_eq_of_isSimpleOLSOn htreat hpos

The treated intercept as a center of mass: $β_{0} + β_{1} = {\bar{y}}_{1} - (β_{2} + β_{3}) {\bar{x}}_{1}$ . Only the first of the two within-group normal equations is used, which is why the slope $β_{2} + β_{3}$ survives in the conclusion as a free quantity rather than being solved for.

Goal

htreat : IsSimpleOLSOn z x y (β 0 + β 1) (β 2 + β 3)
hctrl : IsSimpleOLSOn zᶜ x y (β 0) (β 2)
hzc : 0 < #zᶜ
h1 : β 0 + β 1 = meanOn z y - (β 2 + β 3) * meanOn z x
⊢ β 1 = meanOn z y - (β 2 + β 3) * meanOn z x - (meanOn zᶜ y - β 2 * meanOn zᶜ x)

have h0 := intercept_eq_of_isSimpleOLSOn hctrl hzc

And the control intercept, $β_{0} = {\bar{y}}_{0} - β_{2} {\bar{x}}_{0}$ .

```
linarith
```
Subtract the second from the first. The whole content of the corollary is that subtraction.

Rearranged, the same identity says what adjustment is doing:

ExerciseDifference in means, minus a covariate correction (stretch)

$β_{1} = ({\bar{y}}_{1} - {\bar{y}}_{0}) - (b_{1} {\bar{x}}_{1} - b_{0} {\bar{x}}_{0})$ . The first bracket is the raw difference in observed group means; the second is what the fitted slopes say the covariate imbalance alone would have produced.

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

theorem lin_coeff_eq_DiM_sub_adjustment (z : Assignment n) (x y : Fin n → ℝ) (β : Fin 4 → ℝ)
    (hpos : 0 < #z) (hlt : #z < n) (h : IsOLS (interactedMatrix z x) y β) :
    β 1 = (meanOn z y - meanOn zᶜ y)
      - ((β 2 + β 3) * meanOn z x - β 2 * meanOn zᶜ x) := by
  sorry

Hint

One have := lin_coeff z x y β hpos hlt h and a linarith — the exercise is reading the statement, not proving it.

Show solution

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

theorem lin_coeff_eq_DiM_sub_adjustment (z : Assignment n) (x y : Fin n → ℝ) (β : Fin 4 → ℝ)
    (hpos : 0 < #z) (hlt : #z < n) (h : IsOLS (interactedMatrix z x) y β) :
    β 1 = (meanOn z y - meanOn zᶜ y)
      - ((β 2 + β 3) * meanOn z x - β 2 * meanOn zᶜ x) := by
  have := lin_coeff z x y β hpos hlt h
  linarith

Both are stated for an arbitrary response vector y. Feeding them the observed outcome vector puts the identity in the vocabulary of the rest of the tutorial.

ExerciseGroup means of the observed outcome (warm-up)

On the treated units the observed outcome is $y (1)$ , and meanOn z only looks at those units, so the two means coincide. The treated case is done for you; prove the control-group mirror image.

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

theorem meanOn_Yobs_treated (P : Population n) (z : Assignment n) :
    meanOn z (Yobs P z) = meanOn z P.y1 := by
  rw [meanOn, meanOn, Finset.sum_congr rfl fun i hi => Yobs_of_mem hi]

theorem meanOn_Yobs_control (P : Population n) (z : Assignment n) :
    meanOn zᶜ (Yobs P z) = meanOn zᶜ P.y0 := by
  sorry

Hint

The same three-step rewrite as the treated case — meanOn, meanOn, then Finset.sum_congr rfl with Yobs_of_not_mem — except that i ∈ zᶜ has to become i ∉ z before Yobs_of_not_mem applies. simpa using hi does that.

Show solution

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

theorem meanOn_Yobs_treated (P : Population n) (z : Assignment n) :
    meanOn z (Yobs P z) = meanOn z P.y1 := by
  rw [meanOn, meanOn, Finset.sum_congr rfl fun i hi => Yobs_of_mem hi]

theorem meanOn_Yobs_control (P : Population n) (z : Assignment n) :
    meanOn zᶜ (Yobs P z) = meanOn zᶜ P.y0 := by
  rw [meanOn, meanOn, Finset.sum_congr rfl fun i hi => Yobs_of_not_mem (by simpa using hi)]

Lin's coefficient against the potential outcomesNew tabOpens with 1,252 lines of library preamble — the code above is at the bottom of the editor.

theorem lin_coeff_Yobs (P : Population n) (z : Assignment n) (x : Fin n → ℝ) (β : Fin 4 → ℝ)
    (hpos : 0 < #z) (hlt : #z < n) (h : IsOLS (interactedMatrix z x) (Yobs P z) β) :
    β 1 = DiM P z - ((β 2 + β 3) * meanOn z x - β 2 * meanOn zᶜ x) := by
  rw [lin_coeff_eq_DiM_sub_adjustment z x (Yobs P z) β hpos hlt h,
    meanOn_Yobs_treated, meanOn_Yobs_control, DiM_eq_meanOn_sub_meanOn]

rw [lin_coeff_eq_DiM_sub_adjustment z x (Yobs P z) β hpos hlt h, …]

Apply the rearranged corollary at $y = Y^{obs}$ . The right-hand side now mentions the group means of the observed outcome; the goal is to recognize their difference as ${\hat{τ}}_{DiM}$ , which is stated in terms of the potential outcomes.

Goal

⊢ meanOn z (Yobs P z) - meanOn zᶜ (Yobs P z)
      - ((β 2 + β 3) * meanOn z x - β 2 * meanOn zᶜ x)
    = DiM P z - ((β 2 + β 3) * meanOn z x - β 2 * meanOn zᶜ x)

```
rw […, meanOn_Yobs_treated, meanOn_Yobs_control, …]
```
The two warm-up lemmas above: on the treated units the observed outcome is $y (1)$ , so its mean over $z$ is the mean of $y (1)$ over $z$ ; on the controls it is $y (0)$ .
```
rw […, DiM_eq_meanOn_sub_meanOn]
```
And DiM_eq_meanOn_sub_meanOn says ${\hat{τ}}_{DiM}$ is exactly that difference of two subgroup means — it holds by rfl, since that is what the estimator is defined to be. The two sides are now identical.

ExercisePerfect balance means no adjustment (stretch)

If both within-group covariate means are zero, Lin’s coefficient is the plain difference in means, whatever the fitted slopes turn out to be.

The result shows that covariate adjustment is not a correction for a bias that exists on average — under complete randomization there is none. It is a correction for the imbalance realized in the assignment you actually drew. centred_covariate_means in the library adds the other half of Lin’s recommendation: with $\sum_{i} x_{i} = 0$ the two group means are tied together by $n_{1} {\bar{x}}_{1} + n_{0} {\bar{x}}_{0} = 0$ , so a single scalar imbalance drives the whole correction.

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

theorem lin_coeff_eq_DiM_of_balanced (P : Population n) (z : Assignment n) (x : Fin n → ℝ)
    (β : Fin 4 → ℝ) (hpos : 0 < #z) (hlt : #z < n)
    (h : IsOLS (interactedMatrix z x) (Yobs P z) β)
    (hx1 : meanOn z x = 0) (hx0 : meanOn zᶜ x = 0) :
    β 1 = DiM P z := by
  sorry

Hint

Rewrite with lin_coeff_Yobs P z x β hpos hlt h, then with the two balance hypotheses hx1 and hx0; what is left is closed by ring.

Show solution

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

theorem lin_coeff_eq_DiM_of_balanced (P : Population n) (z : Assignment n) (x : Fin n → ℝ)
    (β : Fin 4 → ℝ) (hpos : 0 < #z) (hlt : #z < n)
    (h : IsOLS (interactedMatrix z x) (Yobs P z) β)
    (hx1 : meanOn z x = 0) (hx0 : meanOn zᶜ x = 0) :
    β 1 = DiM P z := by
  rw [lin_coeff_Yobs P z x β hpos hlt h, hx1, hx0]
  ring

Reading autoformalized Lean

Three ways this chapter’s results get mis-stated. All three type-check.

Dropping the non-degeneracy hypotheses

theorem ols_coeff_is_dim (P : Population n) (z : Assignment n) :
    ols (simpleMatrix z) (Yobs P z) 1 = DiM P z := by
  sorry

Autoformalization pitfall

0 < #z and #z < n look like the sort of side condition a formaliser drops because “division by zero is fine anyway”. It is not fine here, and the reason is that both sides are total functions that disagree at the degenerate point.

Take $z =$ Finset.univ. The left side is $0$ , by the exercise at the top of this chapter: the Gram matrix is singular and Mathlib’s total inverse returns zero. The right side is DiM P univ $= \bar{y} (1) - 0 / 0 = \bar{y} (1)$ , the population mean of the treated potential outcome, which is not zero in general. So the statement above is false, not merely uninformative — and nothing about its appearance tells you that. The tell is the missing hypothesis, and the way to check is to ask what each side evaluates to where the mathematics would have said “undefined”.

Confusing an identity with unbiasedness

theorem regression_recovers_the_ate {n₁ : ℕ} (hle : n₁ ≤ n) (P : Population n)
    (z : Assignment n) (hz : #z = n₁) :
    ols (simpleMatrix z) (Yobs P z) 1 = P.tau := by
  sorry

A name that promises more than the statement delivers

theorem lin_estimator_is_efficient (z : Assignment n) (x y : Fin n → ℝ)
    (β : Fin 4 → ℝ) (hpos : 0 < #z) (hlt : #z < n)
    (h : IsOLS (interactedMatrix z x) y β) :
    β 1 = (meanOn z y - meanOn zᶜ y)
      - ((β 2 + β 3) * meanOn z x - β 2 * meanOn zᶜ x) := by
  sorry

The habit this chapter is trying to build is narrower than “read Lean carefully”. It is: when a formalization claims a regression coefficient is an estimator, check the hypotheses that keep both sides non-degenerate; and when it claims a regression coefficient has a property, check that the property is stated about a design, because algebra alone can never give you one.