# Theory Pi_pmf

(*
File:    Pi_pmf.thy
Authors: Manuel Eberl, Max W. Haslbeck
*)
section ‹Indexed products of PMFs›
theory Pi_pmf
imports "HOL-Probability.Probability"
begin

text ‹Conflicting notation from \<^theory>‹HOL-Analysis.Infinite_Sum››
no_notation Infinite_Sum.abs_summable_on (infixr "abs'_summable'_on" 46)

subsection ‹Definition›

text ‹
In analogy to @{const PiM}, we define an indexed product of PMFs. In the literature, this
is typically called taking a vector of independent random variables. Note that the components
do not have to be identically distributed.

The operation takes an explicit index set \<^term>‹A :: 'a set› and a function \<^term>‹f :: 'a ⇒ 'b pmf›
that maps each element from \<^term>‹A› to a PMF and defines the product measure
$\bigotimes_{i\in A} f(i)$ , which is represented as a \<^typ>‹('a ⇒ 'b) pmf›.

Note that unlike @{const PiM}, this only works for ∗‹finite› index sets. It could
be extended to countable sets and beyond, but the construction becomes somewhat more involved.
›
definition Pi_pmf :: "'a set ⇒ 'b ⇒ ('a ⇒ 'b pmf) ⇒ ('a ⇒ 'b) pmf" where
"Pi_pmf A dflt p =
embed_pmf (λf. if (∀x. x ∉ A ⟶ f x = dflt) then ∏x∈A. pmf (p x) (f x) else 0)"

text ‹
A technical subtlety that needs to be addressed is this: Intuitively, the functions in the
support of a product distribution have domain ‹A›. However, since HOL is a total logic, these
functions must still return ∗‹some› value for inputs outside ‹A›. The product measure
@{const PiM} simply lets these functions return @{const undefined} in these cases. We chose a
different solution here, which is to supply a default value \<^term>‹dflt :: 'b› that is returned
in these cases.

As one possible application, one could model the result of ‹n› different independent coin
tosses as @{term "Pi_pmf {0..<n} False (λ_. bernoulli_pmf (1 / 2))"}. This returns a function
of type \<^typ>‹nat ⇒ bool› that maps every natural number below ‹n› to the result of the
corresponding coin toss, and every other natural number to \<^term>‹False›.
›

lemma pmf_Pi:
assumes A: "finite A"
shows   "pmf (Pi_pmf A dflt p) f =
(if (∀x. x ∉ A ⟶ f x = dflt) then ∏x∈A. pmf (p x) (f x) else 0)"
unfolding Pi_pmf_def
proof (rule pmf_embed_pmf, goal_cases)
case 2
define S where "S = {f. ∀x. x ∉ A ⟶ f x = dflt}"
define B where "B = (λx. set_pmf (p x))"

have neutral_left: "(∏xa∈A. pmf (p xa) (f xa)) = 0"
if "f ∈ PiE A B - (λf. restrict f A)  S" for f
proof -
have "restrict (λx. if x ∈ A then f x else dflt) A ∈ (λf. restrict f A)  S"
by (intro imageI) (auto simp: S_def)
also have "restrict (λx. if x ∈ A then f x else dflt) A = f"
using that by (auto simp: PiE_def Pi_def extensional_def fun_eq_iff)
finally show ?thesis using that by blast
qed
have neutral_right: "(∏xa∈A. pmf (p xa) (f xa)) = 0"
if "f ∈ (λf. restrict f A)  S - PiE A B" for f
proof -
from that obtain f' where f': "f = restrict f' A" "f' ∈ S" by auto
moreover from this and that have "restrict f' A ∉ PiE A B" by simp
then obtain x where "x ∈ A" "pmf (p x) (f' x) = 0" by (auto simp: B_def set_pmf_eq)
with f' and A show ?thesis by auto
qed

have "(λf. ∏x∈A. pmf (p x) (f x)) abs_summable_on PiE A B"
by (intro abs_summable_on_prod_PiE A) (auto simp: B_def)
also have "?this ⟷ (λf. ∏x∈A. pmf (p x) (f x)) abs_summable_on (λf. restrict f A)  S"
by (intro abs_summable_on_cong_neutral neutral_left neutral_right) auto
also have "… ⟷ (λf. ∏x∈A. pmf (p x) (restrict f A x)) abs_summable_on S"
by (rule abs_summable_on_reindex_iff [symmetric]) (force simp: inj_on_def fun_eq_iff S_def)
also have "… ⟷ (λf. if ∀x. x ∉ A ⟶ f x = dflt then ∏x∈A. pmf (p x) (f x) else 0)
abs_summable_on UNIV"
by (intro abs_summable_on_cong_neutral) (auto simp: S_def)
finally have summable: … .

have "1 = (∏x∈A. 1::real)" by simp
also have "(∏x∈A. 1) = (∏x∈A. ∑⇩ay∈B x. pmf (p x) y)"
unfolding B_def by (subst infsetsum_pmf_eq_1) auto
also have "(∏x∈A. ∑⇩ay∈B x. pmf (p x) y) = (∑⇩af∈Pi⇩E A B. ∏x∈A. pmf (p x) (f x))"
by (intro infsetsum_prod_PiE [symmetric] A) (auto simp: B_def)
also have "… = (∑⇩af∈(λf. restrict f A)  S. ∏x∈A. pmf (p x) (f x))" using A
by (intro infsetsum_cong_neutral neutral_left neutral_right refl)
also have "… = (∑⇩af∈S. ∏x∈A. pmf (p x) (restrict f A x))"
by (rule infsetsum_reindex) (force simp: inj_on_def fun_eq_iff S_def)
also have "… = (∑⇩af∈S. ∏x∈A. pmf (p x) (f x))"
by (intro infsetsum_cong) (auto simp: S_def)
also have "… = (∑⇩af. if ∀x. x ∉ A ⟶ f x = dflt then ∏x∈A. pmf (p x) (f x) else 0)"
by (intro infsetsum_cong_neutral) (auto simp: S_def)
also have "ennreal … = (∫⇧+f. ennreal (if ∀x. x ∉ A ⟶ f x = dflt
then ∏x∈A. pmf (p x) (f x) else 0) ∂count_space UNIV)"
by (intro nn_integral_conv_infsetsum [symmetric] summable) (auto simp: prod_nonneg)
finally show ?case by simp
qed (auto simp: prod_nonneg)

lemma pmf_Pi':
assumes "finite A" "⋀x. x ∉ A ⟹ f x = dflt"
shows   "pmf (Pi_pmf A dflt p) f = (∏x∈A. pmf (p x) (f x))"
using assms by (subst pmf_Pi) auto

lemma pmf_Pi_outside:
assumes "finite A" "∃x. x ∉ A ∧ f x ≠ dflt"
shows   "pmf (Pi_pmf A dflt p) f = 0"
using assms by (subst pmf_Pi) auto

lemma pmf_Pi_empty [simp]: "Pi_pmf {} dflt p = return_pmf (λ_. dflt)"
by (intro pmf_eqI, subst pmf_Pi) (auto simp: indicator_def)

lemma set_Pi_pmf_subset: "finite A ⟹ set_pmf (Pi_pmf A dflt p) ⊆ {f. ∀x. x ∉ A ⟶ f x = dflt}"
by (auto simp: set_pmf_eq pmf_Pi)

lemma Pi_pmf_cong [cong]:
assumes "A = A'" "dflt = dflt'" "⋀x. x ∈ A ⟹ f x = f' x"
shows   "Pi_pmf A dflt f = Pi_pmf A' dflt' f'"
proof -
have "(λg. ∏x∈A. pmf (f x) (g x)) = (λg. ∏x∈A. pmf (f' x) (g x))"
by (intro ext prod.cong) (auto simp: assms)
with assms show ?thesis by (simp add: Pi_pmf_def cong: if_cong)
qed

subsection ‹Dependent product sets with a default›

text ‹
The following describes a dependent product of sets where the functions are required to return
the default value \<^term>‹dflt› outside their domain, in analogy to @{const PiE}, which uses
@{const undefined}.
›
definition PiE_dflt
where "PiE_dflt A dflt B = {f. ∀x. (x ∈ A ⟶ f x ∈ B x) ∧ (x ∉ A ⟶ f x = dflt)}"

lemma restrict_PiE_dflt: "(λh. restrict h A)  PiE_dflt A dflt B = PiE A B"
proof (intro equalityI subsetI)
fix h assume "h ∈ (λh. restrict h A)  PiE_dflt A dflt B"
thus "h ∈ PiE A B"
by (auto simp: PiE_dflt_def)
next
fix h assume h: "h ∈ PiE A B"
hence "restrict (λx. if x ∈ A then h x else dflt) A ∈ (λh. restrict h A)  PiE_dflt A dflt B"
by (intro imageI) (auto simp: PiE_def extensional_def PiE_dflt_def)
also have "restrict (λx. if x ∈ A then h x else dflt) A = h"
using h by (auto simp: fun_eq_iff)
finally show "h ∈ (λh. restrict h A)  PiE_dflt A dflt B" .
qed

lemma dflt_image_PiE: "(λh x. if x ∈ A then h x else dflt)  PiE A B = PiE_dflt A dflt B"
(is "?f  ?X = ?Y")
proof (intro equalityI subsetI)
fix h assume "h ∈ ?f  ?X"
thus "h ∈ ?Y"
by (auto simp: PiE_dflt_def PiE_def)
next
fix h assume h: "h ∈ ?Y"
hence "?f (restrict h A) ∈ ?f  ?X"
by (intro imageI) (auto simp: PiE_def extensional_def PiE_dflt_def)
also have "?f (restrict h A) = h"
using h by (auto simp: fun_eq_iff PiE_dflt_def)
finally show "h ∈ ?f  ?X" .
qed

lemma finite_PiE_dflt [intro]:
assumes "finite A" "⋀x. x ∈ A ⟹ finite (B x)"
shows   "finite (PiE_dflt A d B)"
proof -
have "PiE_dflt A d B = (λf x. if x ∈ A then f x else d)  PiE A B"
by (rule dflt_image_PiE [symmetric])
also have "finite …"
by (intro finite_imageI finite_PiE assms)
finally show ?thesis .
qed

lemma card_PiE_dflt:
assumes "finite A" "⋀x. x ∈ A ⟹ finite (B x)"
shows   "card (PiE_dflt A d B) = (∏x∈A. card (B x))"
proof -
from assms have "(∏x∈A. card (B x)) = card (PiE A B)"
by (intro card_PiE [symmetric]) auto
also have "PiE A B = (λf. restrict f A)  PiE_dflt A d B"
by (rule restrict_PiE_dflt [symmetric])
also have "card … = card (PiE_dflt A d B)"
by (intro card_image) (force simp: inj_on_def restrict_def fun_eq_iff PiE_dflt_def)
finally show ?thesis ..
qed

lemma PiE_dflt_empty_iff [simp]: "PiE_dflt A dflt B = {} ⟷ (∃x∈A. B x = {})"
by (simp add: dflt_image_PiE [symmetric] PiE_eq_empty_iff)

text ‹
The probability of an independent combination of events is precisely the product
of the probabilities of each individual event.
›
lemma measure_Pi_pmf_PiE_dflt:
assumes [simp]: "finite A"
shows   "measure_pmf.prob (Pi_pmf A dflt p) (PiE_dflt A dflt B) =
(∏x∈A. measure_pmf.prob (p x) (B x))"
proof -
define B' where "B' = (λx. B x ∩ set_pmf (p x))"
have "measure_pmf.prob (Pi_pmf A dflt p) (PiE_dflt A dflt B) =
(∑⇩ah∈PiE_dflt A dflt B. pmf (Pi_pmf A dflt p) h)"
by (rule measure_pmf_conv_infsetsum)
also have "… = (∑⇩ah∈PiE_dflt A dflt B. ∏x∈A. pmf (p x) (h x))"
by (intro infsetsum_cong, subst pmf_Pi') (auto simp: PiE_dflt_def)
also have "… = (∑⇩ah∈(λh. restrict h A)  PiE_dflt A dflt B. ∏x∈A. pmf (p x) (h x))"
by (subst infsetsum_reindex) (force simp: inj_on_def PiE_dflt_def fun_eq_iff)+
also have "(λh. restrict h A)  PiE_dflt A dflt B = PiE A B"
by (rule restrict_PiE_dflt)
also have "(∑⇩ah∈PiE A B. ∏x∈A. pmf (p x) (h x)) = (∑⇩ah∈PiE A B'. ∏x∈A. pmf (p x) (h x))"
by (intro infsetsum_cong_neutral) (auto simp: B'_def set_pmf_eq)
also have "(∑⇩ah∈PiE A B'. ∏x∈A. pmf (p x) (h x)) = (∏x∈A. infsetsum (pmf (p x)) (B' x))"
by (intro infsetsum_prod_PiE) (auto simp: B'_def)
also have "… = (∏x∈A. infsetsum (pmf (p x)) (B x))"
by (intro prod.cong infsetsum_cong_neutral) (auto simp: B'_def set_pmf_eq)
also have "… = (∏x∈A. measure_pmf.prob (p x) (B x))"
by (subst measure_pmf_conv_infsetsum) (rule refl)
finally show ?thesis .
qed

lemma set_Pi_pmf_subset':
assumes "finite A"
shows   "set_pmf (Pi_pmf A dflt p) ⊆ PiE_dflt A dflt (set_pmf ∘ p)"
using assms by (auto simp: set_pmf_eq pmf_Pi PiE_dflt_def)

lemma Pi_pmf_return_pmf [simp]:
assumes "finite A"
shows   "Pi_pmf A dflt (λx. return_pmf (f x)) = return_pmf (λx. if x ∈ A then f x else dflt)"
proof -
have "set_pmf (Pi_pmf A dflt (λx. return_pmf (f x))) ⊆
PiE_dflt A dflt (set_pmf ∘ (λx. return_pmf (f x)))"
by (intro set_Pi_pmf_subset' assms)
also have "… ⊆ {λx. if x ∈ A then f x else dflt}"
by (auto simp: PiE_dflt_def)
finally show ?thesis
by (simp add: set_pmf_subset_singleton)
qed

lemma Pi_pmf_return_pmf' [simp]:
assumes "finite A"
shows   "Pi_pmf A dflt (λ_. return_pmf dflt) = return_pmf (λ_. dflt)"
using assms by simp

lemma measure_Pi_pmf_Pi:
fixes t::nat
assumes [simp]: "finite A"
shows   "measure_pmf.prob (Pi_pmf A dflt p) (Pi A B) =
(∏x∈A. measure_pmf.prob (p x) (B x))" (is "?lhs = ?rhs")
proof -
have "?lhs = measure_pmf.prob (Pi_pmf A dflt p) (PiE_dflt A dflt B)"
by (intro measure_prob_cong_0)
(auto simp: PiE_dflt_def PiE_def intro!: pmf_Pi_outside)+
also have "… = ?rhs"
using assms by (simp add: measure_Pi_pmf_PiE_dflt)
finally show ?thesis
by simp
qed

subsection ‹Common PMF operations on products›

text ‹
@{const Pi_pmf} distributes over the bind' operation in the Giry monad:
›
lemma Pi_pmf_bind:
assumes "finite A"
shows   "Pi_pmf A d (λx. bind_pmf (p x) (q x)) =
do {f ← Pi_pmf A d' p; Pi_pmf A d (λx. q x (f x))}" (is "?lhs = ?rhs")
proof (rule pmf_eqI, goal_cases)
case (1 f)
show ?case
proof (cases "∃x∈-A. f x ≠ d")
case False
define B where "B = (λx. set_pmf (p x))"
have [simp]: "countable (B x)" for x by (auto simp: B_def)

{
fix x :: 'a
have "(λa. pmf (p x) a * 1) abs_summable_on B x"
by (simp add: pmf_abs_summable)
moreover have "norm (pmf (p x) a * 1) ≥ norm (pmf (p x) a * pmf (q x a) (f x))" for a
unfolding norm_mult by (intro mult_left_mono) (auto simp: pmf_le_1)
ultimately have "(λa. pmf (p x) a * pmf (q x a) (f x)) abs_summable_on B x"
by (rule abs_summable_on_comparison_test)
} note summable = this

have "pmf ?rhs f = (∑⇩ag. pmf (Pi_pmf A d' p) g * (∏x∈A. pmf (q x (g x)) (f x)))"
by (subst pmf_bind, subst pmf_Pi')
(insert assms False, simp_all add: pmf_expectation_eq_infsetsum)
also have "… = (∑⇩ag∈PiE_dflt A d' B.
pmf (Pi_pmf A d' p) g * (∏x∈A. pmf (q x (g x)) (f x)))" unfolding B_def
using assms by (intro infsetsum_cong_neutral) (auto simp: pmf_Pi PiE_dflt_def set_pmf_eq)
also have "… = (∑⇩ag∈PiE_dflt A d' B.
(∏x∈A. pmf (p x) (g x) * pmf (q x (g x)) (f x)))"
using assms by (intro infsetsum_cong) (auto simp: pmf_Pi PiE_dflt_def prod.distrib)
also have "… = (∑⇩ag∈(λg. restrict g A)  PiE_dflt A d' B.
(∏x∈A. pmf (p x) (g x) * pmf (q x (g x)) (f x)))"
by (subst infsetsum_reindex) (force simp: PiE_dflt_def inj_on_def fun_eq_iff)+
also have "(λg. restrict g A)  PiE_dflt A d' B = PiE A B"
by (rule restrict_PiE_dflt)
also have "(∑⇩ag∈…. (∏x∈A. pmf (p x) (g x) * pmf (q x (g x)) (f x))) =
(∏x∈A. ∑⇩aa∈B x. pmf (p x) a * pmf (q x a) (f x))"
using assms summable by (subst infsetsum_prod_PiE) simp_all
also have "… = (∏x∈A. ∑⇩aa. pmf (p x) a * pmf (q x a) (f x))"
by (intro prod.cong infsetsum_cong_neutral) (auto simp: B_def set_pmf_eq)
also have "… = pmf ?lhs f"
using False assms by (subst pmf_Pi') (simp_all add: pmf_bind pmf_expectation_eq_infsetsum)
finally show ?thesis ..
next
case True
have "pmf ?rhs f =
measure_pmf.expectation (Pi_pmf A d' p) (λx. pmf (Pi_pmf A d (λxa. q xa (x xa))) f)"
using assms by (simp add: pmf_bind)
also have "… = measure_pmf.expectation (Pi_pmf A d' p) (λx. 0)"
using assms True by (intro Bochner_Integration.integral_cong pmf_Pi_outside) auto
also have "… = pmf ?lhs f"
using assms True by (subst pmf_Pi_outside) auto
finally show ?thesis ..
qed
qed

text ‹
Analogously any componentwise mapping can be pulled outside the product:
›
lemma Pi_pmf_map:
assumes [simp]: "finite A" and "f dflt = dflt'"
shows   "Pi_pmf A dflt' (λx. map_pmf f (g x)) = map_pmf (λh. f ∘ h) (Pi_pmf A dflt g)"
proof -
have "Pi_pmf A dflt' (λx. map_pmf f (g x)) =
Pi_pmf A dflt' (λx. g x ⤜ (λx. return_pmf (f x)))"
using assms by (simp add: map_pmf_def Pi_pmf_bind)
also have "… = Pi_pmf A dflt g ⤜ (λh. return_pmf (λx. if x ∈ A then f (h x) else dflt'))"
by (subst Pi_pmf_bind[where d' = dflt]) auto
also have "… = map_pmf (λh. f ∘ h) (Pi_pmf A dflt g)"
unfolding map_pmf_def using set_Pi_pmf_subset'[of A dflt g]
by (intro bind_pmf_cong refl arg_cong[of _ _ return_pmf])
(auto dest: simp: fun_eq_iff PiE_dflt_def assms(2))
finally show ?thesis .
qed

text ‹
We can exchange the default value in a product of PMFs like this:
›
lemma Pi_pmf_default_swap:
assumes "finite A"
shows   "map_pmf (λf x. if x ∈ A then f x else dflt') (Pi_pmf A dflt p) =
Pi_pmf A dflt' p" (is "?lhs = ?rhs")
proof (rule pmf_eqI, goal_cases)
case (1 f)
let ?B = "(λf x. if x ∈ A then f x else dflt') - {f} ∩ PiE_dflt A dflt (λ_. UNIV)"
show ?case
proof (cases "∃x∈-A. f x ≠ dflt'")
case False
let ?f' = "λx. if x ∈ A then f x else dflt"
from False have "pmf ?lhs f = measure_pmf.prob (Pi_pmf A dflt p) ?B"
using assms unfolding pmf_map
by (intro measure_prob_cong_0) (auto simp: PiE_dflt_def pmf_Pi_outside)
also from False have "?B = {?f'}"
by (auto simp: fun_eq_iff PiE_dflt_def)
also have "measure_pmf.prob (Pi_pmf A dflt p) {?f'} = pmf (Pi_pmf A dflt p) ?f'"
by (simp add: measure_pmf_single)
also have "… = pmf ?rhs f"
using False assms by (subst (1 2) pmf_Pi) auto
finally show ?thesis .
next
case True
have "pmf ?lhs f = measure_pmf.prob (Pi_pmf A dflt p) ?B"
using assms unfolding pmf_map
by (intro measure_prob_cong_0) (auto simp: PiE_dflt_def pmf_Pi_outside)
also from True have "?B = {}" by auto
also have "measure_pmf.prob (Pi_pmf A dflt p) … = 0"
by simp
also have "0 = pmf ?rhs f"
using True assms by (intro pmf_Pi_outside [symmetric]) auto
finally show ?thesis .
qed
qed

text ‹
The following rule allows reindexing the product:
›
lemma Pi_pmf_bij_betw:
assumes "finite A" "bij_betw h A B" "⋀x. x ∉ A ⟹ h x ∉ B"
shows "Pi_pmf A dflt (λ_. f) = map_pmf (λg. g ∘ h) (Pi_pmf B dflt (λ_. f))"
(is "?lhs = ?rhs")
proof -
have B: "finite B"
using assms bij_betw_finite by auto
have "pmf ?lhs g = pmf ?rhs g" for g
proof (cases "∀a. a ∉ A ⟶ g a = dflt")
case True
define h' where "h' = the_inv_into A h"
have h': "h' (h x) = x" if "x ∈ A" for x
unfolding h'_def using that assms by (auto simp add: bij_betw_def the_inv_into_f_f)
have h: "h (h' x) = x" if "x ∈ B" for x
unfolding h'_def using that assms f_the_inv_into_f_bij_betw by fastforce
have "pmf ?rhs g = measure_pmf.prob (Pi_pmf B dflt (λ_. f)) ((λg. g ∘ h) - {g})"
unfolding pmf_map by simp
also have "… = measure_pmf.prob (Pi_pmf B dflt (λ_. f))
(((λg. g ∘ h) - {g}) ∩ PiE_dflt B dflt (λ_. UNIV))"
using B by (intro measure_prob_cong_0) (auto simp: PiE_dflt_def pmf_Pi_outside)
also have "… = pmf (Pi_pmf B dflt (λ_. f)) (λx. if x ∈ B then g (h' x) else dflt)"
proof -
have "(if h x ∈ B then g (h' (h x)) else dflt) = g x" for x
using h' assms True by (cases "x ∈ A") (auto simp add: bij_betwE)
then have "(λg. g ∘ h) - {g} ∩ PiE_dflt B dflt (λ_. UNIV) =
{(λx. if x ∈ B then g (h' x) else dflt)}"
using assms h' h True unfolding PiE_dflt_def by auto
then show ?thesis
by (simp add: measure_pmf_single)
qed
also have "… = pmf (Pi_pmf A dflt (λ_. f)) g"
using B assms True  h'_def
by (auto simp add: pmf_Pi intro!: prod.reindex_bij_betw bij_betw_the_inv_into)
finally show ?thesis
by simp
next
case False
have "pmf ?rhs g = infsetsum (pmf (Pi_pmf B dflt (λ_. f))) ((λg. g ∘ h) - {g})"
using assms by (auto simp add: measure_pmf_conv_infsetsum pmf_map)
also have "… = infsetsum (λ_. 0) ((λg x. g (h x)) - {g})"
using B False assms by (intro infsetsum_cong pmf_Pi_outside) fastforce+
also have "… = 0"
by simp
finally show ?thesis
using assms False by (auto simp add: pmf_Pi pmf_map)
qed
then show ?thesis
by (rule pmf_eqI)
qed

text ‹
A product of uniform random choices is again a uniform distribution.
›
lemma Pi_pmf_of_set:
assumes "finite A" "⋀x. x ∈ A ⟹ finite (B x)" "⋀x. x ∈ A ⟹ B x ≠ {}"
shows   "Pi_pmf A d (λx. pmf_of_set (B x)) = pmf_of_set (PiE_dflt A d B)" (is "?lhs = ?rhs")
proof (rule pmf_eqI, goal_cases)
case (1 f)
show ?case
proof (cases "∃x. x ∉ A ∧ f x ≠ d")
case True
hence "pmf ?lhs f = 0"
using assms by (intro pmf_Pi_outside) (auto simp: PiE_dflt_def)
also from True have "f ∉ PiE_dflt A d B"
by (auto simp: PiE_dflt_def)
hence "0 = pmf ?rhs f"
using assms by (subst pmf_of_set) auto
finally show ?thesis .
next
case False
hence "pmf ?lhs f = (∏x∈A. pmf (pmf_of_set (B x)) (f x))"
using assms by (subst pmf_Pi') auto
also have "… = (∏x∈A. indicator (B x) (f x) / real (card (B x)))"
by (intro prod.cong refl, subst pmf_of_set) (use assms False in auto)
also have "… = (∏x∈A. indicator (B x) (f x)) / real (∏x∈A. card (B x))"
by (subst prod_dividef) simp_all
also have "(∏x∈A. indicator (B x) (f x) :: real) = indicator (PiE_dflt A d B) f"
using assms False by (auto simp: indicator_def PiE_dflt_def)
also have "(∏x∈A. card (B x)) = card (PiE_dflt A d B)"
using assms by (intro card_PiE_dflt [symmetric]) auto
also have "indicator (PiE_dflt A d B) f / … = pmf ?rhs f"
using assms by (intro pmf_of_set [symmetric]) auto
finally show ?thesis .
qed
qed

subsection ‹Merging and splitting PMF products›

text ‹
The following lemma shows that we can add a single PMF to a product:
›
lemma Pi_pmf_insert:
assumes "finite A" "x ∉ A"
shows   "Pi_pmf (insert x A) dflt p = map_pmf (λ(y,f). f(x:=y)) (pair_pmf (p x) (Pi_pmf A dflt p))"
proof (intro pmf_eqI)
fix f
let ?M = "pair_pmf (p x) (Pi_pmf A dflt p)"
have "pmf (map_pmf (λ(y, f). f(x := y)) ?M) f =
measure_pmf.prob ?M ((λ(y, f). f(x := y)) - {f})"
by (subst pmf_map) auto
also have "((λ(y, f). f(x := y)) - {f}) = (⋃y'. {(f x, f(x := y'))})"
by (auto simp: fun_upd_def fun_eq_iff)
also have "measure_pmf.prob ?M … = measure_pmf.prob ?M {(f x, f(x := dflt))}"
using assms by (intro measure_prob_cong_0) (auto simp: pmf_pair pmf_Pi split: if_splits)
also have "… = pmf (p x) (f x) * pmf (Pi_pmf A dflt p) (f(x := dflt))"
by (simp add: measure_pmf_single pmf_pair pmf_Pi)
also have "… = pmf (Pi_pmf (insert x A) dflt p) f"
proof (cases "∀y. y ∉ insert x A ⟶ f y = dflt")
case True
with assms have "pmf (p x) (f x) * pmf (Pi_pmf A dflt p) (f(x := dflt)) =
pmf (p x) (f x) * (∏xa∈A. pmf (p xa) ((f(x := dflt)) xa))"
by (subst pmf_Pi') auto
also have "(∏xa∈A. pmf (p xa) ((f(x := dflt)) xa)) = (∏xa∈A. pmf (p xa) (f xa))"
using assms by (intro prod.cong) auto
also have "pmf (p x) (f x) * … = pmf (Pi_pmf (insert x A) dflt p) f"
using assms True by (subst pmf_Pi') auto
finally show ?thesis .
qed (insert assms, auto simp: pmf_Pi)
finally show "… = pmf (map_pmf (λ(y, f). f(x := y)) ?M) f" ..
qed

lemma Pi_pmf_insert':
assumes "finite A"  "x ∉ A"
shows   "Pi_pmf (insert x A) dflt p =
do {y ← p x; f ← Pi_pmf A dflt p; return_pmf (f(x := y))}"
using assms
by (subst Pi_pmf_insert)
(auto simp add: map_pmf_def pair_pmf_def case_prod_beta' bind_return_pmf bind_assoc_pmf)

lemma Pi_pmf_singleton:
"Pi_pmf {x} dflt p = map_pmf (λa b. if b = x then a else dflt) (p x)"
proof -
have "Pi_pmf {x} dflt p = map_pmf (fun_upd (λ_. dflt) x) (p x)"
by (subst Pi_pmf_insert) (simp_all add: pair_return_pmf2 pmf.map_comp o_def)
also have "fun_upd (λ_. dflt) x = (λz y. if y = x then z else dflt)"
by (simp add: fun_upd_def fun_eq_iff)
finally show ?thesis .
qed

text ‹
Projecting a product of PMFs onto a component yields the expected result:
›
lemma Pi_pmf_component:
assumes "finite A"
shows   "map_pmf (λf. f x) (Pi_pmf A dflt p) = (if x ∈ A then p x else return_pmf dflt)"
proof (cases "x ∈ A")
case True
define A' where "A' = A - {x}"
from assms and True have A': "A = insert x A'"
by (auto simp: A'_def)
from assms have "map_pmf (λf. f x) (Pi_pmf A dflt p) = p x" unfolding A'
by (subst Pi_pmf_insert)
(auto simp: A'_def pmf.map_comp o_def case_prod_unfold map_fst_pair_pmf)
with True show ?thesis by simp
next
case False
have "map_pmf (λf. f x) (Pi_pmf A dflt p) = map_pmf (λ_. dflt) (Pi_pmf A dflt p)"
using assms False set_Pi_pmf_subset[of A dflt p]
by (intro pmf.map_cong refl) (auto simp: set_pmf_eq pmf_Pi_outside)
with False show ?thesis by simp
qed

text ‹
We can take merge two PMF products on disjoint sets like this:
›
lemma Pi_pmf_union:
assumes "finite A" "finite B" "A ∩ B = {}"
shows   "Pi_pmf (A ∪ B) dflt p =
map_pmf (λ(f,g) x. if x ∈ A then f x else g x)
(pair_pmf (Pi_pmf A dflt p) (Pi_pmf B dflt p))" (is "_ = map_pmf (?h A) (?q A)")
using assms(1,3)
proof (induction rule: finite_induct)
case (insert x A)
have "map_pmf (?h (insert x A)) (?q (insert x A)) =
do {v ← p x; (f, g) ← pair_pmf (Pi_pmf A dflt p) (Pi_pmf B dflt p);
return_pmf (λy. if y ∈ insert x A then (f(x := v)) y else g y)}"
by (subst Pi_pmf_insert)
(insert insert.hyps insert.prems,
simp_all add: pair_pmf_def map_bind_pmf bind_map_pmf bind_assoc_pmf bind_return_pmf)
also have "… = do {v ← p x; (f, g) ← ?q A; return_pmf ((?h A (f,g))(x := v))}"
by (intro bind_pmf_cong refl) (auto simp: fun_eq_iff)
also have "… = do {v ← p x; f ← map_pmf (?h A) (?q A); return_pmf (f(x := v))}"
by (simp add: bind_map_pmf map_bind_pmf case_prod_unfold cong: if_cong)
also have "… = do {v ← p x; f ← Pi_pmf (A ∪ B) dflt p; return_pmf (f(x := v))}"
using insert.hyps and insert.prems by (intro bind_pmf_cong insert.IH [symmetric] refl) auto
also have "… = Pi_pmf (insert x (A ∪ B)) dflt p"
by (subst Pi_pmf_insert)
(insert assms insert.hyps insert.prems, auto simp: pair_pmf_def map_bind_pmf)
also have "insert x (A ∪ B) = insert x A ∪ B"
by simp
finally show ?case ..
qed (simp_all add: case_prod_unfold map_snd_pair_pmf)

text ‹
We can also project a product to a subset of the indices by mapping all the other
indices to the default value:
›
lemma Pi_pmf_subset:
assumes "finite A" "A' ⊆ A"
shows   "Pi_pmf A' dflt p = map_pmf (λf x. if x ∈ A' then f x else dflt) (Pi_pmf A dflt p)"
proof -
let ?P = "pair_pmf (Pi_pmf A' dflt p) (Pi_pmf (A - A') dflt p)"
from assms have [simp]: "finite A'"
by (blast dest: finite_subset)
from assms have "A = A' ∪ (A - A')"
by blast
also have "Pi_pmf … dflt p = map_pmf (λ(f,g) x. if x ∈ A' then f x else g x) ?P"
using assms by (intro Pi_pmf_union) auto
also have "map_pmf (λf x. if x ∈ A' then f x else dflt) … = map_pmf fst ?P"
unfolding map_pmf_comp o_def case_prod_unfold
using set_Pi_pmf_subset[of A' dflt p] by (intro map_pmf_cong refl) (auto simp: fun_eq_iff)
also have "… = Pi_pmf A' dflt p"
by (simp add: map_fst_pair_pmf)
finally show ?thesis ..
qed

lemma Pi_pmf_subset':
fixes f :: "'a ⇒ 'b pmf"
assumes "finite A" "B ⊆ A" "⋀x. x ∈ A - B ⟹ f x = return_pmf dflt"
shows "Pi_pmf A dflt f = Pi_pmf B dflt f"
proof -
have "Pi_pmf (B ∪ (A - B)) dflt f =
map_pmf (λ(f, g) x. if x ∈ B then f x else g x)
(pair_pmf (Pi_pmf B dflt f) (Pi_pmf (A - B) dflt f))"
using assms by (intro Pi_pmf_union) (auto dest: finite_subset)
also have "Pi_pmf (A - B) dflt f = Pi_pmf (A - B) dflt (λ_. return_pmf dflt)"
using assms by (intro Pi_pmf_cong) auto
also have "… = return_pmf (λ_. dflt)"
using assms by simp
also have "map_pmf (λ(f, g) x. if x ∈ B then f x else g x)
(pair_pmf (Pi_pmf B dflt f) (return_pmf (λ_. dflt))) =
map_pmf (λf x. if x ∈ B then f x else dflt) (Pi_pmf B dflt f)"
by (simp add: map_pmf_def pair_pmf_def bind_assoc_pmf bind_return_pmf bind_return_pmf')
also have "… = Pi_pmf B dflt f"
using assms by (intro Pi_pmf_default_swap) (auto dest: finite_subset)
also have "B ∪ (A - B) = A"
using assms by auto
finally show ?thesis .
qed

lemma Pi_pmf_if_set:
assumes "finite A"
shows "Pi_pmf A dflt (λx. if b x then f x else return_pmf dflt) =
Pi_pmf {x∈A. b x} dflt f"
proof -
have "Pi_pmf A dflt (λx. if b x then f x else return_pmf dflt) =
Pi_pmf {x∈A. b x} dflt (λx. if b x then f x else return_pmf dflt)"
using assms by (intro Pi_pmf_subset') auto
also have "… = Pi_pmf {x∈A. b x} dflt f"
by (intro Pi_pmf_cong) auto
finally show ?thesis .
qed

lemma Pi_pmf_if_set':
assumes "finite A"
shows "Pi_pmf A dflt (λx. if b x then return_pmf dflt else f x) =
Pi_pmf {x∈A. ¬b x} dflt f"
proof -
have "Pi_pmf A dflt (λx. if b x then return_pmf dflt else  f x) =
Pi_pmf {x∈A. ¬b x} dflt (λx. if b x then return_pmf dflt else  f x)"
using assms by (intro Pi_pmf_subset') auto
also have "… = Pi_pmf {x∈A. ¬b x} dflt f"
by (intro Pi_pmf_cong) auto
finally show ?thesis .
qed

text ‹
Lastly, we can delete a single component from a product:
›
lemma Pi_pmf_remove:
assumes "finite A"
shows   "Pi_pmf (A - {x}) dflt p = map_pmf (λf. f(x := dflt)) (Pi_pmf A dflt p)"
proof -
have "Pi_pmf (A - {x}) dflt p =
map_pmf (λf xa. if xa ∈ A - {x} then f xa else dflt) (Pi_pmf A dflt p)"
using assms by (intro Pi_pmf_subset) auto
also have "… = map_pmf (λf. f(x := dflt)) (Pi_pmf A dflt p)"
using set_Pi_pmf_subset[of A dflt p] assms
by (intro map_pmf_cong refl) (auto simp: fun_eq_iff)
finally show ?thesis .
qed

subsection ‹Applications›

text ‹
Choosing a subset of a set uniformly at random is equivalent to tossing a fair coin
independently for each element and collecting all the elements that came up heads.
›
lemma pmf_of_set_Pow_conv_bernoulli:
assumes "finite (A :: 'a set)"
shows "map_pmf (λb. {x∈A. b x}) (Pi_pmf A P (λ_. bernoulli_pmf (1/2))) = pmf_of_set (Pow A)"
proof -
have "Pi_pmf A P (λ_. bernoulli_pmf (1/2)) = pmf_of_set (PiE_dflt A P (λx. UNIV))"
using assms by (simp add: bernoulli_pmf_half_conv_pmf_of_set Pi_pmf_of_set)
also have "map_pmf (λb. {x∈A. b x}) … = pmf_of_set (Pow A)"
proof -
have "bij_betw (λb. {x ∈ A. b x}) (PiE_dflt A P (λ_. UNIV)) (Pow A)"
by (rule bij_betwI[of _ _ _ "λB b. if b ∈ A then b ∈ B else P"]) (auto simp add: PiE_dflt_def)
then show ?thesis
using assms by (intro map_pmf_of_set_bij_betw) auto
qed
finally show ?thesis
by simp
qed

text ‹
A binomial distribution can be seen as the number of successes in ‹n› independent coin tosses.
›
lemma binomial_pmf_altdef':
fixes A :: "'a set"
assumes "finite A" and "card A = n" and p: "p ∈ {0..1}"
shows   "binomial_pmf n p =
map_pmf (λf. card {x∈A. f x}) (Pi_pmf A dflt (λ_. bernoulli_pmf p))" (is "?lhs = ?rhs")
proof -
from assms have "?lhs = binomial_pmf (card A) p"
by simp
also have "… = ?rhs"
using assms(1)
proof (induction rule: finite_induct)
case empty
with p show ?case by (simp add: binomial_pmf_0)
next
case (insert x A)
from insert.hyps have "card (insert x A) = Suc (card A)"
by simp
also have "binomial_pmf … p = do {
b ← bernoulli_pmf p;
f ← Pi_pmf A dflt (λ_. bernoulli_pmf p);
return_pmf ((if b then 1 else 0) + card {y ∈ A. f y})
}"
using p by (simp add: binomial_pmf_Suc insert.IH bind_map_pmf)
also have "… = do {
b ← bernoulli_pmf p;
f ← Pi_pmf A dflt (λ_. bernoulli_pmf p);
return_pmf (card {y ∈ insert x A. (f(x := b)) y})
}"
proof (intro bind_pmf_cong refl, goal_cases)
case (1 b f)
have "(if b then 1 else 0) + card {y∈A. f y} = card ((if b then {x} else {}) ∪ {y∈A. f y})"
using insert.hyps by auto
also have "(if b then {x} else {}) ∪ {y∈A. f y} = {y∈insert x A. (f(x := b)) y}"
using insert.hyps by auto
finally show ?case by simp
qed
also have "… = map_pmf (λf. card {y∈insert x A. f y})
(Pi_pmf (insert x A) dflt (λ_. bernoulli_pmf p))"
using insert.hyps by (subst Pi_pmf_insert) (simp_all add: pair_pmf_def map_bind_pmf)
finally show ?case .
qed
finally show ?thesis .
qed

end
`