Projective.lean

import QuantumComputing.Measurement.Generalized

/-!
# Projective Measurements

Projective measurements generated by vector projectors and their bridge to
generalized measurements.
-/

namespace QuantumComputing

namespace Measurement

/-- Projection operators associated with an arbitrary family of state vectors. -/
noncomputable def projectiveOperators {n outcomes : ℕ}
    (u : Fin outcomes → Vector n) : Fin outcomes → Square n :=
  fun i => Matrix.proj (u i)

/-- A family of vectors defines a complete projective measurement when its projectors are complete. -/
def IsProjectiveComplete {n outcomes : ℕ} (u : Fin outcomes → Vector n) : Prop :=
  IsComplete (projectiveOperators u)

/-- Probability for a projective measurement generated by vector projectors. -/
noncomputable def projProb {n outcomes : ℕ}
    (u : Fin outcomes → Vector n) (s : Vector n) (i : Fin outcomes) : ℝ :=
  ((s† ⬝ Matrix.proj (u i) ⬝ s) 0 0).re

/-- Post-measurement state for a projective measurement generated by vector projectors. -/
noncomputable def projPostMeasure {n outcomes : ℕ}
    (u : Fin outcomes → Vector n) (s : Vector n) (i : Fin outcomes) : Vector n :=
  ((1 / Real.sqrt (projProb u s i) : ℝ) : ℂ) • (Matrix.proj (u i) ⬝ s)

/--
Matrix whose rows are the adjoints of a projective-measurement vector family.
Applying it before a computational-basis measurement gives the same outcome
probabilities as measuring the original vector projectors.
-/
noncomputable def projectiveMeasurementOp {n : ℕ} (u : Fin n → Vector n) : Square n :=
  fun i j => (u i)† 0 j

/--
Restore operator for the `i`-th computational-basis outcome after using
`projectiveMeasurementOp`.
-/
noncomputable def projectiveRestoreOp {n : ℕ} (u : Fin n → Vector n) (i : Fin n) : Square n :=
  u i ⬝ (Vector.basis i)†

/-- Complete projective measurement generated by a family of state-vector projectors. -/
structure Projective (n outcomes : ℕ) where
  vector : Fin outcomes → Vector n
  isNormalized : ∀ i, Vector.IsNormalized (vector i)
  isComplete : IsProjectiveComplete vector

@[simp]
theorem projectiveOperators_apply {n outcomes : ℕ}
    (u : Fin outcomes → Vector n) (i : Fin outcomes) :
    projectiveOperators u i = Matrix.proj (u i) :=
  rfl

theorem generalizedProb_projectiveOperators_of_isNormalized {n outcomes : ℕ}
    {u : Fin outcomes → Vector n} (s : Vector n) {i : Fin outcomes}
    (hu : Vector.IsNormalized (u i)) :
    generalizedProb (projectiveOperators u) s i = projProb u s i := by
  simp [generalizedProb, projProb, projectiveOperators, Matrix.proj_mul_proj_of_isNormalized hu]

theorem generalizedPostMeasure_projectiveOperators_of_isNormalized {n outcomes : ℕ}
    {u : Fin outcomes → Vector n} (s : Vector n) {i : Fin outcomes}
    (hu : Vector.IsNormalized (u i)) :
    generalizedPostMeasure (projectiveOperators u) s i = projPostMeasure u s i := by
  simp [generalizedPostMeasure, projPostMeasure, projectiveOperators,
    generalizedProb_projectiveOperators_of_isNormalized s hu]

theorem projectiveOperators_eq_projectors (n : ℕ) :
    projectiveOperators (fun i : Fin n => Vector.basis i) = projectors n :=
  rfl

theorem projectiveOperators_computationalBasis_isComplete (n : ℕ) :
    IsProjectiveComplete (fun i : Fin n => Vector.basis i) := by
  simpa [projectiveOperators_eq_projectors] using projectors_isComplete n

namespace Projective

variable {n outcomes : ℕ}

instance : CoeFun (Projective n outcomes) (fun _ => Fin outcomes → Vector n) where
  coe M := M.vector

@[simp]
theorem coe_apply (M : Projective n outcomes) (m : Fin outcomes) :
    (M : Fin outcomes → Vector n) m = M.vector m :=
  rfl

noncomputable def operator (M : Projective n outcomes) : Fin outcomes → Square n :=
  projectiveOperators M.vector

/-- View a complete projective measurement as a generalized measurement. -/
noncomputable def toGeneralized (M : Projective n outcomes) : Generalized n outcomes where
  operator := M.operator
  isComplete := M.isComplete

/-- Probability of projective measurement outcome `m`. -/
noncomputable def prob (M : Projective n outcomes) (s : Vector n) (m : Fin outcomes) : ℝ :=
  projProb M.vector s m

/-- Normalized post-measurement state after projective measurement outcome `m`. -/
noncomputable def postMeasure
    (M : Projective n outcomes) (s : Vector n) (m : Fin outcomes) : Vector n :=
  projPostMeasure M.vector s m

@[simp]
theorem operator_apply (M : Projective n outcomes) (m : Fin outcomes) :
    M.operator m = Matrix.proj (M.vector m) :=
  rfl

@[simp]
theorem toGeneralized_operator (M : Projective n outcomes) :
    M.toGeneralized.operator = M.operator :=
  rfl

@[simp]
theorem prob_eq_projProb (M : Projective n outcomes)
    (s : Vector n) (m : Fin outcomes) :
    M.prob s m = projProb M.vector s m :=
  rfl

@[simp]
theorem postMeasure_eq_projPostMeasure (M : Projective n outcomes)
    (s : Vector n) (m : Fin outcomes) :
    M.postMeasure s m = projPostMeasure M.vector s m :=
  rfl

@[simp]
theorem prob_eq_generalized_prob (M : Projective n outcomes)
    (s : Vector n) (m : Fin outcomes) :
    M.prob s m = M.toGeneralized.prob s m :=
  (generalizedProb_projectiveOperators_of_isNormalized s (M.isNormalized m)).symm

@[simp]
theorem postMeasure_eq_generalizedPostMeasure (M : Projective n outcomes)
    (s : Vector n) (m : Fin outcomes) :
    M.postMeasure s m = M.toGeneralized.postMeasure s m :=
  (generalizedPostMeasure_projectiveOperators_of_isNormalized s (M.isNormalized m)).symm

theorem postMeasure_isNormalized (M : Projective n outcomes)
    (s : Vector n) (m : Fin outcomes) (h : M.prob s m ≠ 0) :
    Vector.IsNormalized (M.postMeasure s m) := by
  have hg : generalizedProb (projectiveOperators M.vector) s m ≠ 0 := by
    rw [generalizedProb_projectiveOperators_of_isNormalized s (M.isNormalized m)]
    simpa [Projective.prob] using h
  have hnorm := generalizedPostMeasure_isNormalized (projectiveOperators M.vector) s m hg
  rw [generalizedPostMeasure_projectiveOperators_of_isNormalized s (M.isNormalized m)] at hnorm
  simpa [Projective.postMeasure] using hnorm

theorem operator_mul_postMeasure (M : Projective n outcomes) (s : Vector n) (m : Fin outcomes) :
    M.operator m ⬝ M.postMeasure s m = M.postMeasure s m := by
  unfold Projective.operator Projective.postMeasure projPostMeasure projectiveOperators
  rw [Matrix.mul_smul]
  congr 1
  rw [show Matrix.proj (M.vector m) ⬝ (Matrix.proj (M.vector m) ⬝ s) =
      (Matrix.proj (M.vector m) ⬝ Matrix.proj (M.vector m)) ⬝ s by
    simp [Matrix.mul, _root_.Matrix.mul_assoc]]
  rw [Matrix.proj_mul_proj_of_isNormalized (M.isNormalized m)]

theorem sum_prob_of_isNormalized (M : Projective n outcomes)
    {s : Vector n} (hs : Vector.IsNormalized s) :
    (∑ m : Fin outcomes, M.prob s m) = 1 := by
  calc
    (∑ m : Fin outcomes, M.prob s m) = ∑ m : Fin outcomes, M.toGeneralized.prob s m := by
      apply Finset.sum_congr rfl
      intro m _
      exact M.prob_eq_generalized_prob s m
    _ = 1 := M.toGeneralized.sum_prob_of_isNormalized hs

noncomputable def computationalBasis (n : ℕ) : Projective n n where
  vector := fun i => Vector.basis i
  isNormalized := fun i => Vector.basis_isNormalized i
  isComplete := projectiveOperators_computationalBasis_isComplete n

@[simp]
theorem computationalBasis_apply {n : ℕ} (i : Fin n) :
    (computationalBasis n).vector i = Vector.basis i :=
  rfl

@[simp]
theorem computationalBasis_prob {n : ℕ} (s : Vector n) (i : Fin n) :
    (computationalBasis n).prob s i = _root_.QuantumComputing.Measurement.prob s i := by
  unfold Projective.prob projProb
  simp only [computationalBasis_apply]
  rw [← proj_def i]
  exact quadratic_proj s i

@[simp]
theorem computationalBasis_postMeasure {n : ℕ} (s : Vector n) (i : Fin n) :
    (computationalBasis n).postMeasure s i =
      _root_.QuantumComputing.Measurement.postMeasure s i := by
  unfold Projective.postMeasure projPostMeasure _root_.QuantumComputing.Measurement.postMeasure
  simp only [computationalBasis_apply]
  rw [show projProb (computationalBasis n).vector s i =
      _root_.QuantumComputing.Measurement.prob s i by
    simpa [Projective.prob] using computationalBasis_prob s i]
  rw [← proj_def i]

end Projective

theorem projProb_eq_quadratic_projector {n outcomes : ℕ}
    (u : Fin outcomes → Vector n) (s : Vector n) (i : Fin outcomes) :
    projProb u s i = ((s† ⬝ Matrix.proj (u i) ⬝ s) 0 0).re := by
  simp [projProb]

theorem projective_quadratic_eq_inner_square {n outcomes : ℕ}
    (u : Fin outcomes → Vector n) (s : Vector n) (i : Fin outcomes) :
    s† ⬝ Matrix.proj (u i) ⬝ s = ((u i)† ⬝ s)† ⬝ ((u i)† ⬝ s) := by
  simp [Matrix.proj, Matrix.mul, Matrix.adjoint, _root_.Matrix.mul_assoc]

theorem prob_projectiveMeasurementOp {n : ℕ} (u : Fin n → Vector n) (s : Vector n)
    (i : Fin n) :
    prob (projectiveMeasurementOp u ⬝ s) i = projProb u s i := by
  let z : ℂ := ((u i)† ⬝ s) 0 0
  have hop : (projectiveMeasurementOp u ⬝ s) i 0 = z := by
    simp [z, projectiveMeasurementOp, Matrix.mul, _root_.Matrix.mul_apply]
  have hquad : projProb u s i = Complex.normSq z := by
    rw [projProb]
    rw [projective_quadratic_eq_inner_square]
    have hsum := sum_prob (((u i)†) ⬝ s)
    simpa [prob, z, Matrix.adjoint, Matrix.mul, _root_.Matrix.mul_assoc] using hsum.symm
  rw [hquad]
  simp [prob, hop]

theorem postMeasure_projectiveMeasurementOp {n : ℕ} (u : Fin n → Vector n) (s : Vector n)
    (i : Fin n) :
    projPostMeasure u s i =
      projectiveRestoreOp u i ⬝ postMeasure (projectiveMeasurementOp u ⬝ s) i := by
  unfold projPostMeasure projectiveRestoreOp postMeasure
  rw [prob_projectiveMeasurementOp]
  ext j k
  fin_cases k
  simp [projectiveMeasurementOp, proj, Matrix.proj, Matrix.mul, Matrix.adjoint, Vector.basis,
    _root_.Matrix.mul_apply, Finset.mul_sum, mul_left_comm, mul_comm]

end Measurement

end QuantumComputing