第20章 潜在狄利克雷分配 - 狄利克雷分布采样

import numpy as np

def dirichlet_sample(alpha, size=1):
    if np.isscalar(alpha):
        alpha = np.array([alpha])
    
    samples = np.zeros((size, len(alpha)))
    for i in range(size):
        gamma_samples = np.random.gamma(alpha, 1.0)
        samples[i] = gamma_samples / gamma_samples.sum()
    
    return samples.squeeze()

def dirichlet_pdf(x, alpha):
    from scipy.special import gamma as gamma_func
    
    if np.isscalar(alpha):
        alpha = np.array([alpha])
    
    B = np.prod(gamma_func(alpha)) / gamma_func(np.sum(alpha))
    return np.prod(np.power(x, alpha - 1)) / B

def multinomial_sample(p, n=1):
    return np.random.multinomial(n, p)

def estimate_topic_distribution(doc, topic_word, doc_topic_prior):
    n_topics = topic_word.shape[0]
    topic_dist = np.zeros(n_topics)
    
    for word in doc:
        for t in range(n_topics):
            topic_dist[t] += np.log(topic_word[t, word] + 1e-10)
    
    topic_dist = np.exp(topic_dist - topic_dist.max())
    topic_dist *= doc_topic_prior
    topic_dist /= topic_dist.sum()
    
    return topic_dist