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

# ## 多元变量（Multinomial Variables）

# ### 1. 伯努利分布的推广
# 二元变量可以用来描述两种可能值的随机变量，当遇到K个互斥状态中的某一种离散变量时，我们可以使用多元变量的概率分布来表示。

# 用K维向量x，向量中的一个元素$x_k$等于1，剩余元素等于0。这样的向量满足$\sum_{k=1}^Kx_k=1$。如果我们用参数$\mu_k$表示$x_k=1$的概率，那么x的分布就是$$p(\textbf{x}|\textbf{$\mu$})=\prod_\limits{k=1}^K\mu_k^{x_k}$$
# 
# 其中$\textbf{$\mu$}=(\mu_1,...,\mu_K)^T$，参数$\mu_k$满足$\mu_k\ge0,\sum_k\mu_k=1$。该概率分布可以看成是伯努利分布对于多个输出的推广。

# 现考虑一个有N个独立观测值$\textbf{x}_1,...,\textbf{x}_N$的数据集$\mathcal{D}$。对应的似然函数的形式为$$p(\mathcal{D}|\textbf{$\mu$})=\prod_\limits{n=1}^N\prod_\limits{k=1}^K \mu_k^{x_{nk}}=\prod_\limits{k=1}^K \mu_k^{\sum_nx_{nk}}=\prod_\limits{k=1}^K \mu_k^{m_k}$$
# 

# ### 2. 多项式分布
# 考虑$m_1,...,m_K$在参数$\mu$和观测总数N条件下的联合分布，这个分布的形式为$$Mult(m_1,m_2,...,m_K|\textbf{$\mu$},N)=\binom{N}{m_1m_2...m_K}\prod_\limits{k=1}^K\mu_k^{m_k}$$
# 
# 这被称为多项式分布（$multinomial$ $distribution$），$m_k$满足限制$\sum_\limits{k=1}^Km_k=N$

# ### 3. 狄利克雷分布
# 现在我们介绍多项式分布的参数$\{\mu_k\}$的一组先验分布。通过观察多项式分布的形式，我们可以看到先验分布为$$p(\textbf{$\mu$}|\textbf{$\alpha$})\propto\prod_\limits{k=1}^K\mu_k^{\alpha_k-1}$$
# 其中$0\le\mu_k\le1$且$\sum_k\mu_k=1$。$\alpha_1,...,\alpha_K$是分布的参数，$\textbf{$\alpha$}$表示$(\alpha_1,...,\alpha_K)^T$。由于加和的限制，$\{\mu_k\}$空间上的分布被限制在K-1维单纯形(simplex)中。
# 
# 下图表示了$\mu_1,\mu_2,\mu_3$上的狄利克雷分布被限制在一个单纯形中，限制条件是$0\le\mu_k\le1$且$\sum_k\mu_k=1$。

# ![](http://research.microsoft.com/en-us/um/people/cmbishop/prml/prmlfigs-jpg/Figure2.4.jpg)

# 狄利克雷分布的形式是$$Dir(\textbf{$\mu$}|\textbf{$\alpha$})=\frac{\Gamma(\alpha_0)}{\Gamma(\alpha_1)...\Gamma(\alpha_K)}\prod_\limits{k=1}^K\mu_k^{\alpha_k-1}$$
# 其中$\alpha_0=\sum_\limits{k=1}^K\alpha_k$。
# 用似然函数乘以先验，得到参数$\{\mu_k\}$的后验分布，形式为$$p(\textbf{$\mu$}|\mathcal{D},\textbf{$\alpha$})\propto p(\mathcal{D}|\textbf{$\mu$})p(\textbf{$\mu$}|\textbf{$\alpha$})\propto\prod_\limits{k=1}^K\mu_k^{\alpha_k+m_k-1}$$
# 参数的后验分布的形式也是狄利克雷分布，确定归一化系数，得到
# $$p(\textbf{$\mu$}|\mathcal{D},\textbf{$\alpha$})= Dir(\textbf{$\mu$}|\textbf{$\alpha$+m})=\frac{\Gamma(\alpha_0+N)}{\Gamma(\alpha_1+m_1)...\Gamma(\alpha_K+m_k)}\prod_\limits{k=1}^K\mu_k^{\alpha_k+m_k-1}$$

# ### 4. 狄利克雷分布的图像

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D


# In[2]:


from scipy.stats import dirichlet
import matplotlib.tri as tri
from matplotlib import cm


# In[3]:


corners = np.array([[0, 0], [1, 0], [0.5, 0.75**0.5]])
triangle = tri.Triangulation(corners[:, 0], corners[:, 1])

refiner = tri.UniformTriRefiner(triangle)
trimesh = refiner.refine_triangulation(subdiv=4)

plt.figure(figsize=(8, 4))
for (i, mesh) in enumerate((triangle, trimesh)):
    plt.subplot(1, 2, i+ 1)
    plt.triplot(mesh)
    plt.axis('off')
    plt.axis('equal')
    


# In[4]:


# Mid-points of triangle sides opposite of each corner
midpoints = [(corners[(i + 1) % 3] + corners[(i + 2) % 3]) / 2.0 \
             for i in range(3)]
def xy2bc(xy, tol=1.e-3):
    '''Converts 2D Cartesian coordinates to barycentric.'''
    s = [(corners[i] - midpoints[i]).dot(xy - midpoints[i]) / 0.75 \
         for i in range(3)]
    return np.clip(s, tol, 1.0 - tol)


# In[5]:


def dirichlet_pdf(x, al):
    from operator import mul
    from math import gamma
    alpha = np.array(al)
    coef = gamma(np.sum(alpha)) / reduce(mul, [gamma(a) for a in alpha])
        
    return coef * reduce(mul, [xx ** (aa - 1) for (xx, aa)in zip(x, alpha)])


# In[6]:


def draw_dirichlet(alpha, nlevels=200, subdiv=8, **kwargs):
    refiner = tri.UniformTriRefiner(triangle)
    trimesh = refiner.refine_triangulation(subdiv=subdiv)
    pvals = [dirichlet_pdf(xy2bc(xy) , alpha) for xy in zip(trimesh.x, trimesh.y)]

    fig = plt.figure(figsize=(10,8))
    ax = fig.gca(projection='3d')
    ax.plot_trisurf(trimesh.x, trimesh.y, pvals, cmap=cm.jet, linewidth=0.01)
    plt.axis('equal')
    plt.show()


# **下面是狄利克雷分布的图像，其中两个水平轴是单纯形平面上的坐标轴，垂直轴对应于概率密度的值。**

# $\{\alpha_k\}=10$的情形：

# In[7]:


draw_dirichlet([10,10,10])


# $\{\alpha_k\}=1$的情形：

# In[8]:


draw_dirichlet([1,1,1])


# $\{\alpha_k\}=0.1$的情形：

# In[9]:


draw_dirichlet([0.1,0.1,0.1])