Regression as linear algebra

/-- Three facts about Mathlib's `Matrix` that a statistician needs before reading any of the
regression file.

* `Xᵀ` is the transpose and `X * Y` is matrix multiplication, so `Xᵀ * X` is the Gram matrix,
  whose `(j, k)` entry is the inner product of columns `j` and `k` — the "sums of squares and
  cross products" matrix.
* `X *ᵥ v` (`Matrix.mulVec`) is the matrix–vector product, whose `i`-th entry is the fitted
  value `∑ j, X i j * v j`.
* `*ᵥ` binds **tighter** than `*`.  `Xᵀ * X *ᵥ β` therefore parses as `Xᵀ * (X *ᵥ β)`, which
  is a matrix times a vector and does not even type-check.  The Gram matrix applied to a
  coefficient vector must be written `(Xᵀ * X) *ᵥ β`, and that is exactly what `IsOLS` says. -/
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

/-- Literal matrices and vectors.  `![a, b, …]` is a vector (a function out of `Fin k`) and
`!![a, b; c, d]` is a matrix given row by row; `Matrix.of` relabels a plain function
`Fin n → Fin k → ℝ` as a matrix, which is how `simpleMatrix` is built.

The two computations below are the entire content of `simpleMatrix_gram_det` in miniature:
a `2 × 2` determinant is `ad - bc`, and `Matrix.det_fin_two_of` is the lemma that says so. -/
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]

/-- `IsOLS X y β` says that `β` **solves the normal equations** `XᵀX β = Xᵀy`.

This is the definition of "an ordinary least squares fit" that requires no inverses, no
minimisation, and no rank condition: it is a plain linear equation, and it may have zero, one
or many solutions.  (It always has at least one, and it is exactly the first-order condition
for minimising `‖y - Xβ‖²`, but we will not need either fact.)

Everything in this file is phrased in terms of `IsOLS` rather than a formula, precisely
because the interesting theorems — "the coefficient equals the difference in means" — hold
for *any* solution, whether or not the design matrix has full rank. -/
def IsOLS {k : ℕ} (X : Matrix (Fin n) (Fin k) ℝ) (y : Fin n → ℝ) (β : Fin k → ℝ) : Prop :=
  (Xᵀ * X) *ᵥ β = Xᵀ *ᵥ y

/-- The **textbook OLS formula** `β̂ = (XᵀX)⁻¹Xᵀy`.

Mathlib's `Matrix.inv` is total: for a singular matrix it returns `0`.  So this definition
always type-checks, but it only *means* anything when `XᵀX` is nonsingular; that hypothesis
appears explicitly in every theorem about it below. -/
noncomputable def ols {k : ℕ} (X : Matrix (Fin n) (Fin k) ℝ) (y : Fin n → ℝ) : Fin k → ℝ :=
  (Xᵀ * X)⁻¹ *ᵥ (Xᵀ *ᵥ y)

/-- The design matrix of the **simple regression on a treatment dummy**: column `0` is the
intercept (all ones), column `1` is the treatment indicator `Z z`.

Small matrices are built in Lean with the vector notation `![a, b, …]`: here each *row* of the
matrix is the two-entry vector `![1, Z z i]`, and `Matrix.of` re-labels the function
`Fin n → Fin 2 → ℝ` as a matrix. -/
noncomputable def simpleMatrix (z : Assignment n) : Matrix (Fin n) (Fin 2) ℝ :=
  Matrix.of fun i => ![1, Z z i]

/-- **The two normal equations of the simple regression, written out.**

Orthogonality to the column of ones says the residuals sum to zero over the *whole* sample;
orthogonality to the treatment column says they sum to zero over the *treated* units.  In
terms of the coefficients:

* `∑ i, y i = n · β₀ + n₁ · β₁`
* `∑ i ∈ z, y i = n₁ · β₀ + n₁ · β₁`

Everything about simple regression follows from this 2 × 2 system.  Note that it is a genuine
`↔`: no rank condition is needed to unpack the normal equations, only to solve them. -/
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]
  -- Expand the two column equations into sums of `y` over everybody and over the treated.
  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
    -- the only content is `Z² = Z`
    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 OLS treatment coefficient is unbiased for the ATE under complete randomisation.**

This is the one theorem in the file with probabilistic content, and it contains no new
argument: on every assignment the design can produce, the regression coefficient equals the
difference in means (`ols_simpleMatrix_eq_DiM`), and the difference in means was already
shown to be unbiased (`DiM_unbiased_completeRandomization`).

That is the whole message of the chapter.  Running `lm(y ~ z)` does not invoke a model of the
outcomes, does not require homoskedasticity, and does not require normal errors; under this
design it just recomputes an estimator whose sampling distribution is governed entirely by
the randomisation. -/
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

/-- A finite population of three units, with both potential outcomes written down — which we
can do here because this is a thought experiment, not data.

Unit `0` has `(y₁, y₀) = (4, 1)`, unit `1` has `(6, 2)`, unit `2` has `(9, 3)`. -/
noncomputable def exP : Population 3 := ⟨![4, 6, 9], ![1, 2, 3]⟩

/-- The realised assignment: units `0` and `1` treated, unit `2` control. -/
def exZ : Assignment 3 := {0, 1}

theorem exZ_card : #exZ = 2 := by decide

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

/-- The estimand.  The average treatment effect of this population is `(3 + 4 + 6)/3 = 13/3`,
a number no experiment will ever observe. -/
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]

/-- The estimate.  On this assignment the observed outcomes are `(4, 6, 3)`, so the
difference in means is `(4 + 6)/2 - 3/1 = 2`.

`2 ≠ 13/3`: an unbiased estimator is not a correct one, and on a single assignment there is
nothing to be done about that. -/
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]

/-- And the closed-form estimator agrees, because both groups are non-empty and therefore the
Gram matrix `!![3, 2; 2, 2]` (determinant `2`) is invertible.

The proof does not invert anything: `ols_simpleMatrix_eq_DiM` has already done the algebra
once and for all, so all that is left is arithmetic on the observed outcomes. -/
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]

/-- The design matrix of the **group-indicator regression** (no intercept): entry `(i, k)` is
`1` when unit `i` belongs to group `k` and `0` otherwise.

Every row has exactly one `1`, so the columns are orthogonal — which is why the coefficients
turn out to be the group means with no matrix algebra at all. -/
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

/-- `IsSimpleOLSOn s x y a b` says that the line `a + b·x` solves the normal equations of the
simple regression of `y` on `(1, x)` **restricted to the units in `s`**:

* the residuals sum to zero over `s`;
* the residuals are orthogonal to `x` over `s`.

This is the same pair of equations as `isOLS_simpleMatrix_iff`, written for a subgroup and an
arbitrary regressor instead of a dummy. -/
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)

/-- The design matrix of **Lin's fully interacted regression**: the four columns are the
intercept, the treatment dummy, the covariate, and their product. -/
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]

/-- **Lin's coefficient is the difference of the two within-group fitted intercepts.**

For any solution `β` of the interacted normal equations, with both groups non-empty,

`β 1 = (ȳ₁ - b₁·x̄₁) - (ȳ₀ - b₀·x̄₀)`,

where `b₁ = β 2 + β 3` and `b₀ = β 2` are the within-treated and within-control slopes and
the bars denote within-group means.

In words: fit a line to the treated units, fit a line to the control units, and report the
gap between the two fitted *intercepts*.  When `x` has been centred so that `x̄ = 0`, the
intercepts are the two fitted lines evaluated at the grand mean of the covariate, and the gap
is exactly Lin's covariate-adjusted estimate of the ATE.

This is a finite-sample algebraic identity, true for every assignment.  It says nothing about
bias, consistency, or variance; see `lin_coeff_eq_DiM_sub_adjustment` for the form that makes
the connection to the unadjusted difference in means visible. -/
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

/-- **Lin's coefficient, written against the potential outcomes.**

`lin_coeff_eq_DiM_sub_adjustment` is stated for an arbitrary response vector `y`.  Feeding it
the observed outcome vector turns the two group means of `y` into the difference in means of
the population, so the identity reads

`β₁ = DiM − (b₁·x̄₁ − b₀·x̄₀)`,

with `b₁ = β 2 + β 3` the within-treated slope and `b₀ = β 2` the within-control slope.

Lin's estimator is the difference in means, minus what the two fitted lines say the covariate
imbalance alone would have produced.  This is an identity about the realised assignment; it
is not a bias, variance or consistency claim. -/
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]