Tweets-DP-BMM-FP.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\anacoda\\lib\\site-packages\\gensim\\utils.py:1197: UserWarning: detected Windows; aliasing chunkize to chunkize_serial\n",
      "  warnings.warn(\"detected Windows; aliasing chunkize to chunkize_serial\")\n"
     ]
    }
   ],
   "source": [
    "from sklearn.datasets import fetch_20newsgroups\n",
    "from nltk.stem.wordnet import WordNetLemmatizer\n",
    "from nltk.tokenize import RegexpTokenizer\n",
    "from nltk.corpus import stopwords\n",
    "from collections import defaultdict\n",
    "import gensim\n",
    "from InitKmeans import InitKmeans\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.mixture import GaussianMixture\n",
    "     \n",
    "def pre_processData(newsgroups_train):\n",
    "    tokenizer = RegexpTokenizer(r'\\w+')\n",
    "    for i in range(len(newsgroups_train)):\n",
    "        newsgroups_train[i] = newsgroups_train[i].lower()\n",
    "        #newsgroups_train[i] = tokenizer.tokenize(newsgroups_train[i])\n",
    "    #newsgroups_train = [[token for token in doc if not token.isdigit()] for doc in newsgroups_train]\n",
    "    newsgroups_train = [doc.split(' ') for doc in newsgroups_train]\n",
    "    return newsgroups_train\n",
    "from sklearn.datasets import fetch_20newsgroups\n",
    "from nltk.stem.wordnet import WordNetLemmatizer\n",
    "from nltk.tokenize import RegexpTokenizer\n",
    "from nltk.corpus import stopwords\n",
    "from collections import defaultdict\n",
    "import gensim\n",
    "from InitKmeans import InitKmeans\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.mixture import GaussianMixture\n",
    "import sys  \n",
    "import json\n",
    "if __name__ == \"__main__\":\n",
    "    newsgroups_train = list()\n",
    "    news_labels = {}\n",
    "    docID_username = {}\n",
    "    \n",
    "    i = 0\n",
    "    file_path=\"D:/part_news4_biterm.txt\"\n",
    "    with open(file_path) as fp:\n",
    "        lines = fp.read().split(\"\\n\")\n",
    "        for line in lines:\n",
    "            if line:\n",
    "                text = json.loads(line)[\"bitermText\"].strip()\n",
    "                label = json.loads(line)[\"clusterNo\"]\n",
    "                newsgroups_train.append(text)\n",
    "                news_labels[i] = label\n",
    "                i+=1\n",
    "    fp.close()\n",
    "    \n",
    "    corpus = pre_processData(newsgroups_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_biterm = {}\n",
    "for docID in range(len(corpus)):\n",
    "    words = corpus[docID]\n",
    "    duplicate_biterm = []\n",
    "    doc_biterm[docID] = {}\n",
    "    for biterm in words:\n",
    "        hash_biterm = set(biterm.split(','))\n",
    "        if hash_biterm not in duplicate_biterm:\n",
    "            duplicate_biterm.append(hash_biterm)\n",
    "            doc_biterm[docID][biterm] = 1\n",
    "        else:\n",
    "            if biterm in doc_biterm[docID].keys():\n",
    "                doc_biterm[docID][biterm] += 1\n",
    "            else:\n",
    "                doc_biterm[docID][biterm] = 1\n",
    "doc_words = {}\n",
    "i_biterm = 0\n",
    "for docID in range(len(corpus)):\n",
    "    doc_words[docID] = {}\n",
    "    words = corpus[docID]\n",
    "    for biterm in words:\n",
    "        for word in biterm.split(','):\n",
    "            if word not in doc_words[docID].keys():\n",
    "                doc_words[docID][word] = 0\n",
    "            doc_words[docID][word] += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "start  0  end  1895 topics:  77 truth:  89  NMI:  0.8174738724577841\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-22-d8c6944d1a0f>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m     95\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     96\u001b[0m                 \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 97\u001b[1;33m                     \u001b[0mselect_k\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mnp\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mrandom\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mchoice\u001b[0m\u001b[1;33m(\u001b[0m \u001b[1;33m(\u001b[0m\u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mTopics\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m+\u001b[0m\u001b[1;36m1\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m,\u001b[0m \u001b[1;36m1\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mp\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mnormalized_posterior\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m[\u001b[0m\u001b[1;36m0\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     98\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     99\u001b[0m                 \u001b[1;32mif\u001b[0m \u001b[0mselect_k\u001b[0m \u001b[1;33m==\u001b[0m \u001b[0mlen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mTopics\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "total_batch = 16\n",
    "NMI_batch = {}\n",
    "start = 0\n",
    "end = 0\n",
    "topic_batch = {}\n",
    "for batch_i in range(1,17):\n",
    "    if batch_i == 16:\n",
    "        end = int(len(corpus))\n",
    "    else:\n",
    "        end = int(len(corpus) / total_batch) * batch_i\n",
    "    \n",
    "    docID_assign_z = {}\n",
    "    m_z = {}\n",
    "    n_z = {}\n",
    "    n_w = {}\n",
    "    n_b = {}\n",
    "    Topics = []\n",
    "    V = set()\n",
    "    D = set()\n",
    "    alpha = 0.3\n",
    "    beta = 0.002\n",
    "    NMI_batch[batch_i] = 0\n",
    "    topic_batch[batch_i] = 0\n",
    "    \n",
    "    total_iter = 10\n",
    "    for iter in range(total_iter):\n",
    "        for docID in range(start,end):\n",
    "            words = corpus[docID]\n",
    "            D.discard(docID)\n",
    "            if docID in docID_assign_z.keys():\n",
    "                before_k = docID_assign_z[docID]\n",
    "                m_z[before_k].discard(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        n_z[before_k][word] -= 1\n",
    "                        n_w[before_k] -=1\n",
    "            else:\n",
    "                before_k = -1\n",
    "            if len(D) == 0 and len(V) == 0:\n",
    "                choose_k = 0\n",
    "                D.add(docID)\n",
    "                docID_assign_z[docID] = choose_k\n",
    "                if choose_k not in m_z.keys():\n",
    "                    m_z[choose_k] = set()\n",
    "                m_z[choose_k].add(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        if choose_k not in n_w.keys():\n",
    "                            n_w[choose_k] = 0\n",
    "                        if choose_k not in n_z.keys():\n",
    "                            n_z[choose_k] = {}\n",
    "                        if word not in n_z[choose_k].keys():\n",
    "                            n_z[choose_k][word] = 0\n",
    "                        n_z[choose_k][word] += 1\n",
    "                        n_w[choose_k] += 1\n",
    "                        V.add(word)\n",
    "                if choose_k not in Topics:\n",
    "                    Topics.append(choose_k)\n",
    "            else:\n",
    "                log_pro = []\n",
    "                for k in Topics:\n",
    "                    pro_k = len(m_z[k])\n",
    "                    if pro_k != 0:\n",
    "                        i = 0\n",
    "                        for biterm in words:\n",
    "                            each_word = biterm.strip().split(',')\n",
    "                            for word in each_word:\n",
    "                                if word not in n_z[k].keys():\n",
    "                                    n_z[k][word] = 0\n",
    "                            for j in range(doc_biterm[docID][biterm]):\n",
    "                                pro_k *= (n_z[k][each_word[0]] + n_z[k][each_word[-1]] + beta + j) / ( n_w[k] + len(V)*beta + i) \n",
    "                                i += 1\n",
    "\n",
    "                    if pro_k == 0:\n",
    "                        pro_k = sys.float_info.min\n",
    "                    log_pro.append(pro_k)\n",
    "\n",
    "                pro_new_k = alpha*(len(D))\n",
    "                i = 0\n",
    "                for biterm in words:\n",
    "                    for j in range(doc_biterm[docID][biterm]):\n",
    "                        pro_new_k *= ( beta + j) / ( len(V)*beta + i) \n",
    "                        i += 1\n",
    "\n",
    "                if pro_new_k == 0:\n",
    "                    pro_new_k = sys.float_info.min        \n",
    "                log_pro.append(pro_new_k)\n",
    "\n",
    "                sum_pro=sum(log_pro)\n",
    "\n",
    "                normalized_posterior = [i/sum_pro for i in log_pro]    \n",
    "                select_k = None\n",
    "                if iter == (total_iter - 1):\n",
    "                    select_k = normalized_posterior.index(max(normalized_posterior))\n",
    "\n",
    "                else:\n",
    "                    select_k = np.random.choice( (len(Topics)+1) , 1, p=normalized_posterior)[0]  \n",
    "\n",
    "                if select_k == len(Topics):\n",
    "                    choose_k = np.max(Topics) + 1\n",
    "                else:\n",
    "                    choose_k = Topics[select_k]\n",
    "\n",
    "                D.add(docID)\n",
    "                docID_assign_z[docID] = choose_k\n",
    "                if choose_k not in m_z.keys():\n",
    "                    m_z[choose_k] = set()\n",
    "                m_z[choose_k].add(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        if choose_k not in n_w.keys():\n",
    "                            n_w[choose_k] = 0\n",
    "                        if choose_k not in n_z.keys():\n",
    "                            n_z[choose_k] = {}\n",
    "                        if word not in n_z[choose_k].keys():\n",
    "                            n_z[choose_k][word] = 0\n",
    "                        n_z[choose_k][word] += 1\n",
    "                        n_w[choose_k] += 1\n",
    "                        V.add(word)\n",
    "                if choose_k not in Topics:\n",
    "                    Topics.append(choose_k)\n",
    "\n",
    "            count_k = []\n",
    "            for k in Topics:\n",
    "                if k in m_z.keys() and len(m_z[k]) == 0:\n",
    "                    m_z.pop(k, None)\n",
    "                    n_z.pop(k, None)\n",
    "                    n_w.pop(k, None)\n",
    "                    count_k.append(k)\n",
    "            for k in count_k:\n",
    "                Topics.remove(k)   \n",
    "        if iter == 9:\n",
    "            from sklearn.metrics.cluster import normalized_mutual_info_score\n",
    "            nmi_sample = []\n",
    "            nmi_result = []\n",
    "            for key, value in news_labels.items():\n",
    "                if key < end and key >= start:\n",
    "                    nmi_sample.append(value)\n",
    "                    nmi_result.append(docID_assign_z[key])\n",
    "            NMI_batch[batch_i] = normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result))\n",
    "            topic_batch[batch_i] = len(Topics)\n",
    "            print(\"start \",start, \" end \",end, \"topics: \",len(Topics),\"truth: \",len(np.unique(nmi_sample)), \" NMI: \",normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result)))\n",
    "    if batch_i != 1:\n",
    "        start = int(len(corpus) / total_batch) * (batch_i -1) + 1\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "start  0  end  1895 topics:  170 truth:  89  NMI:  0.82200209571536\n",
      "start  0  end  3790 topics:  220 truth:  98  NMI:  0.823954185001369\n",
      "start  1896  end  5685 topics:  213 truth:  96  NMI:  0.8199637234493342\n",
      "start  3791  end  7580 topics:  218 truth:  91  NMI:  0.8110248917572214\n",
      "start  5686  end  9475 topics:  217 truth:  143  NMI:  0.8315142185043477\n",
      "start  7581  end  11370 topics:  184 truth:  127  NMI:  0.8715208980672882\n",
      "start  9476  end  13265 topics:  139 truth:  90  NMI:  0.8950836140923991\n",
      "start  11371  end  15160 topics:  148 truth:  93  NMI:  0.8636605624730506\n",
      "start  13266  end  17055 topics:  172 truth:  97  NMI:  0.8481595064881445\n",
      "start  15161  end  18950 topics:  158 truth:  92  NMI:  0.8740242179908456\n",
      "start  17056  end  20845 topics:  145 truth:  92  NMI:  0.8882630760417017\n",
      "start  18951  end  22740 topics:  154 truth:  89  NMI:  0.9008843321539031\n",
      "start  20846  end  24635 topics:  186 truth:  115  NMI:  0.8583873328536438\n",
      "start  22741  end  26530 topics:  208 truth:  103  NMI:  0.8290221293591598\n",
      "start  24636  end  28425 topics:  137 truth:  44  NMI:  0.806728535675004\n",
      "start  26531  end  30322 topics:  135 truth:  45  NMI:  0.7837130812105415\n"
     ]
    }
   ],
   "source": [
    "total_batch = 16\n",
    "NMI_batch = {}\n",
    "start = 0\n",
    "end = 0\n",
    "topic_batch = {}\n",
    "for batch_i in range(1,17):\n",
    "    if batch_i == 16:\n",
    "        end = int(len(corpus))\n",
    "    else:\n",
    "        end = int(len(corpus) / total_batch) * batch_i\n",
    "    \n",
    "    docID_assign_z = {}\n",
    "    m_z = {}\n",
    "    n_z = {}\n",
    "    n_w = {}\n",
    "    n_b = {}\n",
    "    Topics = []\n",
    "    V = set()\n",
    "    D = set()\n",
    "    alpha = 0.3\n",
    "    beta = 0.002\n",
    "    NMI_batch[batch_i] = 0\n",
    "    topic_batch[batch_i] = 0\n",
    "    \n",
    "    total_iter = 10\n",
    "    for iter in range(total_iter):\n",
    "        for docID in range(start,end):\n",
    "            words = corpus[docID]\n",
    "            D.discard(docID)\n",
    "            if docID in docID_assign_z.keys():\n",
    "                before_k = docID_assign_z[docID]\n",
    "                m_z[before_k].discard(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        n_z[before_k][word] -= 1\n",
    "                        n_w[before_k] -=1\n",
    "            else:\n",
    "                before_k = -1\n",
    "            if len(D) == 0 and len(V) == 0:\n",
    "                choose_k = 0\n",
    "                D.add(docID)\n",
    "                docID_assign_z[docID] = choose_k\n",
    "                if choose_k not in m_z.keys():\n",
    "                    m_z[choose_k] = set()\n",
    "                m_z[choose_k].add(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        if choose_k not in n_w.keys():\n",
    "                            n_w[choose_k] = 0\n",
    "                        if choose_k not in n_z.keys():\n",
    "                            n_z[choose_k] = {}\n",
    "                        if word not in n_z[choose_k].keys():\n",
    "                            n_z[choose_k][word] = 0\n",
    "                        n_z[choose_k][word] += 1\n",
    "                        n_w[choose_k] += 1\n",
    "                        V.add(word)\n",
    "                if choose_k not in Topics:\n",
    "                    Topics.append(choose_k)\n",
    "            else:\n",
    "                log_pro = []\n",
    "                for k in Topics:\n",
    "                    pro_k = len(m_z[k])\n",
    "                    if pro_k != 0:\n",
    "                        i = 0\n",
    "                        for biterm in words:\n",
    "                            each_word = biterm.strip().split(',')\n",
    "                            for word in each_word:\n",
    "                                if word not in n_z[k].keys():\n",
    "                                    n_z[k][word] = 0\n",
    "                            for j in range(doc_biterm[docID][biterm]):\n",
    "                                pro_k *= (n_z[k][each_word[0]] + n_z[k][each_word[-1]] + beta + j) / ( n_w[k] + len(V)*beta + i) \n",
    "                                i += 1\n",
    "\n",
    "                    if pro_k == 0:\n",
    "                        pro_k = sys.float_info.min\n",
    "                    log_pro.append(pro_k)\n",
    "\n",
    "                pro_new_k = alpha*(len(D))\n",
    "                i = 0\n",
    "                for biterm in words:\n",
    "                    for j in range(doc_biterm[docID][biterm]):\n",
    "                        pro_new_k *= ( beta + j) / ( len(V)*beta + i) \n",
    "                        i += 1\n",
    "\n",
    "                if pro_new_k == 0:\n",
    "                    pro_new_k = sys.float_info.min        \n",
    "                log_pro.append(pro_new_k)\n",
    "\n",
    "                sum_pro=sum(log_pro)\n",
    "\n",
    "                normalized_posterior = [i/sum_pro for i in log_pro]    \n",
    "                select_k = None\n",
    "                if iter == (total_iter - 1):\n",
    "                    select_k = normalized_posterior.index(max(normalized_posterior))\n",
    "\n",
    "                else:\n",
    "                    select_k = np.random.choice( (len(Topics)+1) , 1, p=normalized_posterior)[0]  \n",
    "\n",
    "                if select_k == len(Topics):\n",
    "                    choose_k = np.max(Topics) + 1\n",
    "                else:\n",
    "                    choose_k = Topics[select_k]\n",
    "\n",
    "                D.add(docID)\n",
    "                docID_assign_z[docID] = choose_k\n",
    "                if choose_k not in m_z.keys():\n",
    "                    m_z[choose_k] = set()\n",
    "                m_z[choose_k].add(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        if choose_k not in n_w.keys():\n",
    "                            n_w[choose_k] = 0\n",
    "                        if choose_k not in n_z.keys():\n",
    "                            n_z[choose_k] = {}\n",
    "                        if word not in n_z[choose_k].keys():\n",
    "                            n_z[choose_k][word] = 0\n",
    "                        n_z[choose_k][word] += 1\n",
    "                        n_w[choose_k] += 1\n",
    "                        V.add(word)\n",
    "                if choose_k not in Topics:\n",
    "                    Topics.append(choose_k)\n",
    "\n",
    "            count_k = []\n",
    "            for k in Topics:\n",
    "                if k in m_z.keys() and len(m_z[k]) == 0:\n",
    "                    m_z.pop(k, None)\n",
    "                    n_z.pop(k, None)\n",
    "                    n_w.pop(k, None)\n",
    "                    count_k.append(k)\n",
    "            for k in count_k:\n",
    "                Topics.remove(k)   \n",
    "        if iter == 9:\n",
    "            from sklearn.metrics.cluster import normalized_mutual_info_score\n",
    "            nmi_sample = []\n",
    "            nmi_result = []\n",
    "            for key, value in news_labels.items():\n",
    "                if key < end and key >= start:\n",
    "                    nmi_sample.append(value)\n",
    "                    nmi_result.append(docID_assign_z[key])\n",
    "            NMI_batch[batch_i] = normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result))\n",
    "            topic_batch[batch_i] = len(Topics)\n",
    "            print(\"start \",start, \" end \",end, \"topics: \",len(Topics),\"truth: \",len(np.unique(nmi_sample)), \" NMI: \",normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result)))\n",
    "    if batch_i != 1:\n",
    "        start = int(len(corpus) / total_batch) * (batch_i -1) + 1\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]\n",
      "\n",
      "[0.82200209571536, 0.823954185001369, 0.8199637234493342, 0.8110248917572214, 0.8315142185043477, 0.8715208980672882, 0.8950836140923991, 0.8636605624730506, 0.8481595064881445, 0.8740242179908456, 0.8882630760417017, 0.9008843321539031, 0.8583873328536438, 0.8290221293591598, 0.806728535675004, 0.7837130812105415]\n",
      "\n",
      "0.8454941500520822\n"
     ]
    }
   ],
   "source": [
    "common_keys = []\n",
    "\n",
    "item_NMI = []\n",
    "for key,value in sorted(NMI_batch.items(), key=lambda x: x[0]):\n",
    "    common_keys.append(key)\n",
    "    item_NMI.append(value)\n",
    "print(common_keys)\n",
    "\n",
    "print()\n",
    "\n",
    "print(item_NMI)\n",
    "print()\n",
    "\n",
    "print(np.mean(item_NMI))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "start  0  end  1895 topics:  177 truth:  89  NMI:  0.8223536740332866\n",
      "[216.0135057865846]\n",
      "start  0  end  3790 topics:  224 truth:  98  NMI:  0.8177250805428806\n",
      "[216.0135057865846, 509.70495043912064]\n",
      "start  1896  end  5685 topics:  220 truth:  96  NMI:  0.8272888867666193\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729]\n",
      "start  3791  end  7580 topics:  205 truth:  91  NMI:  0.8151190449987157\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076]\n",
      "start  5686  end  9475 topics:  224 truth:  143  NMI:  0.8258527312415648\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191]\n",
      "start  7581  end  11370 topics:  175 truth:  127  NMI:  0.8684109674328697\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191, 592.1550351563897]\n",
      "start  9476  end  13265 topics:  145 truth:  90  NMI:  0.8898665292074673\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191, 592.1550351563897, 520.0472253329121]\n",
      "start  11371  end  15160 topics:  156 truth:  93  NMI:  0.8643331553586218\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191, 592.1550351563897, 520.0472253329121, 563.4922895116147]\n",
      "start  13266  end  17055 topics:  176 truth:  97  NMI:  0.8519742907647204\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191, 592.1550351563897, 520.0472253329121, 563.4922895116147, 536.2143372447404]\n",
      "start  15161  end  18950 topics:  164 truth:  92  NMI:  0.8714305009716794\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191, 592.1550351563897, 520.0472253329121, 563.4922895116147, 536.2143372447404, 508.4277754858658]\n",
      "start  17056  end  20845 topics:  151 truth:  92  NMI:  0.8888032583925958\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191, 592.1550351563897, 520.0472253329121, 563.4922895116147, 536.2143372447404, 508.4277754858658, 433.1330822813661]\n",
      "start  18951  end  22740 topics:  165 truth:  89  NMI:  0.8852987294404487\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191, 592.1550351563897, 520.0472253329121, 563.4922895116147, 536.2143372447404, 508.4277754858658, 433.1330822813661, 474.6584970113836]\n",
      "start  20846  end  24635 topics:  194 truth:  115  NMI:  0.8647326495455325\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191, 592.1550351563897, 520.0472253329121, 563.4922895116147, 536.2143372447404, 508.4277754858658, 433.1330822813661, 474.6584970113836, 533.1079717763341]\n",
      "start  22741  end  26530 topics:  196 truth:  103  NMI:  0.8162282383628474\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191, 592.1550351563897, 520.0472253329121, 563.4922895116147, 536.2143372447404, 508.4277754858658, 433.1330822813661, 474.6584970113836, 533.1079717763341, 555.9029541823656]\n",
      "start  24636  end  28425 topics:  136 truth:  44  NMI:  0.8077132353743824\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191, 592.1550351563897, 520.0472253329121, 563.4922895116147, 536.2143372447404, 508.4277754858658, 433.1330822813661, 474.6584970113836, 533.1079717763341, 555.9029541823656, 362.1147518620828]\n",
      "start  26531  end  30322 topics:  139 truth:  45  NMI:  0.7804754075461054\n",
      "[216.0135057865846, 509.70495043912064, 487.1394209779729, 443.31294940791076, 693.9138134369191, 592.1550351563897, 520.0472253329121, 563.4922895116147, 536.2143372447404, 508.4277754858658, 433.1330822813661, 474.6584970113836, 533.1079717763341, 555.9029541823656, 362.1147518620828, 276.6750330290688]\n"
     ]
    }
   ],
   "source": [
    "import timeit\n",
    "total_batch = 16\n",
    "NMI_batch = {}\n",
    "start = 0\n",
    "end = 0\n",
    "topic_batch = {}\n",
    "time_iteration_list = []\n",
    "for batch_i in range(1,17):\n",
    "    if batch_i == 16:\n",
    "        end = int(len(corpus))\n",
    "    else:\n",
    "        end = int(len(corpus) / total_batch) * batch_i\n",
    "    \n",
    "    docID_assign_z = {}\n",
    "    m_z = {}\n",
    "    n_z = {}\n",
    "    n_w = {}\n",
    "    n_b = {}\n",
    "    Topics = []\n",
    "    V = set()\n",
    "    D = set()\n",
    "    alpha = 0.3\n",
    "    beta = 0.002\n",
    "    NMI_batch[batch_i] = 0\n",
    "    topic_batch[batch_i] = 0\n",
    "    \n",
    "    time_start = timeit.default_timer()\n",
    "    total_iter = 10\n",
    "    for iter in range(total_iter):\n",
    "        for docID in range(start,end):\n",
    "            words = corpus[docID]\n",
    "            D.discard(docID)\n",
    "            if docID in docID_assign_z.keys():\n",
    "                before_k = docID_assign_z[docID]\n",
    "                m_z[before_k].discard(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        n_z[before_k][word] -= 1\n",
    "                        n_w[before_k] -=1\n",
    "            else:\n",
    "                before_k = -1\n",
    "            if len(D) == 0 and len(V) == 0:\n",
    "                choose_k = 0\n",
    "                D.add(docID)\n",
    "                docID_assign_z[docID] = choose_k\n",
    "                if choose_k not in m_z.keys():\n",
    "                    m_z[choose_k] = set()\n",
    "                m_z[choose_k].add(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        if choose_k not in n_w.keys():\n",
    "                            n_w[choose_k] = 0\n",
    "                        if choose_k not in n_z.keys():\n",
    "                            n_z[choose_k] = {}\n",
    "                        if word not in n_z[choose_k].keys():\n",
    "                            n_z[choose_k][word] = 0\n",
    "                        n_z[choose_k][word] += 1\n",
    "                        n_w[choose_k] += 1\n",
    "                        V.add(word)\n",
    "                if choose_k not in Topics:\n",
    "                    Topics.append(choose_k)\n",
    "            else:\n",
    "                log_pro = []\n",
    "                for k in Topics:\n",
    "                    pro_k = len(m_z[k])\n",
    "                    if pro_k != 0:\n",
    "                        i = 0\n",
    "                        for biterm in words:\n",
    "                            each_word = biterm.strip().split(',')\n",
    "                            for word in each_word:\n",
    "                                if word not in n_z[k].keys():\n",
    "                                    n_z[k][word] = 0\n",
    "                            for j in range(doc_biterm[docID][biterm]):\n",
    "                                pro_k *= (n_z[k][each_word[0]] + n_z[k][each_word[-1]] + beta + j) / ( n_w[k] + len(V)*beta + i) \n",
    "                                i += 1\n",
    "\n",
    "                    if pro_k == 0:\n",
    "                        pro_k = sys.float_info.min\n",
    "                    log_pro.append(pro_k)\n",
    "\n",
    "                pro_new_k = alpha*(len(D))\n",
    "                i = 0\n",
    "                for biterm in words:\n",
    "                    for j in range(doc_biterm[docID][biterm]):\n",
    "                        pro_new_k *= ( beta + j) / ( len(V)*beta + i) \n",
    "                        i += 1\n",
    "\n",
    "                if pro_new_k == 0:\n",
    "                    pro_new_k = sys.float_info.min        \n",
    "                log_pro.append(pro_new_k)\n",
    "\n",
    "                sum_pro=sum(log_pro)\n",
    "\n",
    "                normalized_posterior = [i/sum_pro for i in log_pro]    \n",
    "                select_k = None\n",
    "                if iter == (total_iter - 1):\n",
    "                    select_k = normalized_posterior.index(max(normalized_posterior))\n",
    "\n",
    "                else:\n",
    "                    select_k = np.random.choice( (len(Topics)+1) , 1, p=normalized_posterior)[0]  \n",
    "\n",
    "                if select_k == len(Topics):\n",
    "                    choose_k = np.max(Topics) + 1\n",
    "                else:\n",
    "                    choose_k = Topics[select_k]\n",
    "\n",
    "                D.add(docID)\n",
    "                docID_assign_z[docID] = choose_k\n",
    "                if choose_k not in m_z.keys():\n",
    "                    m_z[choose_k] = set()\n",
    "                m_z[choose_k].add(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        if choose_k not in n_w.keys():\n",
    "                            n_w[choose_k] = 0\n",
    "                        if choose_k not in n_z.keys():\n",
    "                            n_z[choose_k] = {}\n",
    "                        if word not in n_z[choose_k].keys():\n",
    "                            n_z[choose_k][word] = 0\n",
    "                        n_z[choose_k][word] += 1\n",
    "                        n_w[choose_k] += 1\n",
    "                        V.add(word)\n",
    "                if choose_k not in Topics:\n",
    "                    Topics.append(choose_k)\n",
    "\n",
    "            count_k = []\n",
    "            for k in Topics:\n",
    "                if k in m_z.keys() and len(m_z[k]) == 0:\n",
    "                    m_z.pop(k, None)\n",
    "                    n_z.pop(k, None)\n",
    "                    n_w.pop(k, None)\n",
    "                    count_k.append(k)\n",
    "            for k in count_k:\n",
    "                Topics.remove(k)   \n",
    "        if iter == 9:\n",
    "            from sklearn.metrics.cluster import normalized_mutual_info_score\n",
    "            nmi_sample = []\n",
    "            nmi_result = []\n",
    "            for key, value in news_labels.items():\n",
    "                if key < end and key >= start:\n",
    "                    nmi_sample.append(value)\n",
    "                    nmi_result.append(docID_assign_z[key])\n",
    "            NMI_batch[batch_i] = normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result))\n",
    "            topic_batch[batch_i] = len(Topics)\n",
    "            print(\"start \",start, \" end \",end, \"topics: \",len(Topics),\"truth: \",len(np.unique(nmi_sample)), \" NMI: \",normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result)))\n",
    "    if batch_i != 1:\n",
    "        start = int(len(corpus) / total_batch) * (batch_i -1) + 1\n",
    "    time_stop = timeit.default_timer()\n",
    "    time_iteration_list.append(time_stop-time_start)\n",
    "    print(time_iteration_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "0.8574291521104326"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]\n",
      "\n",
      "[0.8323176994066535, 0.8321619534393847, 0.8260085749718509, 0.8363347268544267, 0.8466603284526395, 0.8880658911998899, 0.9046168269190166, 0.8685988542157779, 0.8631199199499224, 0.883549612851989, 0.8928763264307683, 0.9066801013701994, 0.8779024429094671, 0.8429377961401533, 0.828917548973653, 0.7881178296811294]\n",
      "\n",
      "0.8574291521104326\n"
     ]
    }
   ],
   "source": [
    "common_keys = []\n",
    "\n",
    "item_NMI = []\n",
    "for key,value in sorted(NMI_batch.items(), key=lambda x: x[0]):\n",
    "    common_keys.append(key)\n",
    "    item_NMI.append(value)\n",
    "print(common_keys)\n",
    "\n",
    "print()\n",
    "\n",
    "print(item_NMI)\n",
    "print()\n",
    "\n",
    "print(np.mean(item_NMI))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]\n",
      "\n",
      "[0.8323176994066535, 0.8321619534393847, 0.8260085749718509, 0.8363347268544267, 0.8466603284526395, 0.8880658911998899, 0.9046168269190166, 0.8685988542157779, 0.8631199199499224, 0.883549612851989, 0.8928763264307683, 0.9066801013701994, 0.8779024429094671, 0.8429377961401533, 0.828917548973653, 0.7881178296811294]\n",
      "\n",
      "0.8574291521104326\n"
     ]
    }
   ],
   "source": [
    "common_keys = []\n",
    "\n",
    "item_NMI = []\n",
    "for key,value in sorted(NMI_batch.items(), key=lambda x: x[0]):\n",
    "    common_keys.append(key)\n",
    "    item_NMI.append(value)\n",
    "print(common_keys)\n",
    "\n",
    "print()\n",
    "\n",
    "print(item_NMI)\n",
    "print()\n",
    "\n",
    "print(np.mean(item_NMI))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{1: 0.8278218254942188, 2: 0.8301163585429583, 3: 0.8276996464223356, 4: 0.8266799384278927, 5: 0.8361987253008136, 6: 0.8918132625586427, 7: 0.9143589365596025, 8: 0.8755057317179911, 9: 0.8761813446904034, 10: 0.8937849598731826, 11: 0.9038236974191605, 12: 0.9076669340448583, 13: 0.8717577659131356, 14: 0.8359866418427073, 15: 0.8274625028703796, 16: 0.7925454624242205}\n",
      "0.02\n",
      "dict_keys([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16])\n",
      "dict_values([0.8278218254942188, 0.8301163585429583, 0.8276996464223356, 0.8266799384278927, 0.8361987253008136, 0.8918132625586427, 0.9143589365596025, 0.8755057317179911, 0.8761813446904034, 0.8937849598731826, 0.9038236974191605, 0.9076669340448583, 0.8717577659131356, 0.8359866418427073, 0.8274625028703796, 0.7925454624242205])\n"
     ]
    }
   ],
   "source": [
    "print(NMI_batch)\n",
    "print(beta)\n",
    "print(NMI_batch.keys())\n",
    "print(NMI_batch.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "start  0  end  500 topics:  103 truth:  68  NMI:  0.8474288340533979\n",
      "start  500  end  1000 topics:  92 truth:  72  NMI:  0.8534737854841058\n",
      "start  1000  end  1500 topics:  98 truth:  69  NMI:  0.8428782378743177\n",
      "start  1500  end  2000 topics:  98 truth:  69  NMI:  0.856304660474624\n",
      "start  2000  end  2500 topics:  98 truth:  68  NMI:  0.8523640295032716\n",
      "start  2500  end  3000 topics:  86 truth:  68  NMI:  0.8860828017208211\n",
      "start  3000  end  3500 topics:  94 truth:  70  NMI:  0.8604437694483016\n",
      "start  3500  end  4000 topics:  92 truth:  72  NMI:  0.848645947842599\n",
      "start  4000  end  4500 topics:  86 truth:  72  NMI:  0.8691456112814188\n",
      "start  4500  end  5000 topics:  92 truth:  68  NMI:  0.8494774982467017\n",
      "start  5000  end  5500 topics:  92 truth:  58  NMI:  0.8848503817879018\n",
      "start  5500  end  6000 topics:  99 truth:  60  NMI:  0.8360336655931694\n",
      "start  6000  end  6500 topics:  90 truth:  61  NMI:  0.8215379347028716\n",
      "start  6500  end  7000 topics:  97 truth:  54  NMI:  0.8254042311193724\n",
      "start  7000  end  7500 topics:  51 truth:  43  NMI:  0.8444627793160893\n",
      "start  7500  end  8000 topics:  75 truth:  42  NMI:  0.8344013482122143\n",
      "start  8000  end  8500 topics:  60 truth:  44  NMI:  0.8464515618826784\n",
      "start  8500  end  9000 topics:  51 truth:  40  NMI:  0.8607882472150461\n",
      "start  9000  end  9500 topics:  80 truth:  56  NMI:  0.8749153294529834\n",
      "start  9500  end  10000 topics:  63 truth:  51  NMI:  0.8940796474350877\n",
      "start  10000  end  10500 topics:  69 truth:  60  NMI:  0.9070813908670783\n",
      "start  10500  end  11000 topics:  64 truth:  54  NMI:  0.8371736923048919\n",
      "start  11000  end  11500 topics:  74 truth:  58  NMI:  0.8810047661205772\n",
      "start  11500  end  12000 topics:  47 truth:  38  NMI:  0.826919043154165\n",
      "start  12000  end  12500 topics:  56 truth:  44  NMI:  0.902442159481986\n",
      "start  12500  end  13000 topics:  46 truth:  40  NMI:  0.8423354466586958\n",
      "start  13000  end  13500 topics:  46 truth:  44  NMI:  0.6791729818435948\n",
      "start  13500  end  14000 topics:  69 truth:  58  NMI:  0.8183968922783766\n",
      "start  14000  end  14500 topics:  73 truth:  65  NMI:  0.9141425537322692\n",
      "start  14500  end  15000 topics:  80 truth:  67  NMI:  0.8962373568774769\n",
      "start  15000  end  15500 topics:  84 truth:  61  NMI:  0.8131660109078047\n",
      "start  15500  end  16000 topics:  70 truth:  55  NMI:  0.8380398250494697\n",
      "start  16000  end  16500 topics:  90 truth:  66  NMI:  0.8893618937496223\n",
      "start  16500  end  17000 topics:  71 truth:  57  NMI:  0.9160824814403975\n",
      "start  17000  end  17500 topics:  55 truth:  44  NMI:  0.8907418727550385\n",
      "start  17500  end  18000 topics:  45 truth:  44  NMI:  0.944305858526963\n",
      "start  18000  end  18500 topics:  66 truth:  55  NMI:  0.8929317345623905\n",
      "start  18500  end  19000 topics:  84 truth:  67  NMI:  0.8879111091100845\n",
      "start  19000  end  19500 topics:  74 truth:  52  NMI:  0.9108470126712936\n",
      "start  19500  end  20000 topics:  77 truth:  59  NMI:  0.8999046104353755\n",
      "start  20000  end  20500 topics:  63 truth:  57  NMI:  0.9475372307047819\n",
      "start  20500  end  21000 topics:  64 truth:  55  NMI:  0.9023659397800806\n",
      "start  21000  end  21500 topics:  77 truth:  62  NMI:  0.9172529938675\n",
      "start  21500  end  22000 topics:  68 truth:  57  NMI:  0.8718172673303971\n",
      "start  22000  end  22500 topics:  42 truth:  33  NMI:  0.7379505050723713\n",
      "start  22500  end  23000 topics:  58 truth:  40  NMI:  0.879428069716036\n",
      "start  23000  end  23500 topics:  69 truth:  46  NMI:  0.8867116531778506\n",
      "start  23500  end  24000 topics:  84 truth:  67  NMI:  0.900130999851982\n",
      "start  24000  end  24500 topics:  64 truth:  35  NMI:  0.8817268553046393\n",
      "start  24500  end  25000 topics:  71 truth:  37  NMI:  0.8719628015275858\n",
      "start  25000  end  25500 topics:  65 truth:  37  NMI:  0.8272025960107012\n",
      "start  25500  end  26000 topics:  74 truth:  36  NMI:  0.8217122280540714\n",
      "start  26000  end  26500 topics:  77 truth:  37  NMI:  0.8112675635094753\n",
      "start  26500  end  27000 topics:  71 truth:  35  NMI:  0.8202820884090798\n",
      "start  27000  end  27500 topics:  67 truth:  35  NMI:  0.8733863746579276\n",
      "start  27500  end  28000 topics:  51 truth:  30  NMI:  0.9032069170040824\n",
      "start  28000  end  28500 topics:  66 truth:  36  NMI:  0.8423441912558113\n",
      "start  28500  end  29000 topics:  46 truth:  30  NMI:  0.7992065253243146\n",
      "start  29000  end  29500 topics:  48 truth:  28  NMI:  0.8032661106986742\n",
      "start  29500  end  30000 topics:  49 truth:  29  NMI:  0.7551789896348232\n",
      "start  30000  end  30322 topics:  46 truth:  30  NMI:  0.8272459163091533\n",
      "dict_keys([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60])\n",
      "dict_values([0.8474288340533979, 0.8534737854841058, 0.8428782378743177, 0.856304660474624, 0.8523640295032716, 0.8860828017208211, 0.8604437694483016, 0.848645947842599, 0.8691456112814188, 0.8494774982467017, 0.8848503817879018, 0.8360336655931694, 0.8215379347028716, 0.8254042311193724, 0.8444627793160893, 0.8344013482122143, 0.8464515618826784, 0.8607882472150461, 0.8749153294529834, 0.8940796474350877, 0.9070813908670783, 0.8371736923048919, 0.8810047661205772, 0.826919043154165, 0.902442159481986, 0.8423354466586958, 0.6791729818435948, 0.8183968922783766, 0.9141425537322692, 0.8962373568774769, 0.8131660109078047, 0.8380398250494697, 0.8893618937496223, 0.9160824814403975, 0.8907418727550385, 0.944305858526963, 0.8929317345623905, 0.8879111091100845, 0.9108470126712936, 0.8999046104353755, 0.9475372307047819, 0.9023659397800806, 0.9172529938675, 0.8718172673303971, 0.7379505050723713, 0.879428069716036, 0.8867116531778506, 0.900130999851982, 0.8817268553046393, 0.8719628015275858, 0.8272025960107012, 0.8217122280540714, 0.8112675635094753, 0.8202820884090798, 0.8733863746579276, 0.9032069170040824, 0.8423441912558113, 0.7992065253243146, 0.8032661106986742, 0.7551789896348232, 0.8272459163091533])\n"
     ]
    }
   ],
   "source": [
    "nmi_batch = {}\n",
    "topic_batch = {}\n",
    "start = 0\n",
    "end = 0\n",
    "total_batch = None\n",
    "if len(corpus) % 500 == 0:\n",
    "    total_batch = int(len(corpus) / 500)\n",
    "else:\n",
    "    total_batch = int(len(corpus) / 500) +1\n",
    "\n",
    "for batch in range(total_batch):\n",
    "    docID_assign_z = {}\n",
    "    m_z = {}\n",
    "    n_z = {}\n",
    "    n_w = {}\n",
    "    n_b = {}\n",
    "    Topics = []\n",
    "    V = set()\n",
    "    D = set()\n",
    "    alpha = 0.3\n",
    "    beta = 0.002\n",
    "    end = 500 * (batch + 1)\n",
    "    if end > len(corpus):\n",
    "        end = len(corpus)\n",
    "    nmi_batch[batch] = 0\n",
    "    topic_batch[batch] = 0\n",
    "    \n",
    "    total_iter = 10\n",
    "    for iter in range(total_iter):\n",
    "        for docID in range(start,end):\n",
    "            words = corpus[docID]\n",
    "            D.discard(docID)\n",
    "            if docID in docID_assign_z.keys():\n",
    "                before_k = docID_assign_z[docID]\n",
    "                m_z[before_k].discard(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        n_z[before_k][word] -= 1\n",
    "                        n_w[before_k] -=1\n",
    "            else:\n",
    "                before_k = -1\n",
    "            if len(D) == 0 and len(V) == 0:\n",
    "                choose_k = 0\n",
    "                D.add(docID)\n",
    "                docID_assign_z[docID] = choose_k\n",
    "                if choose_k not in m_z.keys():\n",
    "                    m_z[choose_k] = set()\n",
    "                m_z[choose_k].add(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        if choose_k not in n_w.keys():\n",
    "                            n_w[choose_k] = 0\n",
    "                        if choose_k not in n_z.keys():\n",
    "                            n_z[choose_k] = {}\n",
    "                        if word not in n_z[choose_k].keys():\n",
    "                            n_z[choose_k][word] = 0\n",
    "                        n_z[choose_k][word] += 1\n",
    "                        n_w[choose_k] += 1\n",
    "                        V.add(word)\n",
    "                if choose_k not in Topics:\n",
    "                    Topics.append(choose_k)\n",
    "            else:\n",
    "                log_pro = []\n",
    "                for k in Topics:\n",
    "                    pro_k = len(m_z[k])\n",
    "                    if pro_k != 0:\n",
    "                        i = 0\n",
    "                        for biterm in words:\n",
    "                            each_word = biterm.strip().split(',')\n",
    "                            for word in each_word:\n",
    "                                if word not in n_z[k].keys():\n",
    "                                    n_z[k][word] = 0\n",
    "                            for j in range(doc_biterm[docID][biterm]):\n",
    "                                pro_k *= (n_z[k][each_word[0]] + n_z[k][each_word[-1]] + beta + j) / ( n_w[k] + len(V)*beta + i) \n",
    "                                i += 1\n",
    "\n",
    "                    if pro_k == 0:\n",
    "                        pro_k = sys.float_info.min\n",
    "                    log_pro.append(pro_k)\n",
    "\n",
    "                pro_new_k = alpha*(len(D))\n",
    "                i = 0\n",
    "                for biterm in words:\n",
    "                    for j in range(doc_biterm[docID][biterm]):\n",
    "                        pro_new_k *= ( beta + j) / ( len(V)*beta + i) \n",
    "                        i += 1\n",
    "\n",
    "                if pro_new_k == 0:\n",
    "                    pro_new_k = sys.float_info.min        \n",
    "                log_pro.append(pro_new_k)\n",
    "\n",
    "                sum_pro=sum(log_pro)\n",
    "\n",
    "                normalized_posterior = [i/sum_pro for i in log_pro]    \n",
    "                select_k = None\n",
    "                if iter == (total_iter - 1):\n",
    "                    select_k = normalized_posterior.index(max(normalized_posterior))\n",
    "\n",
    "                else:\n",
    "                    select_k = np.random.choice( (len(Topics)+1) , 1, p=normalized_posterior)[0]  \n",
    "\n",
    "                if select_k == len(Topics):\n",
    "                    choose_k = np.max(Topics) + 1\n",
    "                else:\n",
    "                    choose_k = Topics[select_k]\n",
    "\n",
    "                D.add(docID)\n",
    "                docID_assign_z[docID] = choose_k\n",
    "                if choose_k not in m_z.keys():\n",
    "                    m_z[choose_k] = set()\n",
    "                m_z[choose_k].add(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        if choose_k not in n_w.keys():\n",
    "                            n_w[choose_k] = 0\n",
    "                        if choose_k not in n_z.keys():\n",
    "                            n_z[choose_k] = {}\n",
    "                        if word not in n_z[choose_k].keys():\n",
    "                            n_z[choose_k][word] = 0\n",
    "                        n_z[choose_k][word] += 1\n",
    "                        n_w[choose_k] += 1\n",
    "                        V.add(word)\n",
    "                if choose_k not in Topics:\n",
    "                    Topics.append(choose_k)\n",
    "\n",
    "            count_k = []\n",
    "            for k in Topics:\n",
    "                if k in m_z.keys() and len(m_z[k]) == 0:\n",
    "                    m_z.pop(k, None)\n",
    "                    n_z.pop(k, None)\n",
    "                    n_w.pop(k, None)\n",
    "                    count_k.append(k)\n",
    "            for k in count_k:\n",
    "                Topics.remove(k)   \n",
    "        if iter == 9:\n",
    "            from sklearn.metrics.cluster import normalized_mutual_info_score\n",
    "            nmi_sample = []\n",
    "            nmi_result = []\n",
    "            for key, value in news_labels.items():\n",
    "                if key < end and key >= start:\n",
    "                    nmi_sample.append(value)\n",
    "                    nmi_result.append(docID_assign_z[key])\n",
    "            nmi_batch[batch] = normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result))\n",
    "            topic_batch[batch] = len(Topics)\n",
    "            print(\"start \",start, \" end \",end, \"topics: \",len(Topics),\"truth: \",len(np.unique(nmi_sample)), \" NMI: \",normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result)))\n",
    "    start = end \n",
    " \n",
    "print(nmi_batch.keys())\n",
    "print(nmi_batch.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "iter  0  topic:  139\n",
      "iter  1  topic:  156\n",
      "iter  2  topic:  165\n",
      "iter  3  topic:  165\n",
      "iter  4  topic:  172\n",
      "iter  5  topic:  173\n",
      "iter  6  topic:  169\n",
      "iter  7  topic:  171\n",
      "iter  8  topic:  177\n",
      "iter  9  topic:  170\n",
      "start  0  end  1895 topics:  170 truth:  89  NMI:  0.8234967923953466\n",
      "iter  0  topic:  175\n",
      "iter  1  topic:  186\n",
      "iter  2  topic:  194\n",
      "iter  3  topic:  203\n",
      "iter  4  topic:  205\n",
      "iter  5  topic:  216\n",
      "iter  6  topic:  215\n",
      "iter  7  topic:  208\n",
      "iter  8  topic:  207\n",
      "iter  9  topic:  203\n",
      "start  0  end  3790 topics:  203 truth:  98  NMI:  0.8311417812056228\n",
      "iter  0  topic:  173\n",
      "iter  1  topic:  195\n",
      "iter  2  topic:  202\n",
      "iter  3  topic:  204\n",
      "iter  4  topic:  205\n",
      "iter  5  topic:  210\n",
      "iter  6  topic:  213\n",
      "iter  7  topic:  218\n",
      "iter  8  topic:  217\n",
      "iter  9  topic:  216\n",
      "start  1896  end  5685 topics:  216 truth:  96  NMI:  0.823754584299913\n",
      "iter  0  topic:  169\n",
      "iter  1  topic:  182\n",
      "iter  2  topic:  186\n",
      "iter  3  topic:  193\n",
      "iter  4  topic:  204\n",
      "iter  5  topic:  206\n",
      "iter  6  topic:  211\n",
      "iter  7  topic:  212\n",
      "iter  8  topic:  218\n",
      "iter  9  topic:  209\n",
      "start  3791  end  7580 topics:  209 truth:  91  NMI:  0.8224176781096412\n",
      "iter  0  topic:  182\n",
      "iter  1  topic:  201\n",
      "iter  2  topic:  215\n",
      "iter  3  topic:  215\n",
      "iter  4  topic:  219\n",
      "iter  5  topic:  214\n",
      "iter  6  topic:  223\n",
      "iter  7  topic:  219\n",
      "iter  8  topic:  217\n",
      "iter  9  topic:  214\n",
      "start  5686  end  9475 topics:  214 truth:  143  NMI:  0.8342011829877818\n",
      "iter  0  topic:  151\n",
      "iter  1  topic:  168\n",
      "iter  2  topic:  172\n",
      "iter  3  topic:  173\n",
      "iter  4  topic:  172\n",
      "iter  5  topic:  175\n",
      "iter  6  topic:  177\n",
      "iter  7  topic:  173\n",
      "iter  8  topic:  176\n",
      "iter  9  topic:  169\n",
      "start  7581  end  11370 topics:  169 truth:  127  NMI:  0.8807817337139942\n",
      "iter  0  topic:  125\n",
      "iter  1  topic:  133\n",
      "iter  2  topic:  138\n",
      "iter  3  topic:  138\n",
      "iter  4  topic:  145\n",
      "iter  5  topic:  146\n",
      "iter  6  topic:  143\n",
      "iter  7  topic:  144\n",
      "iter  8  topic:  151\n",
      "iter  9  topic:  141\n",
      "start  9476  end  13265 topics:  141 truth:  90  NMI:  0.892479261642998\n",
      "iter  0  topic:  133\n",
      "iter  1  topic:  140\n",
      "iter  2  topic:  147\n",
      "iter  3  topic:  147\n",
      "iter  4  topic:  150\n",
      "iter  5  topic:  144\n",
      "iter  6  topic:  147\n",
      "iter  7  topic:  153\n",
      "iter  8  topic:  153\n",
      "iter  9  topic:  152\n",
      "start  11371  end  15160 topics:  152 truth:  93  NMI:  0.8593393266744448\n",
      "iter  0  topic:  144\n",
      "iter  1  topic:  154\n",
      "iter  2  topic:  157\n",
      "iter  3  topic:  159\n",
      "iter  4  topic:  168\n",
      "iter  5  topic:  170\n",
      "iter  6  topic:  167\n",
      "iter  7  topic:  169\n",
      "iter  8  topic:  169\n",
      "iter  9  topic:  170\n",
      "start  13266  end  17055 topics:  170 truth:  97  NMI:  0.8513263110415927\n",
      "iter  0  topic:  138\n",
      "iter  1  topic:  152\n",
      "iter  2  topic:  152\n",
      "iter  3  topic:  153\n",
      "iter  4  topic:  156\n",
      "iter  5  topic:  161\n",
      "iter  6  topic:  161\n",
      "iter  7  topic:  161\n",
      "iter  8  topic:  159\n",
      "iter  9  topic:  158\n",
      "start  15161  end  18950 topics:  158 truth:  92  NMI:  0.8706397677696075\n",
      "iter  0  topic:  129\n",
      "iter  1  topic:  132\n",
      "iter  2  topic:  135\n",
      "iter  3  topic:  141\n",
      "iter  4  topic:  144\n",
      "iter  5  topic:  144\n",
      "iter  6  topic:  148\n",
      "iter  7  topic:  144\n",
      "iter  8  topic:  145\n",
      "iter  9  topic:  144\n",
      "start  17056  end  20845 topics:  144 truth:  92  NMI:  0.8869756524583201\n",
      "iter  0  topic:  137\n",
      "iter  1  topic:  145\n",
      "iter  2  topic:  147\n",
      "iter  3  topic:  148\n",
      "iter  4  topic:  148\n",
      "iter  5  topic:  154\n",
      "iter  6  topic:  150\n",
      "iter  7  topic:  152\n",
      "iter  8  topic:  161\n",
      "iter  9  topic:  154\n",
      "start  18951  end  22740 topics:  154 truth:  89  NMI:  0.887975516312373\n",
      "iter  0  topic:  154\n",
      "iter  1  topic:  168\n",
      "iter  2  topic:  169\n",
      "iter  3  topic:  178\n",
      "iter  4  topic:  176\n",
      "iter  5  topic:  180\n",
      "iter  6  topic:  186\n",
      "iter  7  topic:  186\n",
      "iter  8  topic:  193\n",
      "iter  9  topic:  182\n",
      "start  20846  end  24635 topics:  182 truth:  115  NMI:  0.8750403623475929\n",
      "iter  0  topic:  167\n",
      "iter  1  topic:  188\n",
      "iter  2  topic:  190\n",
      "iter  3  topic:  198\n",
      "iter  4  topic:  197\n",
      "iter  5  topic:  202\n",
      "iter  6  topic:  198\n",
      "iter  7  topic:  201\n",
      "iter  8  topic:  202\n",
      "iter  9  topic:  196\n",
      "start  22741  end  26530 topics:  196 truth:  103  NMI:  0.8292225144619587\n",
      "iter  0  topic:  132\n",
      "iter  1  topic:  142\n",
      "iter  2  topic:  140\n",
      "iter  3  topic:  143\n",
      "iter  4  topic:  144\n",
      "iter  5  topic:  142\n",
      "iter  6  topic:  141\n",
      "iter  7  topic:  141\n",
      "iter  8  topic:  139\n",
      "iter  9  topic:  137\n",
      "start  24636  end  28425 topics:  137 truth:  44  NMI:  0.795579346957113\n",
      "iter  0  topic:  117\n",
      "iter  1  topic:  120\n",
      "iter  2  topic:  130\n",
      "iter  3  topic:  137\n",
      "iter  4  topic:  131\n",
      "iter  5  topic:  129\n",
      "iter  6  topic:  126\n",
      "iter  7  topic:  134\n",
      "iter  8  topic:  130\n",
      "iter  9  topic:  129\n",
      "start  26531  end  30322 topics:  129 truth:  45  NMI:  0.7820984109771273\n"
     ]
    }
   ],
   "source": [
    "total_batch = 16\n",
    "NMI_batch = {}\n",
    "start = 0\n",
    "end = 0\n",
    "topic_batch = {}\n",
    "for batch_i in range(1,17):\n",
    "    if batch_i == 16:\n",
    "        end = int(len(corpus))\n",
    "    else:\n",
    "        end = int(len(corpus) / total_batch) * batch_i\n",
    "    \n",
    "    docID_assign_z = {}\n",
    "    m_z = {}\n",
    "    n_z = {}\n",
    "    n_w = {}\n",
    "    n_b = {}\n",
    "    Topics = []\n",
    "    V = set()\n",
    "    D = set()\n",
    "    alpha = 0.3\n",
    "    beta = 0.003\n",
    "    NMI_batch[batch_i] = 0\n",
    "    topic_batch[batch_i] = 0\n",
    "    \n",
    "    total_iter = 10\n",
    "    for iter in range(total_iter):\n",
    "        for docID in range(start,end):\n",
    "            words = corpus[docID]\n",
    "            D.discard(docID)\n",
    "            if docID in docID_assign_z.keys():\n",
    "                before_k = docID_assign_z[docID]\n",
    "                m_z[before_k].discard(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        n_z[before_k][word] -= 1\n",
    "                        n_w[before_k] -=1\n",
    "            else:\n",
    "                before_k = -1\n",
    "            if len(D) == 0 and len(V) == 0:\n",
    "                choose_k = 0\n",
    "                D.add(docID)\n",
    "                docID_assign_z[docID] = choose_k\n",
    "                if choose_k not in m_z.keys():\n",
    "                    m_z[choose_k] = set()\n",
    "                m_z[choose_k].add(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        if choose_k not in n_w.keys():\n",
    "                            n_w[choose_k] = 0\n",
    "                        if choose_k not in n_z.keys():\n",
    "                            n_z[choose_k] = {}\n",
    "                        if word not in n_z[choose_k].keys():\n",
    "                            n_z[choose_k][word] = 0\n",
    "                        n_z[choose_k][word] += 1\n",
    "                        n_w[choose_k] += 1\n",
    "                        V.add(word)\n",
    "                if choose_k not in Topics:\n",
    "                    Topics.append(choose_k)\n",
    "            else:\n",
    "                log_pro = []\n",
    "                for k in Topics:\n",
    "                    pro_k = len(m_z[k])\n",
    "                    if pro_k != 0:\n",
    "                        i = 0\n",
    "                        for biterm in words:\n",
    "                            each_word = biterm.strip().split(',')\n",
    "                            for word in each_word:\n",
    "                                if word not in n_z[k].keys():\n",
    "                                    n_z[k][word] = 0\n",
    "                            for j in range(doc_biterm[docID][biterm]):\n",
    "                                pro_k *= (n_z[k][each_word[0]] + n_z[k][each_word[-1]] + beta + j) / ( n_w[k] + len(V)*beta + i) \n",
    "                                i += 1\n",
    "\n",
    "                    if pro_k == 0:\n",
    "                        pro_k = sys.float_info.min\n",
    "                    log_pro.append(pro_k)\n",
    "\n",
    "                pro_new_k = alpha*(len(D))\n",
    "                i = 0\n",
    "                for biterm in words:\n",
    "                    for j in range(doc_biterm[docID][biterm]):\n",
    "                        pro_new_k *= ( beta + j) / ( len(V)*beta + i) \n",
    "                        i += 1\n",
    "\n",
    "                if pro_new_k == 0:\n",
    "                    pro_new_k = sys.float_info.min        \n",
    "                log_pro.append(pro_new_k)\n",
    "\n",
    "                sum_pro=sum(log_pro)\n",
    "\n",
    "                normalized_posterior = [i/sum_pro for i in log_pro]    \n",
    "                select_k = None\n",
    "                if iter == (total_iter - 1):\n",
    "                    select_k = normalized_posterior.index(max(normalized_posterior))\n",
    "\n",
    "                else:\n",
    "                    select_k = np.random.choice( (len(Topics)+1) , 1, p=normalized_posterior)[0]  \n",
    "\n",
    "                if select_k == len(Topics):\n",
    "                    choose_k = np.max(Topics) + 1\n",
    "                else:\n",
    "                    choose_k = Topics[select_k]\n",
    "\n",
    "                D.add(docID)\n",
    "                docID_assign_z[docID] = choose_k\n",
    "                if choose_k not in m_z.keys():\n",
    "                    m_z[choose_k] = set()\n",
    "                m_z[choose_k].add(docID)\n",
    "                for biterm in words:\n",
    "                    for word in biterm.split(','):\n",
    "                        if choose_k not in n_w.keys():\n",
    "                            n_w[choose_k] = 0\n",
    "                        if choose_k not in n_z.keys():\n",
    "                            n_z[choose_k] = {}\n",
    "                        if word not in n_z[choose_k].keys():\n",
    "                            n_z[choose_k][word] = 0\n",
    "                        n_z[choose_k][word] += 1\n",
    "                        n_w[choose_k] += 1\n",
    "                        V.add(word)\n",
    "                if choose_k not in Topics:\n",
    "                    Topics.append(choose_k)\n",
    "\n",
    "            count_k = []\n",
    "            for k in Topics:\n",
    "                if k in m_z.keys() and len(m_z[k]) == 0:\n",
    "                    m_z.pop(k, None)\n",
    "                    n_z.pop(k, None)\n",
    "                    n_w.pop(k, None)\n",
    "                    count_k.append(k)\n",
    "            for k in count_k:\n",
    "                Topics.remove(k)   \n",
    "        print(\"iter \",iter, \" topic: \", len(Topics))\n",
    "        if iter == 9:\n",
    "            from sklearn.metrics.cluster import normalized_mutual_info_score\n",
    "            nmi_sample = []\n",
    "            nmi_result = []\n",
    "            for key, value in news_labels.items():\n",
    "                if key < end and key >= start:\n",
    "                    nmi_sample.append(value)\n",
    "                    nmi_result.append(docID_assign_z[key])\n",
    "            NMI_batch[batch_i] = normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result))\n",
    "            topic_batch[batch_i] = len(Topics)\n",
    "            print(\"start \",start, \" end \",end, \"topics: \",len(Topics),\"truth: \",len(np.unique(nmi_sample)), \" NMI: \",normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result)))\n",
    "    if batch_i != 1:\n",
    "        start = int(len(corpus) / total_batch) * (batch_i -1) + 1\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]\n",
      "\n",
      "[0.8234967923953466, 0.8311417812056228, 0.823754584299913, 0.8224176781096412, 0.8342011829877818, 0.8807817337139942, 0.892479261642998, 0.8593393266744448, 0.8513263110415927, 0.8706397677696075, 0.8869756524583201, 0.887975516312373, 0.8750403623475929, 0.8292225144619587, 0.795579346957113, 0.7820984109771273]\n",
      "0.8466543889597142\n"
     ]
    }
   ],
   "source": [
    "common_keys = []\n",
    "\n",
    "item_NMI = []\n",
    "for key,value in sorted(NMI_batch.items(), key=lambda x: x[0]):\n",
    "    common_keys.append(key)\n",
    "    item_NMI.append(value)\n",
    "print(common_keys)\n",
    "\n",
    "print()\n",
    "\n",
    "print(item_NMI)\n",
    "print(np.mean(item_NMI))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]\n",
      "\n",
      "[0.8298927392436527, 0.8269852334907385, 0.8275801293034326, 0.8325725508167305, 0.8296941181191704, 0.8858025131592013, 0.9112544715777952, 0.882279072937081, 0.8625284434677941, 0.8870335110637954, 0.9027410653955591, 0.8983125116090817, 0.8757577934488942, 0.8330259394498267, 0.8076724003551009, 0.7923083439379209]\n",
      "0.855340052335986\n"
     ]
    }
   ],
   "source": [
    "common_keys = []\n",
    "\n",
    "item_NMI = []\n",
    "for key,value in sorted(NMI_batch.items(), key=lambda x: x[0]):\n",
    "    common_keys.append(key)\n",
    "    item_NMI.append(value)\n",
    "print(common_keys)\n",
    "\n",
    "print()\n",
    "\n",
    "print(item_NMI)\n",
    "print(np.mean(item_NMI))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "alpha  0.1  NMI:  0.857268408981715\n",
      "alpha  0.1  Topics:  123\n",
      "alpha  0.2  NMI:  0.8592148717596345\n",
      "alpha  0.2  Topics:  126\n",
      "alpha  0.3  NMI:  0.8575133376480126\n",
      "alpha  0.3  Topics:  127\n",
      "alpha  0.4  NMI:  0.8562499910589776\n",
      "alpha  0.4  Topics:  131\n",
      "alpha  0.5  NMI:  0.857533270578133\n",
      "alpha  0.5  Topics:  135\n",
      "alpha  0.6  NMI:  0.8575866690811085\n",
      "alpha  0.6  Topics:  136\n",
      "alpha  0.7  NMI:  0.8582720702474204\n",
      "alpha  0.7  Topics:  139\n",
      "alpha  0.8  NMI:  0.8570571645878108\n",
      "alpha  0.8  Topics:  139\n",
      "alpha  0.9  NMI:  0.8573575831418399\n",
      "alpha  0.9  Topics:  147\n",
      "alpha  1.0  NMI:  0.8579458438287924\n",
      "alpha  1.0  Topics:  147\n"
     ]
    }
   ],
   "source": [
    "alpha_NMI = {}\n",
    "alpha_topics = {}\n",
    "for alpha in np.arange(0.1,1.1,0.1):\n",
    "    alpha = np.around(alpha,decimals=1)\n",
    "    \n",
    "    total_batch = 16\n",
    "    NMI_batch = {}\n",
    "    start = 0\n",
    "    end = 0\n",
    "    topic_batch = {}\n",
    "    time_iteration_list = []\n",
    "    for batch_i in range(1,17):\n",
    "        if batch_i == 16:\n",
    "            end = int(len(corpus))\n",
    "        else:\n",
    "            end = int(len(corpus) / total_batch) * batch_i\n",
    "\n",
    "        docID_assign_z = {}\n",
    "        m_z = {}\n",
    "        n_z = {}\n",
    "        n_w = {}\n",
    "        n_b = {}\n",
    "        Topics = []\n",
    "        V = set()\n",
    "        D = set()\n",
    "        #alpha = 0.6\n",
    "        beta = 0.02\n",
    "        NMI_batch[batch_i] = 0\n",
    "        topic_batch[batch_i] = 0\n",
    "\n",
    "        time_start = timeit.default_timer()\n",
    "        total_iter = 10\n",
    "        for iter in range(total_iter):\n",
    "            for docID in range(start,end):\n",
    "                words = corpus[docID]\n",
    "                D.discard(docID)\n",
    "                if docID in docID_assign_z.keys():\n",
    "                    before_k = docID_assign_z[docID]\n",
    "                    m_z[before_k].discard(docID)\n",
    "                    for biterm in words:\n",
    "                        for word in biterm.split(','):\n",
    "                            n_z[before_k][word] -= 1\n",
    "                            n_w[before_k] -=1\n",
    "                else:\n",
    "                    before_k = -1\n",
    "                if len(D) == 0 and len(V) == 0:\n",
    "                    choose_k = 0\n",
    "                    D.add(docID)\n",
    "                    docID_assign_z[docID] = choose_k\n",
    "                    if choose_k not in m_z.keys():\n",
    "                        m_z[choose_k] = set()\n",
    "                    m_z[choose_k].add(docID)\n",
    "                    for biterm in words:\n",
    "                        for word in biterm.split(','):\n",
    "                            if choose_k not in n_w.keys():\n",
    "                                n_w[choose_k] = 0\n",
    "                            if choose_k not in n_z.keys():\n",
    "                                n_z[choose_k] = {}\n",
    "                            if word not in n_z[choose_k].keys():\n",
    "                                n_z[choose_k][word] = 0\n",
    "                            n_z[choose_k][word] += 1\n",
    "                            n_w[choose_k] += 1\n",
    "                            V.add(word)\n",
    "                    if choose_k not in Topics:\n",
    "                        Topics.append(choose_k)\n",
    "                else:\n",
    "                    log_pro = []\n",
    "                    for k in Topics:\n",
    "                        pro_k = len(m_z[k])\n",
    "                        if pro_k != 0:\n",
    "                            i = 0\n",
    "                            for biterm in words:\n",
    "                                each_word = biterm.strip().split(',')\n",
    "                                for word in each_word:\n",
    "                                    if word not in n_z[k].keys():\n",
    "                                        n_z[k][word] = 0\n",
    "                                for j in range(doc_biterm[docID][biterm]):\n",
    "                                    pro_k *= (n_z[k][each_word[0]] + n_z[k][each_word[-1]] + beta + j) / ( n_w[k] + len(V)*beta + i) \n",
    "                                    i += 1\n",
    "\n",
    "                        if pro_k == 0:\n",
    "                            pro_k = sys.float_info.min\n",
    "                        log_pro.append(pro_k)\n",
    "\n",
    "                    pro_new_k = alpha*(len(D))\n",
    "                    i = 0\n",
    "                    for biterm in words:\n",
    "                        for j in range(doc_biterm[docID][biterm]):\n",
    "                            pro_new_k *= ( beta + j) / ( len(V)*beta + i) \n",
    "                            i += 1\n",
    "\n",
    "                    if pro_new_k == 0:\n",
    "                        pro_new_k = sys.float_info.min        \n",
    "                    log_pro.append(pro_new_k)\n",
    "\n",
    "                    sum_pro=sum(log_pro)\n",
    "\n",
    "                    normalized_posterior = [i/sum_pro for i in log_pro]    \n",
    "                    select_k = None\n",
    "                    if iter == (total_iter - 1):\n",
    "                        select_k = normalized_posterior.index(max(normalized_posterior))\n",
    "\n",
    "                    else:\n",
    "                        select_k = np.random.choice( (len(Topics)+1) , 1, p=normalized_posterior)[0]  \n",
    "\n",
    "                    if select_k == len(Topics):\n",
    "                        choose_k = np.max(Topics) + 1\n",
    "                    else:\n",
    "                        choose_k = Topics[select_k]\n",
    "\n",
    "                    D.add(docID)\n",
    "                    docID_assign_z[docID] = choose_k\n",
    "                    if choose_k not in m_z.keys():\n",
    "                        m_z[choose_k] = set()\n",
    "                    m_z[choose_k].add(docID)\n",
    "                    for biterm in words:\n",
    "                        for word in biterm.split(','):\n",
    "                            if choose_k not in n_w.keys():\n",
    "                                n_w[choose_k] = 0\n",
    "                            if choose_k not in n_z.keys():\n",
    "                                n_z[choose_k] = {}\n",
    "                            if word not in n_z[choose_k].keys():\n",
    "                                n_z[choose_k][word] = 0\n",
    "                            n_z[choose_k][word] += 1\n",
    "                            n_w[choose_k] += 1\n",
    "                            V.add(word)\n",
    "                    if choose_k not in Topics:\n",
    "                        Topics.append(choose_k)\n",
    "\n",
    "                count_k = []\n",
    "                for k in Topics:\n",
    "                    if k in m_z.keys() and len(m_z[k]) == 0:\n",
    "                        m_z.pop(k, None)\n",
    "                        n_z.pop(k, None)\n",
    "                        n_w.pop(k, None)\n",
    "                        count_k.append(k)\n",
    "                for k in count_k:\n",
    "                    Topics.remove(k)   \n",
    "            if iter == 9:\n",
    "                from sklearn.metrics.cluster import normalized_mutual_info_score\n",
    "                nmi_sample = []\n",
    "                nmi_result = []\n",
    "                for key, value in news_labels.items():\n",
    "                    if key < end and key >= start:\n",
    "                        nmi_sample.append(value)\n",
    "                        nmi_result.append(docID_assign_z[key])\n",
    "                NMI_batch[batch_i] = normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result))\n",
    "                topic_batch[batch_i] = len(Topics)\n",
    "                #print(\"start \",start, \" end \",end, \"topics: \",len(Topics),\"truth: \",len(np.unique(nmi_sample)), \" NMI: \",normalized_mutual_info_score(np.array(nmi_sample), np.array(nmi_result)))\n",
    "        if batch_i != 1:\n",
    "            start = int(len(corpus) / total_batch) * (batch_i -1) + 1\n",
    "    \n",
    "    alpha_NMI[alpha] = np.mean(list(NMI_batch.values()))\n",
    "    alpha_topics[alpha] = int(np.mean(list(topic_batch.values())))\n",
    "    print(\"alpha \",alpha, \" NMI: \", alpha_NMI[alpha])\n",
    "    print(\"alpha \",alpha, \" Topics: \", alpha_topics[alpha])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_values([0.857268408981715, 0.8592148717596345, 0.8575133376480126, 0.8562499910589776, 0.857533270578133, 0.8575866690811085, 0.8582720702474204, 0.8570571645878108, 0.8573575831418399, 0.8579458438287924])\n",
      "dict_values([123, 126, 127, 131, 135, 136, 139, 139, 147, 147])\n"
     ]
    }
   ],
   "source": [
    "print(alpha_NMI.values())\n",
    "print(alpha_topics.values())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}