Library probsa.probability.law_of_total_prob

From discprob.prob Require Export prob countable stochastic_order indep.

From probsa.util Require Export indicator.
From probsa.probability Require Export conditional.


Law of Total Probability

In the following section, we show that if we have a countable set of pairwise disjoint sets, then the sum of probabilities of those sets converges to a finite value. The proof of this result is highly technical; therefore, it is recommended to skip over any "Definitions" and "Remarks".
Section SumOfDisjointEventsExists.

  Variable (Ω : countType).
  Variable (μ : measure Ω).

  Variables (I : countType) (B : I pred Ω).
  Hypothesis H_Bi_disjoint : ω, i j : I, B i ω B j ω i = j.

  Definition is_in_cover : pred Ω :=
    λ (ω : Ω), exC (λ i, B i ω).

  Definition get_cover_index (ω : Ω) : option I :=
    match LPO_countb (fun iB i ω) with
    | inleft iSome (sval i)
    | inright _None
    end.

  Definition pickle_cover_of ω : nat :=
    match get_cover_index ω with
    | NoneO
    | Some i(pickle i).+1
    end.

  Remark get_cover_index_valid :
     v a, B v a get_cover_index a = Some v.
  Proof.
    intros × Bva; unfold get_cover_index.
    destruct (LPO_countb) as [[i Bvi]|NO]; last first.
    { by exfalso; eapply NO; exact Bva. }
    { by rewrite /sval //= (H_Bi_disjoint a i v). }
  Qed.

  Definition apt :=
    λ mn, match mn with
          | (S m, S n)
              match (@pickle_inv I m), (@pickle_inv Ω n) with
              | Some i, Some aif B i a then μ a else 0
              | _, _ ⇒ 0
              end
          | _ ⇒ 0
          end.

  Definition σpt : nat nat × nat :=
    λ n, match @pickle_inv [countType of Ω] n with
         | Some a(pickle_cover_of a , S(pickle a))
         | None(O, O)
         end.

  Remark σpt_inj :
     n n', apt (σpt n) 0 σpt n = σpt n' n = n'.
  Proof.
    intros n n'; rewrite /σpt/apt.
    case_eq (@pickle_inv [countType of Ω] n); last by nra.
    case_eq (@pickle_inv [countType of Ω] n'); last first.
    { by intros Heq_none a Heq_some ⇒ //=. }
    intros a' Heq' a Heq Hneq0 ⇒ //=.
    inversion 1 as [[Hp1 Hp2]].
    assert (a = a').
    { apply (f_equal (@pickle_inv Ω)) in Hp2. rewrite ?pickleK_inv in Hp2.
      inversion Hp2; done. }
    subst.
    apply (f_equal (oapp (@pickle _) n)) in Heq.
    apply (f_equal (oapp (@pickle _) n')) in Heq'.
    by rewrite ?pickle_invK //= in Heq Heq'; congruence.
  Qed.

  Remark σpt_cov :
     n, apt n 0 m, σpt m = n.
  Proof.
    intros [n1 n2]; destruct n1, n2 ⇒ //=.
    case_eq (@pickle_inv I n1) ⇒ //= ⇒ v Heq1.
    case_eq (@pickle_inv Ω n2) ⇒ //= ⇒ a Heq2 Hneq0.
     (pickle a).
    rewrite /σpt pickleK_inv; repeat f_equal; last first.
    - apply (f_equal (oapp (@pickle _) n2)) in Heq2.
      by rewrite pickle_invK //= in Heq2.
    - apply (f_equal (oapp (@pickle _) n1)) in Heq1.
      rewrite pickle_invK //= in Heq1; rewrite Heq1.
      move: Hneq0; case: ifP; last by nra.
      moveHeq ?; unfold pickle_cover_of.
      erewrite get_cover_index_valid; last by exact Heq.
      reflexivity.
  Qed.

  Remark apt_row_rewrite :
     j, Series (λ k : nat, apt (S j, S k)) = countable_sum (λ i : I, <μ>{[ B i ]}) j.
  Proof.
    intros j; rewrite /apt/countable_sum.
    destruct (@pickle_inv _ j) as [v|] ⇒ //=; last apply Series_0 ⇒ //=.
  Qed.

  Remark apt_double_summable_by_column :
    double_summable apt.
  Proof.
    have Hext:
       v i x,
        B i v
        (λ n : nat, match @pickle_inv I n with Some v0if B v0 v then 1 else 0 | None ⇒ 0 end) x
        = (λ n : nat, match eq_nat_dec n (@pickle I i) with left _ ⇒ 1 | _ ⇒ 0 end) x.
    { intros v i x Biv.
      destruct Nat.eq_dec as [e|n].
      - by rewrite //= e pickleK_inv Biv //= eq_refl.
      - specialize (@pickle_invK I x) ⇒ //=.
        destruct pickle_inv ⇒ //.
        case: ifP; last by done.
        moveBsv Heq.
        exfalso; apply: n; rewrite -Heq; clear Heq.
        by rewrite (H_Bi_disjoint v s i).
    }
    apply ex_series_columns_ds.
    { intros j; destruct j.
      { 0; apply is_series_0n.
        by destruct n; rewrite Rabs_R0. }
      { rewrite ex_series_incr_1 /apt.
        destruct (pickle_inv j) as [v|]; last first.
        { 0; apply is_series_0n.
          by destruct (pickle_inv _); rewrite Rabs_R0. }
        apply (ex_series_ext (λ k, μ v × (match @pickle_inv I k with Some v0if B v0 v then 1 else 0 | None ⇒ 0 end))).
        { intros k; destruct (pickle_inv) as [s|] ⇒ //=; last by rewrite Rabs_R0; nra.
          case: ifP_; try by rewrite Rabs_R0; nra.
          by rewrite Rmult_1_r; symmetry; apply Rabs_right, pmf_pos. }
        apply: ex_series_scal.
        destruct (is_in_cover v) eqn:EQ.
        { move: EQ ⇒ /exCP [i Biv].
          eapply ex_series_ext.
          - by symmetry; apply: Hext Biv.
          - 1; apply is_series_ext with (a := λ n : nat, if Nat.eq_dec n (pickle i) then 1 else 0);
              last by apply is_series_bump.
            by intros ?; rewrite /is_left //=; destruct (Nat.eq_dec).
        }
        { 0; apply is_series_0.
          intros ?; destruct (pickle_inv); last by done.
          destruct (B s v) eqn:Bsv; last by done.
          exfalso; move : EQ =>/eqP; rewrite eqbF_neg ⇒ /negP NEX; apply:NEX.
          by apply/exCP; s.
        }
      }
    }
    { rewrite ex_series_incr_1.
      have EX_pt: ex_series (countable_sum (fun aif is_in_cover a then μ a else 0)) by apply pr_ex_series.
      destruct EX_pt; eapply ex_series_ext; last by eexists; eauto.
      intros k; rewrite /countable_sum/apt.
      rewrite /oapp//=; destruct (@pickle_inv _ k) as [v|] ⇒ //=; last first.
      { symmetry; apply Series_0; intros [|]; first by rewrite Rabs_R0.
        by destruct (pickle_inv _) ⇒ //=; rewrite Rabs_R0.
      }
      rewrite Series_incr_1_aux; last by apply Rabs_R0.
      rewrite -(Series_ext (λ k, μ v × (match @pickle_inv I k with Some v0if B v0 v then 1 else 0 | None ⇒ 0 end)));
        first last.
      { intros j; destruct (pickle_inv) as [s|] ⇒ //=; last by rewrite Rabs_R0; nra.
        case: ifP_; try by rewrite Rabs_R0; nra.
        by rewrite Rmult_1_r; symmetry; apply Rabs_right, pmf_pos. }
      { destruct (is_in_cover v) eqn:EQ.
        { move: EQ ⇒ /exCP [i Biv].
          rewrite Series_scal_l -{1}[μ v]Rmult_1_r.
          apply Rmult_eq_compat_l; symmetry.
          erewrite Series_ext; [ | intros; apply Hext; eauto].
          rewrite -{2}(Series_bump (pickle i) 1); apply Series_ext.
          by intros ?; rewrite /is_left //=; destruct (Nat.eq_dec). }
        { symmetry; apply Series_0n.
          apply Rmult_eq_0_compat_l.
          destruct (pickle_inv); last by done.
          destruct (B s v) eqn:Bsv; last by done.
          exfalso; move : EQ =>/eqP; rewrite eqbF_neg ⇒ /negP NEX; apply:NEX.
          by apply/exCP; s.
        }
      }
    }
  Qed.

Next we prove that if we have a partition of our measure space Ω into sets B i, then the probabilities of each set B i can be summed up and that sum is equal to the sum of the measure of all ω Ω, where ω is only included in the sum if ω is covered by some set in the partition.
  Remark pr_eq_over_partition_is_series :
    is_series
      (countable_sum (λ (i : I), <μ>{[ B i ]}))
      (Series (countable_sum (λ (ω : Ω), if is_in_cover ω then μ ω else 0))).
  Proof.
    eapply (is_series_ext (λ j, Series (λ k, apt (S j, k)%N))).
    { by moven; rewrite Series_incr_1_aux ⇒ //; apply: apt_row_rewrite. }
    rewrite -(Series_ext (apt \o σpt)); last first.
    { moven; rewrite /apt/σpt/countable_sum//=.
      destruct (@pickle_inv _ n) as [a|] ⇒ //=.
      rewrite ?pickleK_inv; destruct (is_in_cover a) eqn:INCOV.
      { move : INCOV ⇒ /exCP [x Bxa]; move: (Bxa) ⇒ COV.
        apply get_cover_index_valid in Bxa.
        by rewrite /pickle_cover_of Bxa ?pickleK_inv COV. }
      { destruct (pickle_cover_of _) eqn:PICK1; first by done.
        destruct (pickle_inv) eqn:PICK2; last by done.
        destruct (B s a) eqn:Bsa; last by done.
        exfalso; move : INCOV =>/eqP; rewrite eqbF_neg ⇒ /negP NEX; apply:NEX.
        by apply/exCP; s. }
    }
    apply (is_series_incr_1 (λ j, Series (λ k, apt (j, k)))).
    rewrite {3}/apt [a in is_series _ (plus _ a)]Series_0 // plus_zero_r.
    by apply series_double_covering;
      auto using σpt_inj, σpt_cov, apt_double_summable_by_column.
  Qed.

Finally, we show that if we have a countable partition of a probability space (Ω, μ) into (disjoint) sets B i, then the series of probabilities of those sets converges. In other words, if we sum the probability measures of each set in the partition, the resulting series converges.
  Lemma ex_series_pr_eq_over_disjoint :
    ex_series (countable_sum (λ i : I, <μ>{[ B i ]})).
  Proof. by eexists; apply pr_eq_over_partition_is_series. Qed.

End SumOfDisjointEventsExists.

Section SumOfPartitionEventsExists.

  Lemma ex_series_pr_eq_over_partition :
     {Ω} (μ : measure Ω), (P : @Ω_partition Ω μ),
      ex_series (countable_sum (λ i : I P, <μ>{[ p P i]})).
  Proof. by intros; apply ex_series_pr_eq_over_disjoint; destruct P. Qed.

End SumOfPartitionEventsExists.

In the following section we prove the law of total probability. The law states that, given an event P and a countable set of pairwise disjoint events B i, the probability of P is equal to the sum of the probabilities of the intersection P B i for all i.
The proof of this result is highly technical; therefore, it is recommended to skip over any "Definitions" and "Remarks".
Consider a probability space (Ω, μ), ...
  Context {Ω} {μ : measure Ω}.

Consider any predicate P : Ω bool. Notice that any predicate trivially defines a subset of Ω as Ω | P ω }. Also recall that subsets of Ω are also events.
  Variable (P : pred Ω).

Consider a countable type I (this can be either finite or countably infinite set) and a family of predicates B : I pred Ω.
  Variables (I : countType) (B : I pred Ω).

Next, let us assume that B covers all positive-probability elements of Ω. That is, for any element ω : Ω with positive probability, there exists an element i : I such that ω satisfies/belongs to B i.
Second, let us assume that B is a set of disjoint sets. That is, for any element ω : Ω and for any distinct i, j : I, if both B i ω and B j ω are true, then i must equal j. This enforces the requirement that the sets defined by the predicates in B must be disjoint.
  Hypothesis H_Bi_covers : ω, μ ω > 0 i : I, B i ω.
  Hypothesis H_Bi_disjoint : ω, i j : I, B i ω B j ω i = j.

  Definition atotal : nat × nat R :=
    fun mn
      match mn with
      | (m, n)
          match (@pickle_inv I m), (@pickle_inv Ω n) with
          | Some i, Some a ⇒ (if (B i a) && P a then μ a else 0)
          | _, _ ⇒ 0
          end
      end.

  Remark ex_total_column_abs i :
    ex_series
      (λ k : nat,
          Rabs
            match @pickle_inv Ω k with
            | Some aif (B i a) && P a then μ a else 0
            | None ⇒ 0
            end).
  Proof.
    apply: (ex_series_le _ (countable_sum (λ a : Ω, if P a then μ a else 0)));
      last by apply pr_ex_series.
    intros n. rewrite /norm //= /abs //=.
    rewrite /countable_sum. rewrite /pr_eq/pr //=.
    destruct (pickle_inv _) as [s|] ⇒ //=.
    × rewrite Rabs_Rabsolu.
      rewrite Rabs_right; auto.
      ** case: ifP.
         *** move /andP ⇒ [? ->]. reflexivity.
         *** intros; case: ifP; auto using Rge_le, pmf_pos; reflexivity.
      ** case: ifP ⇒ ?.
         *** apply pmf_pos.
         *** nra.
    × rewrite ?Rabs_R0. reflexivity.
  Qed.

  Remark atotal_double_summable_by_column:
    double_summable atotal.
  Proof.
    apply ex_series_rows_ds.
    - intros j.
      × rewrite /atotal. destruct (@pickle_inv _ j) as [v|]; last first.
        { 0. apply is_series_0n. rewrite Rabs_R0; done. }
        apply ex_total_column_abs.
    - rewrite /atotal.
      apply: (ex_series_le _ (countable_sum (λ i : I, pr μ (B i)))).
      { intros n. rewrite /norm //= /abs //=.
        rewrite /countable_sum. rewrite /pr_eq/pr //=.
        destruct (pickle_inv _) as [s|] ⇒ //=.
        × rewrite Rabs_right.
          ** apply: Series_le. intros n'; split.
             *** apply Rabs_pos.
             *** rewrite /countable_sum//=.
                 destruct (pickle_inv _) ⇒ //=; last (rewrite Rabs_R0; reflexivity).
                 case: ifP.
                 **** move /andP ⇒ [-> ?]. rewrite Rabs_right; first reflexivity.
                      apply pmf_pos.
                 **** rewrite Rabs_R0. intros; case: ifP; auto using Rge_le, pmf_pos; reflexivity.
             *** apply pr_ex_series.
          ** apply Rle_ge, Series_pos; auto.
        × rewrite Series_0 // Rabs_R0; reflexivity.
      }
      apply ex_series_pr_eq_over_disjoint.
      eauto.
  Qed.

  Remark atotal_row_rewrite j :
    Series (λ k : nat, atotal (j, k))
    = countable_sum (λ i : I, pr μ (λ a, (B i a) && P a)) j.
  Proof.
    rewrite /atotal/countable_sum.
    destruct (@pickle_inv _ j) as [v|] ⇒ //=; last apply Series_0 ⇒ //=.
  Qed.

  Remark atotal_column_rewrite k :
    Series (λ j : nat, atotal (j, k))
    = countable_sum (λ a, if P a then μ a else 0) k.
  Proof.
    rewrite /atotal/countable_sum.
    destruct (@pickle_inv Ω k) as [ω|]; first last.
    { rewrite //=. apply Series_0n //=. destruct (pickle_inv) ⇒ //=. }
    { rewrite //=. destruct (P ω); first last.
      { apply Series_0n //=. destruct (pickle_inv) ⇒ //=. by rewrite andbF. }
      { have CASE : μ ω = 0 μ ω > 0
          by destruct μ; move: (pmf_pos ω) ⇒ [POS | Z]; [right | left].
        destruct CASE as [Z | POS].
        { rewrite !Z; apply Series_0n.
          by destruct (pickle_inv); first destruct (B s ω). }
        have EX: i, B i ω by apply H_Bi_covers.
        destruct EX as [i IN].
        rewrite -{2}(Series_bump (@pickle I i) (μ ω)).
        apply Series_extn.
        destruct (Nat.eq_dec) as [Heqn|Hneq].
        ** by subst; rewrite pickleK_inv //= IN //=.
        ** specialize (@pickle_invK I n).
           destruct (pickle_inv) as [i'|] ⇒ //=. intros Heq'.
           case: ifP ⇒ //=.
           rewrite andbT. moveHeq. subst.
           exfalso; apply: Hneq; f_equal.
           by eapply H_Bi_disjoint; eauto.
      }
    }
  Qed.

With this, we prove that the probability of an event defined by P is equal to the (possibly infinite) sum of events defined by B i P, where i I.
  Lemma law_of_total_probability :
    <μ>{[ P ]} = ∑[∞]_{i <- I} <μ>{[ B i P ]}.
  Proof.
    rewrite /pr.
    symmetry; etransitivity.
    { by apply Series_extn; rewrite -atotal_row_rewrite. }
    { rewrite Series_double_swap; auto using atotal_double_summable_by_column.
      by apply Series_extk; apply atotal_column_rewrite. }
  Qed.

End LawOfTotalProbability.

Section LawOfTotalProbabilityProd.

  Context {Ω} (μ : measure Ω).

  Variable (A : pred Ω).
  Variable (S : @Ω_partition Ω μ).

  Lemma law_of_total_probability_prod :
    <μ>{[ A ]} = ∑[∞]_{ i <- I S } <μ>{[ A p S i ]}.
  Proof.
    symmetry; erewrite law_of_total_probability.
    - apply SeriesC_exti.
      apply pr_eq_pred ⇒ ω; rewrite !unfold_in.
      by apply andb_comm.
    - destruct S as [? ? COVER ? ].
      by intros; eapply COVER.
    - destruct S as [? ? ? DISJ].
      by intros; eapply DISJ; eauto.
  Qed.

End LawOfTotalProbabilityProd.


Corollary bigop_inf_cdf_le :
   {Ω} {μ : measure Ω} (X : nrvar μ) (n0 : nat),
    ∑[∞]_{n<-nat} (I[n n0] × <μ>{[ λ ω, X ω == n ]})%R = <μ>{[ X ⟨<=⟩ n0 ]}.
Proof.
  clear; intros; erewrite law_of_total_probability with (B := fun n ωX ω == n); first last.
  { by move ⇒ ω i j /eqP EQ1 /eqP EQ2; subst. }
  { by intros ω _; (X ω). }
  { apply SeriesC_extn.
    rewrite -pr_bool_move.
    apply pr_eq_pred ⇒ ω; rewrite !unfold_in.
    unfold "⟨<=⟩", rvar_nat_pred_leqop.
    by destruct (_ == _) eqn:EQ;
    [move: EQ ⇒ /eqP ->; rewrite andb_comm
    | rewrite andb_false_r andb_false_l].
  }
Qed.

Corollary bigop_inf_option_cdf_lt :
   {Ω} {μ : measure Ω} (X : Ω option nat) (n0 : option nat),
    ∑[∞]_{n<-option nat} (I[n0 ⟨<⟩ n] × <μ>{[ λ ω, X ω == n ]})%R = <μ>{[ λ ω, n0 ⟨<⟩ X ω ]}.
Proof.
  intros; erewrite law_of_total_probability with (B := fun n ωX ω == n); first last.
  { by move ⇒ ω i j /eqP EQ1 /eqP EQ2; subst. }
  { by intros ω _; (X ω). }
  { apply SeriesC_extn.
    rewrite -pr_bool_move.
    apply pr_eq_pred ⇒ ω.
    rewrite !unfold_in.
    by destruct (_ == _) eqn:EQ;
    [move: EQ ⇒ /eqP ->; rewrite andb_comm
    | rewrite andb_false_r andb_false_l].
  }
Qed.