#!/usr/bin/env python
# coding: utf-8

# # 2.2 多元变量

# 现在考虑有 $K$ 个状态的问题。
# 
# 我们用一个 $K$ 维的向量 $(x_1, \dots, x_K)$ 来表示这些状态，第 $k$ 个状态用 $x_k = 1, x_j=0, \forall j \neq k$ 表示。
# 
# 例如 $\mathbf x=(0,0,1,0,0,0)^\text{T}$ 表示 $K = 6$ 的第 $3$ 个状态。
# 
# 这些向量都满足 $\sum_{k=1}^K x_k = 1$。
# 
# 假设 $x_k=1$ 的概率为 $\mu_k$，则 $x$ 的分布为：
# 
# $$
# p(\mathbf x~|~\mathbf\mu)=\prod_{k=1}^K \mu_k^{x_k}
# $$
# 
# 其中 $\mathbf\mu=(\mu_1, \dots, \mu_K)^\text{T}$，且满足 $\mu_k \geq 0, \sum_j \mu_k = 1$。
# 
# 这可以看成是伯努利分布的一般化表示。为了验证它是概率分布，我们有
# 
# $$
# \sum_\mathbf{x} p(\mathbf x~|~\mathbf\mu) = \sum_{k=1}^K \mu_k = 1
# $$
# 
# 均值为
# 
# $$
# \mathbb E[\mathbf{x|\mu}] = \sum_\mathbf{x} \mathbf{x} p(\mathbf x~|~\mathbf\mu) =(\mu_1, \dots, \mu_K)^\text{T} = \mathbf\mu
# $$

# ### 最大似然估计

# 考虑一组 i.i.d. 的观测数据 $\mathbf D=(\mathbf x_1, \dots, \mathbf x_N)$，其似然函数为
# 
# $$
# p(\mathcal D|\mathbf\mu) = \prod_{n=1}^N \prod_{k=1}^K \mu_k^{x_{nk}} = \prod_{k=1}^K \mu_k^{\left(\sum_n x_{nk}\right)} = \prod_{k=1}^K \mu_k^{m_k}
# $$
# 
# 其中 
# 
# $$
# m_k =\sum_n x_{nk} 
# $$
# 
# 表示出现 $x_k=1$ 的次数，这些也是这个分布的充分统计量。
# 
# 在 $\sum_{k=1}^K \mu_k = 1$ 的约束条件下，为了最大化对数似然，我们考虑它的拉格朗日函数：
# 
# $$
# \sum_{k=1}^K m_k\ln \mu_k + \lambda (\sum_{k=1}^K \mu_k - 1)
# $$
# 
# 对 $\mu_k$ 的偏导设为 $0$，有：
# 
# $$
# \mu_k = - \frac{m_k}{\lambda}
# $$
# 
# 带入约束条件得到：
# 
# $$
# \mu_k^{\text{ML}} = \frac{m_k}{N}
# $$
# 
# 即 $x_k=1$ 在观测数据中所占的比例。
# 
# 如果只考虑 $m_1,\dots,m_K$，那么我们可以定义多项分布（`multinomial distribution`）：
# 
# $$
# \text{Mult}(m_1, m_2,\dots,m_k|\mathbf\mu, N) = \begin{pmatrix} N \\ m_1m_2\dots m_k\end{pmatrix} \prod_{k=1}^K \mu_k^{m_k}
# $$
# 
# 其中：
# 
# $$
# \begin{pmatrix} N \\ m_1m_2\dots m_k\end{pmatrix} \equiv \frac{N!}{m_1!m_2!\dotsm_K!}
# $$
# 
# 以及
# 
# $$
# N = \sum_{k=1}^K m_k
# $$

# ## 2.2.1 狄利克雷分布

# 与二元分布类似，我们要给多元分布引入一个先验分布。考虑多元分布的形式，为了满足共轭性，先验分布应该满足这样的形式：
# 
# $$
# p(\mathbf{\mu|\alpha}) \propto \prod_{k=1}^K \mu_k^{\alpha_k-1}
# $$
# 
# 其中 $0 \leq \mu_k \leq 1, \sum_k \mu_k = 1$，$\mathbf\alpha=\left(\alpha_1,\dots,\alpha_K\right)^{\text T}$ 是先验分布的参数。
# 
# 事实上，$\mu_k$ 的分布是一个 $K-1$ 的单纯形。
# 
# 归一化这个分布，我们可以得到
# 
# $$
# \text{Dir}(\mathbf{\mu|\alpha}) = \frac{\Gamma(\alpha_0)}{\Gamma(\alpha_1)\dots\Gamma(\alpha_K)} \prod_{k=1}^K\mu_k^{\alpha_k-1}
# $$
# 
# 其中 
# 
# $$
# \alpha_0 = \sum_{k=1}^K \alpha_k
# $$
# 
# 这个分布叫做狄利克雷分布（`Dirichlet distribution`）。
# 
# 使用狄利克雷分布作为先验，后验分布为：
# 
# $$
# p(\mathbf \mu|\mathcal D,\mathbf \alpha) 
# \propto p(\mathcal D|\mathbf \alpha) p(\mathbf{\mu|\alpha})
# \propto \prod_{k=1}^K \mu_k^{\alpha_k + m_k - 1}
# $$
# 
# 由共轭性，我们可以得到后验分布也是个狄利克雷分布：
# 
# $$
# p(\mathbf \mu|\mathcal D,\mathbf \alpha)
# = \text{Dir}(\mathbf{\mu|\alpha+m})
# = \frac{\Gamma(\alpha_0+ N)}{\Gamma(\alpha_1+m_1)\dots\Gamma(\alpha_K+m_K)} \prod_{k=1}^K \mu_k^{\alpha_k+m_k-1}
# $$
# 
# 其中 $\mathbf m=\left(m_1,\dots,m_K\right)^{\text T}$。
# 
# 与二项分布类似，我们也可以将 $a_k$ 看成 $x_k=1$ 的一个有效观测次数。
# 
# 二项分布可以看出是多项分布 $K = 2$ 的特殊情况。