Dirichlet Distribution

On this weekend, I just tried to draw "Dirichlet Distribution" on three dimention, however it was way harder that I thought :( So I will share the way to draw Dirichlet distribution. If someone know more slick and efficient way to do it, I would be thankful to let me know.

1. What is dirichlet distribution¶

First of all, we should be on the same page on what the dirichlet distribution is. Dirichlet Distribution is a family of continuous multivariate probability distribution. It is multivariate generalization of beta distribution. Especially Dirichlet distribution is known as prior confugate of Categorical distribution and Multinominal distribution.
Dirichlet distribution is parametarized by a vector $\alpha$ which is k-dimentional vecotor such that $\alpha \in \mathbb{R }^{+}$ , let $\pi$ be k-dimentaional vector such that $\sum_k^K\pi_k = 1$ and $\pi_k \in (0,1)$, defined by following, $$Dir(\pi|\alpha) = C_D(\alpha)\prod_k^K\pi_k^{\alpha_k -1}$$

2. Implementation¶

Here I've drawn Dirichlet distribution with three different $\alpha$ of [10,10,10], [0.5,0.5,0.5],[1.0,1.0,1.0].

In [1]:

import numpy as np
from scipy.stats import dirichlet
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
% matplotlib inline

In [2]:

# Make box-shaped mesh-grid
X = np.linspace(0.01,0.98,50)
Y = np.linspace(0.01,0.98,50)
X_grid,Y_grid = np.meshgrid(X,Y)

# Make elements where addition of X and Y  is less than 1 0.
_mask = (X_grid+Y_grid) < 1
X_grid_mask = X_grid * _mask

max_row = np.max(X_grid_mask,axis=1)
max_row = max_row[:,np.newaxis]

X_tile = np.tile(max_row,(X.shape[0]))
inv_mask = 1 - _mask
X_tile_grid_mask = X_tile * inv_mask + X_grid_mask

# Prepare 3rd of pai
Z_grid = np.ones_like(X_grid)
Z_grid = Z_grid - (X_tile_grid_mask+Y_grid)

# # グラフの高さを格納する
prob = np.zeros_like(Z_grid)

X_grid_expand = np.expand_dims(X_tile_grid_mask,2)
Y_grid_expand = np.expand_dims(Y_grid,2)
Z_grid_expand = np.expand_dims(Z_grid,2)

tensorXYZ = np.concatenate([X_grid_expand,Y_grid_expand,
                            Z_grid_expand],axis=-1)

# alpha to be drawn
alphas = [[10,10,10],
         [0.5,0.5,0.5],
         [1,1,1]]

# Create probability every alpha
prob_list = []
for alpha in alphas:
    for i in range(_mask.shape[0]):
        for j in range(_mask.shape[1]):
            prob[i,j] = dirichlet.pdf(x=tensorXYZ[i,j,:],alpha=alpha)
    prob_list.append(prob)
    prob = np.zeros_like(Z_grid)
        
# Plot dirichlet distribution for each alpha
fig = plt.figure(figsize=(13,10))
for i in range(len(alpha)):
    ax = fig.add_subplot(221+i,projection='3d')
    ax.plot_surface(X_tile_grid_mask,Y_grid,prob_list[i],
                    cmap=cm.coolwarm,antialiased=False)
    ax.set_title('Alpha = {}'.format(alphas[i]))
    ax.set_zlim(0,25)