Dirichlet pdf in python

dirichlet_pdf.py

import math
import numpy as np


# Another example at http://stackoverflow.com/questions/10658866/calculating-pdf-of-dirichlet-distribution-in-python


def gamma(x):
    return math.gamma(x)

def gamma_beauty_2d(a):
    return (gamma(a[0] + a[1] + a[2])) / ( gamma(a[0]) * gamma(a[1]) * gamma(a[2]) )

def dirichlet_pdf_2d(X,a):
    assert (X[:,0] > 0).all()
    assert (X[:,1] > 0).all()
    assert (1 - np.subtract(X[:,0],X[:,1]) > 0 ).all()
    assert len(a) == 3
    
    points = []
    for d in X:
         x = d[0]
         y = d[1]
         p = (x **(-1 + a[0]) * (1 - x - y)**(-1 + a[2]) * y**(-1 + a[1]) ) * gamma_beauty_2d(a)
         points.append(p)
    return np.array(points)

def beta_fun(alpha):
    product = 1
    alpha_sum = 0 
    for a in alpha:
        product = product * gamma(a)
        alpha_sum = alpha_sum + a
    beta = product /  gamma(alpha_sum)
    return beta
    
def dirichlet_product(x,a):
    K = x.shape[0]
    product = 1
    f_t = 1 
    for i in range(0,K):
        p = x[i] ** (a[i] - 1)
        product = product * p
        f_t = f_t - x[i]
    product = product * (f_t ** (-1 + a[K]))
    return product 


def dirichlet_pdf(X,a):
    beta_component = beta_fun(a)
    points = []
    for d in X:
        product = dirichlet_product(d,a)
        dirichlet = product * (1 / beta_component)
        points.append(dirichlet)
    return points
        
        
d = np.array([[.5,.5],[.2,.4]])
alpha = np.array([4,2,3])

x = dirichlet_pdf_2d(d,alpha)
x2 = dirichlet_pdf(d,alpha)

assert (x2 == x).all()