Library probsa.probability.pmf

From probsa.util Require Export bigop_inf.
From probsa.probability Require Export pred stochastic_order independence.

Section ProbMassFunction.

  Program Definition pmf_zero : distrib [countType of nat] :=
    {|
      pmf := fun (n : nat) ⇒ if (n == 0%nat) then 1 else 0
    |}.
  Next Obligation.
    by case (a == 0%nat); auto with real.
  Qed.
  Next Obligation.
    by apply is_seriesC_bump.
  Qed.

  Program Definition pmf_sum (x1 x2 : distrib [countType of nat]) : distrib [countType of nat] :=
    {|
      pmf := fun (n : nat) ⇒ _{0 i n} pmf x1 i × pmf x2 (n - i)%nat
    |}.
  Next Obligation.
    apply Rle_ge, Rle_big0i _.
    destruct x1, x2 ⇒ //=.
    by apply Rmult_le_0_compat; auto with real.
  Qed.
  Next Obligation.
    move: (x1) (x2) ⇒ [pmf1 pos1 sum1] [pmf2 pos2 sum2].
    rewrite -[1](Rmult_1_l).
    apply: is_series_ext; last apply is_series_mult; first last ⇒ //=.
    { 1; apply: is_series_ext; last by apply sum2.
      intros n; symmetry.
      apply Rabs_pos_eq.
      rewrite /countable_sum /pickle_inv //= eq_refl.
      by apply Rge_le, pos2.
    }
    { 1; apply: is_series_ext; last by apply sum1.
      intros n; symmetry.
      apply Rabs_pos_eq.
      rewrite /countable_sum /pickle_inv //= eq_refl.
      by apply Rge_le, pos1.
    }
    { by apply sum2. }
    { by apply sum1. }
    { intros n.
      have → : f n, sum_f_R0 f n = \big[Rplus/0]_(0 i < n.+1) f i.
      { clear; induction n.
        - by rewrite big_nat1.
        - by rewrite tech5 big_nat_recr //= IHn.
      }
      rewrite /countable_sum /pickle_inv /pickle //= eq_refl //=.
      by apply eq_bigrt _; rewrite !eq_refl //=.
    }
  Qed.

End ProbMassFunction.


Infix "⊕" := pmf_sum (at level 50, left associativity).

Notation "⨁_{ i <- xs | P } F" :=
  (\big[pmf_sum/pmf_zero]_(i <- xs | P) F%R)
    (at level 41, F at level 41, i, xs at level 50,
      format "'[' ⨁_{ i <- xs | P } '/ ' F ']'") : probability_scope.

Notation "⨁_{ m <= i < n } F" :=
  (\big[pmf_sum/pmf_zero]_(i <- index_iota m n) F%R)
    (at level 41, F at level 41, i, m, n at level 50,
      format "'[' ⨁_{ m <= i < n } '/ ' F ']'") : probability_scope.

Notation "⨁_{ i < n } F" :=
  (\big[pmf_sum/pmf_zero]_(i <- iota 0 n) F%R)
    (at level 41, F at level 41, i, n at level 50,
      format "'[' ⨁_{ i < n } '/ ' F ']'") : probability_scope.


We say a random variable X is stochastically dominated by distribution p (written X p) if the CDF of X is pointwise greater than or equal to the CDF of p. Formally, <μ>{X h} <p>{ω | ω h}, where on the RHS, ω ranges over natural numbers sampled from p, so ω h is simply a comparison of natural numbers. Note that a distribution p over natural numbers can be viewed as a probability measure where p(n) gives the probability of observing n.
Definition le_nrvar_ndistrib {Ω} {μ : measure Ω} (X : nrvar μ) (p : distrib [countType of nat]) :=
   h, <μ>{[ X ⟨<=⟩ h ]} <p>{[ leq^~ h ]}.

Instance dominance_nrvar_ndistrib :
   {Ω} {μ : measure Ω}, DominanceRelation (nrvar μ) (distrib [countType of nat]) :=
  { dominates := le_nrvar_ndistrib }.

Definition le_ndistrib_nrvar {Ω} {μ : measure Ω} (p : distrib [countType of nat]) (X : nrvar μ) :=
   h, <p>{[ leq^~ h ]} <μ>{[ X ⟨<=⟩ h ]}.

Instance dominance_ndistrib_nrvar :
   {Ω} {μ : measure Ω}, DominanceRelation (distrib [countType of nat]) (nrvar μ) :=
  { dominates := le_ndistrib_nrvar }.


Section StochasticOrder.

  Context {Ω} {μ : measure Ω}.

  Variable X1 X2 : nrvar μ.
  Hypothesis H_independent : indep2 X1 X2.

  Variable p1 p2 : distrib [countType of nat].
  Hypothesis H_X1_bounded_by_p1 : X1 p1.
  Hypothesis H_X2_bounded_by_p2 : X2 p2.

If X1 p1, X2 p2, and X1, X2 are independent, then their sum is stochastically dominated by the convolution of p1 and p2.
  Lemma addrv_addmpf_respect_stochastic_order :
    X1 ⟨+⟩ X2 p1 p2.
  Proof.
    set (Ωn := [countType of nat × nat]).
    set (μp := distrib_prod p1 p2).
    set (P1 := mkRvar μp (fun '(n1, n2)n1) : nrvar μp).
    set (P2 := mkRvar μp (fun '(n1, n2)n2) : nrvar μp).
    have PREC: indep2 P1 P2.
    { intros x1 x2; unfold pr_eq, P1, P2 ⇒ //=.
      apply: eq_tr4; first apply: pair_joint_pred; first last.
      { apply Rmult_eq_compat.
        { apply: eq_tr3; last apply: pair_marginal1.
          apply pr_eq_pred ⇒ ω; destruct ω as [ω1 ω2].
          by rewrite !unfold_in //=.
        }
        { apply: eq_tr3; last apply: pair_marginal2.
          apply pr_eq_pred ⇒ ω; destruct ω as [ω1 ω2].
          by rewrite !unfold_in //=.
        }
      }
      apply pr_eq_pred ⇒ ω; destruct ω as [ω1 ω2].
      by rewrite !unfold_in //=.
    }
    apply: func_dom_trans.
    apply: (addrv_respects_stochastic_order _ _ P1 P2) ⇒ //.
    { apply: func_dom_trans; first by apply H_X1_bounded_by_p1.
      intros k; unfold cdf.
      apply Rle_ge, Req_le.
      apply: eq_tr3; last apply: pair_marginal1.
      apply pr_eq_pred ⇒ ω; destruct ω as [ω1 ω2].
      by rewrite !unfold_in //=.
    }
    { apply: func_dom_trans; first by apply H_X2_bounded_by_p2.
      intros k; unfold cdf.
      apply Rle_ge, Req_le.
      apply: eq_tr3; last apply: pair_marginal2.
      apply pr_eq_pred ⇒ ω; destruct ω as [ω1 ω2].
      by rewrite !unfold_in //=.
    }
    { intros k.
      apply Rle_ge, Req_le; symmetry.
      rewrite cdf_to_sum_of_preq.
      apply: eq_tr3; last by symmetry; apply (cdf_to_sum_of_preq (mkRvar (p1 p2) (fun ωω))).
      rewrite !big_nat; apply eq_bigri LEi.
      rewrite stochastic_order.addrv_eq_decomposition //.
      apply: eq_tr3.
      { apply eq_bigrx _.
        apply Rmult_eq_compat.
        { apply: eq_tr3; last apply: pair_marginal1.
          apply pr_eq_pred ⇒ ω; destruct ω as [ω1 ω2].
          by rewrite !unfold_in //=.
        }
        { apply: eq_tr3; last apply: pair_marginal2.
          apply pr_eq_pred ⇒ ω; destruct ω as [ω1 ω2].
          by rewrite !unfold_in //=.
        }
      }
      rewrite /pr !SeriesCf_bump.
      unfold "⊕"; rewrite //= big_nat big_nat.
      apply eq_bigrj LEj.
      by rewrite /pr !SeriesCf_bump.
    }
  Qed.

Complement of the previous lemma: the probability that the sum of two independent random variables exceeds a threshold t is bounded by the probability that a sample from the convolution of their bounding distributions exceeds t.
  Corollary addrv_addmpf_respect_ltn :
     t,
      <μ>{[ λ ω, (t < X1 ω + X2 ω)%nat ]} <p1 p2>{[ ltn t ]}.
  Proof.
    by intros t; apply pr_ineq_compl, addrv_addmpf_respect_stochastic_order.
  Qed.

End StochasticOrder.

Section StochasticOrderSum.

  Context {Ω} {μ : measure Ω}.

  Context {X : eqType}.
  Variable (F : X nrvar μ) (P : pred X).

  Variable (xs : seq X).
  Hypothesis H_independent : independent [seq F x | x <- xs].

  Variable p : X distrib [countType of nat].
  Hypothesis H_bound : x, x \in xs P x F x p x.

  Lemma sumrv_sumpmf_respect_stochastic_order :
    ∑[rv]_{ x <- xs | P x } F x _{x <- xs | P x} p x.
  Proof.
    induction xs as [ | y ys].
    { rewrite !big_nilk.
      apply Rle_ge; apply: Rle_trans; first by apply pr_le_1.
      by rewrite pr_xpredT_ext //; apply Rge_refl.
    }
    { rewrite !big_cons.
      destruct (P y) eqn:Py; last first.
      { apply IHys ⇒ ×.
        - apply: independent_tl.
          erewrite <-map_cons.
          by apply H_independent.
        - apply H_bound ⇒ //=.
          by rewrite in_cons; apply/orP; right.
      }
      { apply addrv_addmpf_respect_stochastic_order.
        { move: (H_independent) ⇒ EEQ.
          have EQ : y :: ys = [::y] ++ ys by done.
          rewrite EQ in EEQ.
          eapply indep2_sum with (P0 := P) in EEQb1 b2.
          apply: eq_tr4; first apply EEQ.
          - apply pr_eq_pred ⇒ ω; rewrite !unfold_in; f_equal; f_equal.
            by rewrite big_cons Py big_nil //= addn0.
          - f_equal; apply pr_eq_pred ⇒ ω; rewrite !unfold_in; f_equal.
            by rewrite big_cons Py big_nil //= addn0.
        }
        { by apply H_bound ⇒ //; rewrite in_cons; apply/orP; left. }
        { apply IHys ⇒ ×.
          - apply: independent_tl.
            erewrite <-map_cons.
            by apply H_independent.
          - apply H_bound ⇒ //=.
            by rewrite in_cons; apply/orP; right.
        }
      }
    }
  Qed.

End StochasticOrderSum.