main.py

import numpy as np

'''
Implementation adapted from Agustinus Kristiadi's Blog:
https://agustinus.kristia.de/techblog/2017/09/07/lda-gibbs/
----------------------
All the parameters correspond to Finding scientific topics 2004. 
n_t: Number of topics
W: The collection of all the words
n_w: Number of words in dictionary
n_d: Number of documents

Theta ~ Dirichlet(alpha), document-topic distribution
Phi ~ Dirichlet(beta), topic-word distribution

X: corpus
Z: word-topic assignment, shape (D,n_w)

n_jw: The number of word w assigned to topic j in Z
n_ja: The total number of word in topic j in Z
n_jd: The number of words in document d assigend to j
n_ad: The number of words in document d 
'''

# A helper function that assigns the initial Z with the initial Phi
def wt_from_Phi(Phi, w):
    wt_dist = Phi[:,w]
    wt_dist = wt_dist / np.sum(wt_dist)
    return wt_dist


def lda_gibbs_BOW(n_d, n_t, W, Theta, Phi, X, Z, alpha, beta, iterations=1000):
    n_w = len(W)
    
    # Initialize count matrices
    n_jw = np.zeros((n_t, n_w))  # topic-word
    n_ja = np.zeros(n_t)         # total words per topic
    n_jd = np.zeros((n_d, n_t))  # document-topic
    n_ad = np.zeros(n_d)         # total words per document

    # Fill in initial counts from Z
    for d_idx, doc in enumerate(Z):
        for w_idx, topic_counts in enumerate(doc):
            word_id = X[d_idx][w_idx][0]
            count_list = topic_counts
            for topic, c in enumerate(count_list):
                if c > 0:
                    n_jw[topic, word_id] += c
                    n_ja[topic] += c
                    n_jd[d_idx, topic] += c
                    n_ad[d_idx] += c

    for it in range(iterations):
        for d_idx, doc in enumerate(X):
            for w_idx, (word_id, count) in enumerate(doc):
                current_topics = Z[d_idx][w_idx]
                for topic in range(n_t):
                    c = current_topics[topic]
                    if c > 0:
                        n_jw[topic, word_id] -= c
                        n_ja[topic] -= c
                        n_jd[d_idx, topic] -= c
                        n_ad[d_idx] -= c

                # Re-sample new topic assignment
                p = np.zeros(n_t)
                for t in range(n_t):
                    p[t] = ((n_jw[t, word_id] + beta) / (n_ja[t] + beta * n_w)) * \
                           ((n_jd[d_idx, t] + alpha) / (n_ad[d_idx] + alpha * n_t))
                p /= p.sum()

                new_assignment = np.random.multinomial(count, p)
                Z[d_idx][w_idx] = new_assignment

                for t in range(n_t):
                    c = new_assignment[t]
                    n_jw[t, word_id] += c
                    n_ja[t] += c
                    n_jd[d_idx, t] += c
                    n_ad[d_idx] += c

    # Compute final Theta and Phi
    for d in range(n_d):
        Theta[d] = (n_jd[d] + alpha) / (n_ad[d] + alpha * n_t)
    for t in range(n_t):
        Phi[t] = (n_jw[t] + beta) / (n_ja[t] + beta * n_w)

    return Theta, Phi, Z


def main():
    ### Step 0: Initialize our example corpus. Imagine each integer from 0-4 as a word. For example, 0="football", 1="basketball", 2="tennis", 3="movie", 4="art".
    # Vocabulary - all the words
    W = np.array([0, 1, 2, 3, 4])

    # Toy sample documents
    X = [
        [0, 0, 1, 2, 2],  # Topic 1
        [0, 0, 1, 1, 1],  # Topic 1
        [0, 1, 2, 2, 2],  # Topic 1
        [2, 2, 1, 1, 4],  # Topic 2
        [4, 4, 4, 4, 4],  # Topic 2
        [3, 3, 4, 4, 4],  # Topic 2
        [3, 4, 4, 4, 4],  # Topic 2
        [3, 3, 3, 4, 1],  # Topic 2
        [4, 4, 3, 3, 2],  # Topic 2
    ]
    # X in BOW format:
    # [[(0, 2), (1, 1), (2, 2)],
    #  [(0, 2), (1, 3)],
    #  [(0, 1), (1, 1), (2, 3)],
    #  [(1, 2), (2, 2), (4, 1)],
    #  [(4, 5)],
    #  [(3, 2), (4, 3)],
    #  [(3, 1), (4, 4)],
    #  [(1, 1), (3, 3), (4, 1)],
    #  [(2, 1), (3, 2), (4, 2)]]

    # Changing document to BOW format. Same as BOW format used by Gensim
    X_BOW = []
    for i in range(len(X)):
        doc = np.zeros(W.shape[0])
        for j in X[i]:
            doc[j] += 1
        doc2 = []
        for j in range(W.shape[0]):
            if doc[j] != 0:
                doc2.append((j,int(doc[j])))
        X_BOW.append(doc2)
    X = X_BOW

    # Necesary parameters
    n_d = len(X)  # num of docs
    n_w = W.shape[0]  # num of words / size of vocabulary
    n_t = 2  # num of topics (Note that you can also set it above 2 and give it a try!)



    ### Step 1: Initialize the Dirichlet parameters. 
    # Dirichlet priors
    alpha = 1/n_t   # Parameter for Theta, the document-topic distribution
    beta = 1/n_t    # Phi, topic-word
    iterations = 1000

    # Theta := document-topic distribution. Initialized in Dirichlet(alpha) distribution
    Theta = np.zeros((n_d,n_t), dtype=float)
    for i in range(n_d):
        Theta[i] = np.random.dirichlet(alpha*np.ones(n_t))

    # Phi := word-topic distribution. Initialized in Dirichlet(beta) distribution
    Phi = np.zeros((n_t,n_w), dtype=float)
    for k in range(n_t):
        Phi[k] = np.random.dirichlet(beta*np.ones(n_w))

    # Z := word-topic assignment, aka real assignment to real words in corpus X.
    # Z[i,j]: tuple of length T, containing the number of words assigned to each topic
    # e.g. Z[i,j][0]: num of words assigned to topic 0 in document i word j from corpus X
    Z = []
    for doc in X:     # document index
        Z_doc = []
        for word in doc:     # word index
            w = word[0]
            count = word[1]
            wt_dist = wt_from_Phi(Phi, w)
            assigned_topic = np.random.multinomial(count, wt_dist)
            Z_doc.append(assigned_topic.tolist())
        Z.append(Z_doc)

    ### Step 2: Run Gibbs sampling and print the result
    Theta, Phi, Z = lda_gibbs_BOW(n_d, n_t, W, Theta, Phi, X, Z, alpha, beta, iterations=2000)
    Theta = np.round(Theta, 3)
    Phi = np.round(Phi, 3)

    print("Theta (document-topic distribution) is: ")
    print(Theta)
    print()
    print("Phi (topic-word distribution) is:")
    print(Phi)
    print()


if __name__ == "__main__":
    main()