Panel.Estimand­Characterization.Staggered­TWFEDecomposition.Pairwise

lemma weighted_center_cov_uncentered_right {ι : Type*} [Fintype ι] (p x y : ι → ℝ) (hp : ∑ i, p i = 1) : ∑ i, p i * (x i - ∑ j, p j * x j) * y i = (1 / 2) * ∑ i, ∑ j, p i * p j * (x i - x j) * (y i - y j) := by classical let mx := ∑ j, p j * x j let my := ∑ j, p j * y j have hzero : ∑ i, p i * (x i - mx) = 0 := by calc ∑ i, p i * (x i - mx) = (∑ i, p i * x i) - ∑ i, p i * mx := by simp [mul_sub, Finset.sum_sub_distrib] _ = mx - (∑ i, p i) * mx := by simp [Finset.sum_mul, mx] _ = 0 := by rw [hp] ring have hmy : ∑ i, p i * (x i - mx) * my = 0 := by calc ∑ i, p i * (x i - mx) * my = (∑ i, p i * (x i - mx)) * my := by rw [Finset.sum_mul] _ = 0 := by rw [hzero] ring have hcenter : ∑ i, p i * (x i - mx) * y i = ∑ i, p i * (x i - mx) * (y i - my) := by calc ∑ i, p i * (x i - mx) * y i = ∑ i, (p i * (x i - mx) * (y i - my) + p i * (x i - mx) * my) := by refine Finset.sum_congr rfl ?_ intro i _hi ring _ = ∑ i, p i * (x i - mx) * (y i - my) + ∑ i, p i * (x i - mx) * my := by rw [Finset.sum_add_distrib] _ = ∑ i, p i * (x i - mx) * (y i - my) := by rw [hmy] ring calc ∑ i, p i * (x i - ∑ j, p j * x j) * y i = ∑ i, p i * (x i - mx) * y i := by rfl _ = ∑ i, p i * (x i - mx) * (y i - my) := hcenter _ = ∑ i, p i * (x i - ∑ j, p j * x j) * (y i - ∑ j, p j * y j) := by rfl _ = (1 / 2) * ∑ i, ∑ j, p i * p j * (x i - x j) * (y i - y j) := weighted_center_cov p x y hp (by norm_num)

lemma sum_weight_over_T_commute (P : CohortPanel 𝒢 T) (f : 𝒢 → Fin T → ℝ) : (∑ g, ∑ t, (P.p g / (T : ℝ)) * f g t) = (T : ℝ)⁻¹ * ∑ t, ∑ g, P.p g * f g t := by classical rw [Finset.sum_comm] simp [div_eq_mul_inv, Finset.mul_sum, mul_left_comm, mul_comm]

lemma pairwise_sum_normalize (P : CohortPanel 𝒢 T) (f : 𝒢 → 𝒢 → Fin T → ℝ) : (T : ℝ)⁻¹ * ∑ t, ((1 / 2) * ∑ g, ∑ u, P.p g * P.p u * f g u t) = ∑ g, ∑ u, (P.p g * P.p u / (2 * (T : ℝ))) * ∑ t, f g u t := by classical calc (T : ℝ)⁻¹ * ∑ t, ((1 / 2) * ∑ g, ∑ u, P.p g * P.p u * f g u t) = (T : ℝ)⁻¹ * ∑ t, ∑ g, ∑ u, (1 / 2) * (P.p g * P.p u * f g u t) := by simp [Finset.mul_sum] _ = (T : ℝ)⁻¹ * ∑ g, ∑ u, ∑ t, (1 / 2) * (P.p g * P.p u * f g u t) := by congr 1 rw [Finset.sum_comm] refine Finset.sum_congr rfl ?_ intro g _hg rw [Finset.sum_comm] _ = ∑ g, ∑ u, (P.p g * P.p u / (2 * (T : ℝ))) * ∑ t, f g u t := by simp [div_eq_mul_inv, Finset.mul_sum, Finset.sum_mul, mul_assoc, mul_left_comm, mul_comm]

lemma D_eq_zero_of_isInf (P : CohortPanel 𝒢 T) {u : 𝒢} (hu : AdoptionPath.isInfinite (P.A u)) (t : Fin T) : D P u t = 0 := by classical unfold D rw [show P.A u = ⊤ from hu] simp

lemma binary_time_variance (hT_pos : 0 < T) (x : Fin T → ℝ) (hx : ∀ t, x t ^ 2 = x t) : (T : ℝ)⁻¹ * ∑ t, (x t - ((T : ℝ)⁻¹ * ∑ t, x t)) ^ 2 = ((T : ℝ)⁻¹ * ∑ t, x t) * (1 - ((T : ℝ)⁻¹ * ∑ t, x t)) := by classical let m : ℝ := (T : ℝ)⁻¹ * ∑ t, x t have hTne : (T : ℝ) ≠ 0 := by exact_mod_cast (ne_of_gt hT_pos) have hsum_sq : ∑ t, x t ^ 2 = ∑ t, x t := by exact Finset.sum_congr rfl (by intro t _ht; exact hx t) have hsum_expand : ∑ t, (x t ^ 2 - 2 * m * x t + m ^ 2) = ∑ t, x t - 2 * m * ∑ t, x t + (T : ℝ) * m ^ 2 := by calc ∑ t, (x t ^ 2 - 2 * m * x t + m ^ 2) = ∑ t, x t ^ 2 - ∑ t, 2 * m * x t + ∑ _t : Fin T, m ^ 2 := by simp [Finset.sum_sub_distrib, Finset.sum_add_distrib] _ = ∑ t, x t - 2 * m * ∑ t, x t + (T : ℝ) * m ^ 2 := by rw [hsum_sq] simp [Finset.mul_sum, Fintype.card_fin] calc (T : ℝ)⁻¹ * ∑ t, (x t - ((T : ℝ)⁻¹ * ∑ t, x t)) ^ 2 = (T : ℝ)⁻¹ * ∑ t, (x t ^ 2 - 2 * m * x t + m ^ 2) := by congr 1 refine Finset.sum_congr rfl ?_ intro t _ht simp [m] ring _ = ((T : ℝ)⁻¹ * ∑ t, x t) * (1 - ((T : ℝ)⁻¹ * ∑ t, x t)) := by rw [hsum_expand] simp [m] field_simp [hTne] ring

lemma binary_time_cov_filter_mean (hT_pos : 0 < T) (p : Fin T → Prop) [DecidablePred p] (z : Fin T → ℝ) (hcard1 : ((Finset.univ.filter p).card : ℝ) ≠ 0) : (T : ℝ)⁻¹ * ∑ t, ((if p t then (1 : ℝ) else 0) - ((T : ℝ)⁻¹ * ∑ t, (if p t then (1 : ℝ) else 0))) * z t = ((T : ℝ)⁻¹ * ∑ t, (if p t then (1 : ℝ) else 0)) * (1 - ((T : ℝ)⁻¹ * ∑ t, (if p t then (1 : ℝ) else 0))) * (((Finset.univ.filter p).card : ℝ)⁻¹ * (∑ t ∈ (Finset.univ.filter p), z t) - ((Finset.univ.filter (fun t => ¬ p t)).card : ℝ)⁻¹ * (∑ t ∈ (Finset.univ.filter (fun t => ¬ p t)), z t)) := by classical have hTne : (T : ℝ) ≠ 0 := by exact_mod_cast (ne_of_gt hT_pos) let A : ℝ := ((Finset.univ.filter p).card : ℝ) let B : ℝ := ((Finset.univ.filter (fun t => ¬ p t)).card : ℝ) let Z1 : ℝ := ∑ t ∈ (Finset.univ.filter p), z t let Z0 : ℝ := ∑ t ∈ (Finset.univ.filter (fun t => ¬ p t)), z t have hA : A ≠ 0 := by simpa [A] using hcard1 have hsum_if : (∑ t, (if p t then (1 : ℝ) else 0)) = A := by simp [A] have hsum_if_z : (∑ t, (if p t then (1 : ℝ) else 0) * z t) = Z1 := by simp [Z1, Finset.sum_filter] have hsum_z : (∑ t, z t) = Z1 + Z0 := by simp [Z1, Z0, Finset.sum_filter] rw [← Finset.sum_add_distrib] refine Finset.sum_congr rfl ?_ intro t _ by_cases ht : p t <;> simp [ht] have hcard_total : A + B = (T : ℝ) := by have hnat : (Finset.univ.filter p).card + (Finset.univ.filter (fun t => ¬ p t)).card = T := by calc (Finset.univ.filter p).card + (Finset.univ.filter (fun t => ¬ p t)).card = ((Finset.univ.filter p) ∪ (Finset.univ.filter (fun t => ¬ p t))).card := by rw [Finset.card_union_of_disjoint] simp [Finset.disjoint_left] _ = T := by have hunion : (Finset.univ.filter p) ∪ (Finset.univ.filter (fun t => ¬ p t)) = (Finset.univ : Finset (Fin T)) := by ext t simp [em] simp [hunion] simpa [A, B] using (show (((Finset.univ.filter p).card : ℝ) + ((Finset.univ.filter (fun t => ¬ p t)).card : ℝ) = (T : ℝ)) by exact_mod_cast hnat) have hmain : (T : ℝ)⁻¹ * (Z1 - ((T : ℝ)⁻¹ * A) * (Z1 + Z0)) = ((T : ℝ)⁻¹ * A) * (1 - ((T : ℝ)⁻¹ * A)) * (A⁻¹ * Z1 - B⁻¹ * Z0) := by by_cases hB : B = 0 · have hZT0 : Z0 = 0 := by have hempty : Finset.univ.filter (fun t => ¬ p t) = ∅ := by apply Finset.card_eq_zero.mp have hB' : ((Finset.univ.filter (fun t => ¬ p t)).card : ℝ) = 0 := by simpa [B] using hB exact_mod_cast hB' simp [Z0, hempty] have hAeq : A = (T : ℝ) := by linarith rw [hZT0, hAeq] field_simp [hTne] ring · have hABne : A + B ≠ 0 := by rw [hcard_total] exact hTne rw [← hcard_total] field_simp [hA, hB, hABne] ring calc (T : ℝ)⁻¹ * ∑ t, ((if p t then (1 : ℝ) else 0) - ((T : ℝ)⁻¹ * ∑ t, (if p t then (1 : ℝ) else 0))) * z t = (T : ℝ)⁻¹ * (Z1 - ((T : ℝ)⁻¹ * A) * (Z1 + Z0)) := by rw [hsum_if] calc (T : ℝ)⁻¹ * ∑ t, ((if p t then (1 : ℝ) else 0) - (T : ℝ)⁻¹ * A) * z t = (T : ℝ)⁻¹ * (∑ t, (if p t then (1 : ℝ) else 0) * z t - ∑ t, ((T : ℝ)⁻¹ * A) * z t) := by congr 1 rw [← Finset.sum_sub_distrib] refine Finset.sum_congr rfl ?_ intro t _ ring _ = (T : ℝ)⁻¹ * (Z1 - ((T : ℝ)⁻¹ * A) * (Z1 + Z0)) := by rw [hsum_if_z, ← Finset.mul_sum, hsum_z] _ = ((T : ℝ)⁻¹ * ∑ t, (if p t then (1 : ℝ) else 0)) * (1 - ((T : ℝ)⁻¹ * ∑ t, (if p t then (1 : ℝ) else 0))) * (((Finset.univ.filter p).card : ℝ)⁻¹ * (∑ t ∈ (Finset.univ.filter p), z t) - ((Finset.univ.filter (fun t => ¬ p t)).card : ℝ)⁻¹ * (∑ t ∈ (Finset.univ.filter (fun t => ¬ p t)), z t)) := by rw [hmain, hsum_if]

lemma TN_treated_window_card_ne_zero (P : CohortPanel 𝒢 T) {g : 𝒢} (hg : AdoptionPath.isFinite (P.A g)) : ((Finset.univ.filter (fun t => AdoptionPath.le (P.A g) t)).card : ℝ) ≠ 0 := by cases hA : P.A g with | top => exact False.elim (hg hA) | coe a => have hmem : a ∈ Finset.univ.filter (fun t => AdoptionPath.le (P.A g) t) := by simp [AdoptionPath.le, hA] have hpos : 0 < (Finset.univ.filter (fun t => AdoptionPath.le (P.A g) t)).card := Finset.card_pos.mpr ⟨a, hmem⟩ have hne : ((Finset.univ.filter (fun t => AdoptionPath.le (P.A g) t)).card : ℝ) ≠ 0 := by exact_mod_cast (ne_of_gt hpos) simpa [hA] using hne

lemma TT_middle_window_card_ne_zero (P : CohortPanel 𝒢 T) {e ℓ : 𝒢} (hord : P.A e < P.A ℓ) : ((Finset.univ.filter (fun t => AdoptionPath.le (P.A e) t ∧ AdoptionPath.lt (P.A ℓ) t)).card : ℝ) ≠ 0 := by cases he : P.A e with | top => have hbad := hord simp [he] at hbad | coe a => cases hℓA : P.A ℓ with | top => have hmem : a ∈ Finset.univ.filter (fun t => AdoptionPath.le (P.A e) t ∧ AdoptionPath.lt (P.A ℓ) t) := by simp [AdoptionPath.le, AdoptionPath.lt, he, hℓA] have hpos : 0 < (Finset.univ.filter (fun t => AdoptionPath.le (P.A e) t ∧ AdoptionPath.lt (P.A ℓ) t)).card := Finset.card_pos.mpr ⟨a, hmem⟩ have hne : ((Finset.univ.filter (fun t => AdoptionPath.le (P.A e) t ∧ AdoptionPath.lt (P.A ℓ) t)).card : ℝ) ≠ 0 := by exact_mod_cast (ne_of_gt hpos) simpa [he, hℓA] using hne | coe b => have hab : (a : WithTop (Fin T)) < (b : WithTop (Fin T)) := by simpa [he, hℓA] using hord have habFin : a < b := by exact (WithTop.coe_lt_coe (a := b) (b := a)).mp hab have hmem : a ∈ Finset.univ.filter (fun t => AdoptionPath.le (P.A e) t ∧ AdoptionPath.lt (P.A ℓ) t) := by simp [AdoptionPath.le, AdoptionPath.lt, he, hℓA, habFin] have hpos : 0 < (Finset.univ.filter (fun t => AdoptionPath.le (P.A e) t ∧ AdoptionPath.lt (P.A ℓ) t)).card := Finset.card_pos.mpr ⟨a, hmem⟩ have hne : ((Finset.univ.filter (fun t => AdoptionPath.le (P.A e) t ∧ AdoptionPath.lt (P.A ℓ) t)).card : ℝ) ≠ 0 := by exact_mod_cast (ne_of_gt hpos) simpa [he, hℓA] using hne

lemma Dtilde_eq_centeredD_sub_weighted_mean (P : CohortPanel 𝒢 T) (g : 𝒢) (t : Fin T) : Dtilde P g t = centeredD P g t - ∑ h, P.p h * centeredD P h t := by rw [weighted_centeredD_mean P t] rw [Dtilde_eq] unfold centeredD ring

lemma twfe_numerator_eq_pairwise_centeredD_Y (P : CohortPanel 𝒢 T) : (∑ g, ∑ t, (P.p g / (T : ℝ)) * Dtilde P g t * P.Y g t) = ∑ g, ∑ u, numPairContribution P g u := by classical calc (∑ g, ∑ t, (P.p g / (T : ℝ)) * Dtilde P g t * P.Y g t) = ∑ g, ∑ t, (P.p g / (T : ℝ)) * (Dtilde P g t * P.Y g t) := by refine Finset.sum_congr rfl ?_ intro g _hg refine Finset.sum_congr rfl ?_ intro t _ht ring _ = (T : ℝ)⁻¹ * ∑ t, ∑ g, P.p g * (Dtilde P g t * P.Y g t) := by exact sum_weight_over_T_commute P (fun g t => Dtilde P g t * P.Y g t) _ = (T : ℝ)⁻¹ * ∑ t, ∑ g, P.p g * (centeredD P g t - ∑ h, P.p h * centeredD P h t) * P.Y g t := by congr 1 refine Finset.sum_congr rfl ?_ intro t _ht refine Finset.sum_congr rfl ?_ intro g _hg rw [Dtilde_eq_centeredD_sub_weighted_mean] ring _ = (T : ℝ)⁻¹ * ∑ t, ((1 / 2) * ∑ g, ∑ u, P.p g * P.p u * (centeredD P g t - centeredD P u t) * (P.Y g t - P.Y u t)) := by congr 1 refine Finset.sum_congr rfl ?_ intro t _ht exact weighted_center_cov_uncentered_right (fun g => P.p g) (fun g => centeredD P g t) (fun g => P.Y g t) P.p_sum_one _ = ∑ g, ∑ u, numPairContribution P g u := by unfold numPairContribution simpa [mul_assoc] using pairwise_sum_normalize P (fun g u t => (centeredD P g t - centeredD P u t) * (P.Y g t - P.Y u t))

lemma disjoint_union_mean_eq_card_weighted_mean {α : Type*} [DecidableEq α] (A B : Finset α) (z : α → ℝ) (hdisj : Disjoint A B) (hB : ((B.card : ℝ) ≠ 0)) : (((A ∪ B).card : ℝ)⁻¹ * ∑ t ∈ A ∪ B, z t) = ((A.card : ℝ) / ((A ∪ B).card : ℝ)) * ((A.card : ℝ)⁻¹ * ∑ t ∈ A, z t) + (1 - ((A.card : ℝ) / ((A ∪ B).card : ℝ))) * ((B.card : ℝ)⁻¹ * ∑ t ∈ B, z t) := by classical let a : ℝ := A.card let b : ℝ := B.card let ZA : ℝ := ∑ t ∈ A, z t let ZB : ℝ := ∑ t ∈ B, z t have hsum : ∑ t ∈ A ∪ B, z t = ZA + ZB := by simp [ZA, ZB, Finset.sum_union hdisj] have hcard : ((A ∪ B).card : ℝ) = a + b := by have hnat : (A ∪ B).card = A.card + B.card := Finset.card_union_of_disjoint hdisj simp [a, b, hnat] have hb : b ≠ 0 := by simpa [b] using hB have hbpos : 0 < b := by have hBnat : B.card ≠ 0 := by exact_mod_cast hB dsimp [b] exact_mod_cast (Nat.pos_of_ne_zero hBnat) have hab : a + b ≠ 0 := by intro h have ha_nonneg : 0 ≤ a := by dsimp [a] exact_mod_cast (zero_le A.card) linarith by_cases ha : a = 0 · have hAempty : A = ∅ := by apply Finset.card_eq_zero.mp have : (A.card : ℝ) = 0 := by simpa [a] using ha exact Nat.cast_eq_zero.mp this have hZA : ZA = 0 := by simp [ZA, hAempty] rw [hsum, hcard, hZA, ha] simp [hAempty] field_simp [hb] ring · have haA : ((A.card : ℝ) ≠ 0) := by simpa [a] using ha have hAB : ((A.card : ℝ) + (B.card : ℝ)) ≠ 0 := by simpa [a, b] using hab rw [hsum, hcard] field_simp [haA, hB, hAB] simp [a, b, ZA, ZB] ring_nf

lemma TN_pair_vd_contribution_eq_gap (P : CohortPanel 𝒢 T) {g u : 𝒢} (hu : AdoptionPath.isInfinite (P.A u)) : vdPairContribution P g u + vdPairContribution P u g = P.p g * P.p u * q P g u * (1 - q P g u) := by rw [vdPairContribution_add_swap] have hTne : (T : ℝ) ≠ 0 := by exact_mod_cast (ne_of_gt P.T_pos) have hvar := TN_time_variance_eq_gap P (g := g) hu calc (P.p g * P.p u / (T : ℝ)) * ∑ t, (centeredD P g t - centeredD P u t)^2 = P.p g * P.p u * ((T : ℝ)⁻¹ * ∑ t, (centeredD P g t - centeredD P u t)^2) := by field_simp [hTne] _ = P.p g * P.p u * q P g u * (1 - q P g u) := by rw [hvar] ring

lemma TT_pair_vd_contribution_eq_gap (P : CohortPanel 𝒢 T) {e ℓ : 𝒢} (hord : P.A e < P.A ℓ) : vdPairContribution P e ℓ + vdPairContribution P ℓ e = P.p e * P.p ℓ * q P e ℓ * (1 - q P e ℓ) := by rw [vdPairContribution_add_swap] have hTne : (T : ℝ) ≠ 0 := by exact_mod_cast (ne_of_gt P.T_pos) have hvar := TT_time_variance_eq_gap P hord calc (P.p e * P.p ℓ / (T : ℝ)) * ∑ t, (centeredD P e t - centeredD P ℓ t)^2 = P.p e * P.p ℓ * ((T : ℝ)⁻¹ * ∑ t, (centeredD P e t - centeredD P ℓ t)^2) := by field_simp [hTne] _ = P.p e * P.p ℓ * q P e ℓ * (1 - q P e ℓ) := by rw [hvar] ring

lemma TT_pair_contribution_eq_lambda_delta_sum (P : CohortPanel 𝒢 T) {e ℓ : 𝒢} (hord : P.A e < P.A ℓ) (hℓ : AdoptionPath.isFinite (P.A ℓ)) : numPairContribution P e ℓ + numPairContribution P ℓ e = lambdaEL P e ℓ * Δ_EL P e ℓ + lambdaLE P e ℓ * Δ_LE P e ℓ := by rw [numPairContribution_add_swap] have hTne : (T : ℝ) ≠ 0 := by exact_mod_cast (ne_of_gt P.T_pos) have hcov := TT_time_cov_eq_lambda_delta_core P hord hℓ calc (P.p e * P.p ℓ / (T : ℝ)) * ∑ t, (centeredD P e t - centeredD P ℓ t) * (P.Y e t - P.Y ℓ t) = P.p e * P.p ℓ * ((T : ℝ)⁻¹ * ∑ t, (centeredD P e t - centeredD P ℓ t) * (P.Y e t - P.Y ℓ t)) := by field_simp [hTne] _ = lambdaEL P e ℓ * Δ_EL P e ℓ + lambdaLE P e ℓ * Δ_LE P e ℓ := by rw [hcov] unfold lambdaEL lambdaLE ring