fix: Trying to fix Canonica Ensemble Two State

PhysLean/StatisticalMechanics/CanonicalEnsemble/TwoState.lean

@@ -19,91 +19,79 @@ namespace CanonicalEnsemble
 open Temperature
 open Real MeasureTheory
 
-TODO "EVJNH" "Generalize the results for the two-state canonical ensemble so that the two
-  states have arbitrary energies, rather than one state having energy `0`."
-
 /-- The canonical ensemble corresponding to state system, with one state of energy
-  `E` and the other state of energy `0`. -/
-noncomputable def twoState (E : ℝ) : CanonicalEnsemble (Fin 2) where
-  energy := fun | 0 => 0 | 1 => E
+  `E₀` and the other state of energy `E₁`. -/
+noncomputable def twoState (E₀ E₁ : ℝ) : CanonicalEnsemble (Fin 2) where
+  energy := fun | 0 => E₀ | 1 => E₁
   dof := 0
   μ := Measure.count
   energy_measurable := by fun_prop
 
-instance {E} : IsFinite (twoState E) where
+instance {E₀ E₁} : IsFinite (twoState E₀ E₁) where
   μ_eq_count := rfl
   dof_eq_zero := rfl
   phase_space_unit_eq_one := rfl
 
-lemma twoState_partitionFunction_apply_eq_one_add_exp (E : ℝ) (T : Temperature) :
-    (twoState E).partitionFunction T = 1 + exp (- β T * E) := by
+lemma twoState_partitionFunction_apply (E₀ E₁ : ℝ) (T : Temperature) :
+    (twoState E₀ E₁).partitionFunction T = exp (- β T * E₀) + exp (- β T * E₁) := by
   rw [partitionFunction_of_fintype, twoState]
-  simp
+  simp [Fin.sum_univ_two]
 
-lemma twoState_partitionFunction_apply_eq_cosh (E : ℝ) (T : Temperature) :
-    (twoState E).partitionFunction T = 2 * exp (- β T * E / 2) * cosh (β T * E / 2) := by
-  rw [twoState_partitionFunction_apply_eq_one_add_exp, Real.cosh_eq]
-  field_simp
-  simp only [mul_add, ← exp_add, neg_add_cancel, exp_zero, add_right_inj, exp_eq_exp]
+lemma twoState_partitionFunction_apply_eq_cosh (E₀ E₁ : ℝ) (T : Temperature) :
+    (twoState E₀ E₁).partitionFunction T =
+    2 * exp (- β T * (E₀ + E₁) / 2) * cosh (β T * (E₁ - E₀) / 2) := by
+  rw [twoState_partitionFunction_apply, Real.cosh_eq]
   field_simp
-  ring
+  simp only [mul_add, ← exp_add]
+  ring_nf
 
 @[simp]
-lemma twoState_energy_fst (E : ℝ) : (twoState E).energy 0 = 0 := by
+lemma twoState_energy_fst (E₀ E₁ : ℝ) : (twoState E₀ E₁).energy 0 = E₀ := by
   rfl
 
 @[simp]
-lemma twoState_energy_snd (E : ℝ) : (twoState E).energy 1 = E := by
+lemma twoState_energy_snd (E₀ E₁ : ℝ) : (twoState E₀ E₁).energy 1 = E₁ := by
   rfl
 
-/-- Probability of the excited (energy `E`) state in closed form. -/
-lemma twoState_probability_snd (E : ℝ) (T : Temperature) :
-    (twoState E).probability T 1 = 1 / 2 * (1 - Real.tanh (β T * E / 2)) := by
-  have h_basic :
-      (twoState E).probability T 1 =
-        Real.exp (-β T * E) / (1 + Real.exp (-β T * E)) := by
-    -- The mathematical partition function of `twoState` is `1 + e^{-βE}`.
-    have hZ :
-        (twoState E).mathematicalPartitionFunction T =
-          1 + Real.exp (-β T * E) := by
-      rw [mathematicalPartitionFunction_of_fintype]
-      simp [twoState, Fin.sum_univ_two]
-    simp [probability, hZ]
-  set x : ℝ := β T * E with hx
-  have h_sym :
-      Real.exp (-x) / (1 + Real.exp (-x)) =
-        Real.exp (-x / 2) / (Real.exp (x / 2) + Real.exp (-x / 2)) := by
-    calc
-      _ = (Real.exp (-x) * Real.exp (x / 2)) / ((1 + Real.exp (-x)) * Real.exp (x / 2)) := by
-        field_simp
-      _ = _ := by
-        congr
-        · rw [← Real.exp_add]; ring_nf
-        · rw [add_mul, one_mul, ← Real.exp_add]; ring_nf
-  have h_tanh (y : ℝ) :
-      1 / 2 * (1 - Real.tanh y) = Real.exp (-y) / (Real.exp y + Real.exp (-y)) := by
-    rw [Real.tanh_eq_sinh_div_cosh, Real.sinh_eq, Real.cosh_eq, Real.exp_neg]
-    field_simp
-    ring
-  have h_half :
-      Real.exp (-x / 2) / (Real.exp (x / 2) + Real.exp (-x / 2)) =
-        1 / 2 * (1 - Real.tanh (x / 2)) := by
-    rw [h_tanh]
-    ring_nf
-  calc
-    (twoState E).probability T 1
-        = Real.exp (-x) / (1 + Real.exp (-x)) := by rw [hx, h_basic]; ring_nf
-    _ = Real.exp (-x / 2) / (Real.exp (x / 2) + Real.exp (-x / 2)) := h_sym
-    _ = 1 / 2 * (1 - Real.tanh (x / 2)) := h_half
-    _ = 1 / 2 * (1 - Real.tanh (β T * E / 2)) := by rw [hx]
-
-lemma twoState_meanEnergy_eq (E : ℝ) (T : Temperature) :
-    (twoState E).meanEnergy T = E / 2 * (1 - tanh (β T * E / 2)) := by
-  calc
-    _ = ∑ i : Fin 2, (twoState E).energy i * (twoState E).probability T i :=
-      meanEnergy_of_fintype (twoState E) T
-    _ = E * (twoState E).probability T 1 := by simp [twoState]
-  rw [twoState_probability_snd]
+/-- Probability of the first state (energy `E₀`) in closed form. -/
+lemma twoState_probability_fst (E₀ E₁ : ℝ) (T : Temperature) :
+    (twoState E₀ E₁).probability T 0 = 1 / 2 * (1 + Real.tanh (β T * (E₁ - E₀) / 2)) := by
+  set x := β T * (E₁ - E₀) / 2
+  set C := β T * (E₀ + E₁) / 2
+  have hE0 : - β T * E₀ = x - C := by
+    simp [x, C]; ring
+  have hE1 : - β T * E₁ = -x - C := by
+    simp [x, C]; ring
+  rw [probability, mathematicalPartitionFunction_of_fintype]
+  simp only [twoState, Fin.sum_univ_two, Fin.isValue]
+  rw [hE0, hE1]
+  rw [Real.tanh_eq_sinh_div_cosh, Real.sinh_eq, Real.cosh_eq]
+  simp only [Real.exp_sub, Real.exp_neg]
+  field_simp
+  ring
+
+/-- Probability of the second state (energy `E₁`) in closed form. -/
+lemma twoState_probability_snd (E₀ E₁ : ℝ) (T : Temperature) :
+    (twoState E₀ E₁).probability T 1 = 1 / 2 * (1 - Real.tanh (β T * (E₁ - E₀) / 2)) := by
+  set x := β T * (E₁ - E₀) / 2
+  set C := β T * (E₀ + E₁) / 2
+  have hE0 : - β T * E₀ = x - C := by
+    simp [x, C]; ring
+  have hE1 : - β T * E₁ = -x - C := by
+    simp [x, C]; ring
+  rw [probability, mathematicalPartitionFunction_of_fintype]
+  simp only [twoState, Fin.sum_univ_two, Fin.isValue]
+  rw [hE0, hE1]
+  rw [Real.tanh_eq_sinh_div_cosh, Real.sinh_eq, Real.cosh_eq]
+  simp only [Real.exp_sub, Real.exp_neg]
+  field_simp
+  ring
+
+lemma twoState_meanEnergy_eq (E₀ E₁ : ℝ) (T : Temperature) :
+    (twoState E₀ E₁).meanEnergy T =
+    (E₀ + E₁) / 2 - (E₁ - E₀) / 2 * Real.tanh (β T * (E₁ - E₀) / 2) := by
+  rw [meanEnergy_of_fintype]
+  simp [Fin.sum_univ_two, twoState_probability_fst, twoState_probability_snd]
   ring
 
 /-- A simplification of the `entropy` of the two-state canonical ensemble. -/