em.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pomegranate import *\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Trying bayes net in pomegranate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "#generate example data\n",
    "n = 10\n",
    "h = BernoulliDistribution(0.3).sample(n)\n",
    "w = BernoulliDistribution(0.6).sample(n)\n",
    "l = BernoulliDistribution(0.7).sample(n) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 1. Initialize parameters (uniform)\n",
    "\n",
    "height = DiscreteDistribution({'short':0.5, 'tall':0.5})\n",
    "weight = DiscreteDistribution({'low':0.5, 'high':0.5})\n",
    "sex = ConditionalProbabilityTable([\n",
    "    ['short', 'low', 'm', 0.5],\n",
    "    ['short', 'low', 'f', 0.5],\n",
    "    ['short', 'high', 'm', 0.5],\n",
    "    ['short', 'high', 'f', 0.5],\n",
    "    ['tall', 'low', 'm', 0.5],\n",
    "    ['tall', 'low', 'f', 0.5],\n",
    "    ['tall', 'high', 'm', 0.5],\n",
    "    ['tall', 'high', 'f', 0.5]], [height, weight] )\n",
    "hair = ConditionalProbabilityTable([\n",
    "    ['m', 'sh', 0.5],\n",
    "    ['m', 'long', 0.5],\n",
    "    ['f', 'sh', 0.5],\n",
    "    ['f', 'long', 0.5]], [sex] )\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# create state objects (nodes)\n",
    "\n",
    "s_h = State(height, name='height')\n",
    "s_w = State(weight, name='weight')\n",
    "s_s = State(sex, name='sex')\n",
    "s_l = State(hair, name='hair')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# create the model\n",
    "\n",
    "model = BayesianNetwork(\"Hair Length\")\n",
    "\n",
    "# Add the three states to the network \n",
    "\n",
    "model.add_states(s_h, s_w, s_s, s_l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# build the DAG\n",
    "\n",
    "model.add_transition(s_h, s_s)\n",
    "model.add_transition(s_w, s_s)\n",
    "model.add_transition(s_s, s_l)\n",
    "model.bake()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ {\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"tall\" :0.0,\n",
       "            \"short\" :1.0\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "},\n",
       "       {\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"high\" :0.0,\n",
       "            \"low\" :1.0\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "},\n",
       "       {\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"m\" :0.5,\n",
       "            \"f\" :0.5\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "},\n",
       "       {\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"sh\" :0.5,\n",
       "            \"long\" :0.5\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "}], dtype=object)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.predict_proba({'height':'short','weight':'low'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['height', 'weight', 'sex', 'hair']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "[state.name for state in model.states]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 1, 2], dtype=int32)"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s_s.distribution.column_idxs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([2, 3], dtype=int32)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s_l.distribution.column_idxs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A simple model for trying things "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "A = DiscreteDistribution({1:0.45,0:0.55})\n",
    "B = ConditionalProbabilityTable([\n",
    "    [1, 0, 0.3],\n",
    "    [1, 1, 0.7],\n",
    "    [0, 0, 0.4],\n",
    "    [0, 1, 0.6]], [A] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "s_A = State(A, name='A')\n",
    "s_B = State(B, name='B')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "test_model = BayesianNetwork(\"test\")\n",
    "test_model.add_states(s_A, s_B)\n",
    "test_model.add_transition(s_A, s_B)\n",
    "test_model.bake()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "name_list = [s.name for s in test_model.states]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "test_model.states[0].distribution.parameters[0][0]=0.99"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{0: 0.99, 1: 0.45}]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_model.states[0].distribution.parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{\n",
       "     \"distribution\" : {\n",
       "         \"frozen\" : false,\n",
       "         \"class\" : \"Distribution\",\n",
       "         \"parameters\" : [\n",
       "             {\n",
       "                 \"1\" : 0.45,\n",
       "                 \"0\" : 0.99\n",
       "             }\n",
       "         ],\n",
       "         \"name\" : \"DiscreteDistribution\"\n",
       "     },\n",
       "     \"class\" : \"State\",\n",
       "     \"weight\" : 1.0,\n",
       "     \"name\" : \"A\"\n",
       " }, {\n",
       "     \"distribution\" : {\n",
       "         \"table\" : [\n",
       "             [\n",
       "                 \"0\",\n",
       "                 \"0\",\n",
       "                 \"0.4\"\n",
       "             ],\n",
       "             [\n",
       "                 \"0\",\n",
       "                 \"1\",\n",
       "                 \"0.6\"\n",
       "             ],\n",
       "             [\n",
       "                 \"1\",\n",
       "                 \"0\",\n",
       "                 \"0.3\"\n",
       "             ],\n",
       "             [\n",
       "                 \"1\",\n",
       "                 \"1\",\n",
       "                 \"0.7\"\n",
       "             ]\n",
       "         ],\n",
       "         \"parents\" : [\n",
       "             {\n",
       "                 \"frozen\" : false,\n",
       "                 \"class\" : \"Distribution\",\n",
       "                 \"parameters\" : [\n",
       "                     {\n",
       "                         \"1\" : 0.45,\n",
       "                         \"0\" : 0.99\n",
       "                     }\n",
       "                 ],\n",
       "                 \"name\" : \"DiscreteDistribution\"\n",
       "             }\n",
       "         ],\n",
       "         \"class\" : \"Distribution\",\n",
       "         \"name\" : \"ConditionalProbabilityTable\"\n",
       "     },\n",
       "     \"class\" : \"State\",\n",
       "     \"weight\" : 1.0,\n",
       "     \"name\" : \"B\"\n",
       " }]"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_model.states"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# getting the conditional probability\n",
    "\n",
    "prob_B = test_model.predict_proba({'A':1})[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"1\" :0.6999999999999998,\n",
       "            \"0\" :0.3000000000000001\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prob_B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.69999999999999984"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prob_B.probability(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data = np.asarray([[0,1],[0,0],[1,1],[1,1],[0,0],[1,0],[0,0]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"1\" :0.42857142857142855,\n",
       "            \"0\" :0.5714285714285714\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "}"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s_A.distribution.marginal()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\n",
       "    \"states\" : [\n",
       "        {\n",
       "            \"distribution\" : {\n",
       "                \"frozen\" : false,\n",
       "                \"class\" : \"Distribution\",\n",
       "                \"parameters\" : [\n",
       "                    {\n",
       "                        \"1\" : 0.42857142857142855,\n",
       "                        \"0\" : 0.5714285714285714\n",
       "                    }\n",
       "                ],\n",
       "                \"name\" : \"DiscreteDistribution\"\n",
       "            },\n",
       "            \"class\" : \"State\",\n",
       "            \"weight\" : 1.0,\n",
       "            \"name\" : \"A\"\n",
       "        },\n",
       "        {\n",
       "            \"distribution\" : {\n",
       "                \"table\" : [\n",
       "                    [\n",
       "                        \"1\",\n",
       "                        \"0\",\n",
       "                        \"0.333333333333\"\n",
       "                    ],\n",
       "                    [\n",
       "                        \"1\",\n",
       "                        \"1\",\n",
       "                        \"0.666666666667\"\n",
       "                    ],\n",
       "                    [\n",
       "                        \"0\",\n",
       "                        \"0\",\n",
       "                        \"0.75\"\n",
       "                    ],\n",
       "                    [\n",
       "                        \"0\",\n",
       "                        \"1\",\n",
       "                        \"0.25\"\n",
       "                    ]\n",
       "                ],\n",
       "                \"parents\" : [\n",
       "                    {\n",
       "                        \"frozen\" : false,\n",
       "                        \"class\" : \"Distribution\",\n",
       "                        \"parameters\" : [\n",
       "                            {\n",
       "                                \"1\" : 0.42857142857142855,\n",
       "                                \"0\" : 0.5714285714285714\n",
       "                            }\n",
       "                        ],\n",
       "                        \"name\" : \"DiscreteDistribution\"\n",
       "                    }\n",
       "                ],\n",
       "                \"class\" : \"Distribution\",\n",
       "                \"name\" : \"ConditionalProbabilityTable\"\n",
       "            },\n",
       "            \"class\" : \"State\",\n",
       "            \"weight\" : 1.0,\n",
       "            \"name\" : \"B\"\n",
       "        }\n",
       "    ],\n",
       "    \"class\" : \"BayesianNetwork\",\n",
       "    \"structure\" : [\n",
       "        [],\n",
       "        [\n",
       "            0\n",
       "        ]\n",
       "    ],\n",
       "    \"name\" : \"test\"\n",
       "}"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_model.fit(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ {\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"1\" :0.42857142857142866,\n",
       "            \"0\" :0.5714285714285714\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "},\n",
       "       {\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"1\" :0.4285714285714286,\n",
       "            \"0\" :0.5714285714285713\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "}], dtype=object)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_model.predict_proba()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pomegranate.BayesianNetwork.BayesianNetwork"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(test_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Model for testing 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "A = DiscreteDistribution({'1':0.2, '0':0.8})\n",
    "H = ConditionalProbabilityTable([\n",
    "        ['0', '0', 0.4],\n",
    "        ['0', '1', 0.6],\n",
    "        ['1', '0', 0.3],\n",
    "        ['1', '1', 0.7]], [A])\n",
    "B = ConditionalProbabilityTable([\n",
    "        ['0', '0', 0.6],\n",
    "        ['0', '1', 0.4],\n",
    "        ['1', '0', 0.5],\n",
    "        ['1', '1', 0.5]], [H])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "s_A = State(A, name='A')\n",
    "s_H = State(H, name='H')\n",
    "s_B = State(B, name='B')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "test_model2 = BayesianNetwork(\"test2\")\n",
    "test_model2.add_states(s_A, s_H, s_B)\n",
    "test_model2.add_transition(s_A, s_H)\n",
    "test_model2.add_transition(s_H, s_B)\n",
    "test_model2.bake()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data2 = np.array([[0, np.nan, 0],\n",
    "                   [0, np.nan, 1],\n",
    "                   [1, np.nan, 0],\n",
    "                   [1, np.nan, 1],\n",
    "                   [1, np.nan, 0],\n",
    "                   [0, np.nan, 1],\n",
    "                   [1, np.nan, 1],\n",
    "                   [1, np.nan, 0]], dtype='str')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([['0', '0'],\n",
       "       ['0', '1'],\n",
       "       ['1', '0'],\n",
       "       ['1', '1'],\n",
       "       ['1', '0'],\n",
       "       ['0', '1'],\n",
       "       ['1', '1'],\n",
       "       ['1', '0']], \n",
       "      dtype='|S3')"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data2[:, [0,2]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Test:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mb2 = MarkovBlanket(1)\n",
    "mb2.populate(test_model2)\n",
    "mb2.calculate_prob(test_model2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{1: {('0', '0'): 0.4, ('0', '1'): 0.6, ('1', '0'): 0.3, ('1', '1'): 0.7},\n",
       " 2: {('0', '0'): 0.6, ('0', '1'): 0.4, ('1', '0'): 0.5, ('1', '1'): 0.5}}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mb2.prob_table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "expected_counts2 = create_expected_counts(test_model2, mb2) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{(0, 0, 0): 0,\n",
       " (0, 0, 1): 0,\n",
       " (0, 1, 0): 0,\n",
       " (0, 1, 1): 0,\n",
       " (1, 0, 0): 0,\n",
       " (1, 0, 1): 0,\n",
       " (1, 1, 0): 0,\n",
       " (1, 1, 1): 0}"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expected_counts2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{(0, 1): 0.46, (1, 0): 0.53, (0, 0): 0.54, (1, 1): 0.47}\n"
     ]
    }
   ],
   "source": [
    "expected_counts2 = update_expected_counts(expected_counts2, test_model2, mb2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{(0, 0, 0): 0.4444444444444444,\n",
       " (0, 0, 1): 0.3478260869565218,\n",
       " (0, 1, 0): 0.33962264150943394,\n",
       " (0, 1, 1): 0.2553191489361702,\n",
       " (1, 0, 0): 0.5555555555555555,\n",
       " (1, 0, 1): 0.6521739130434782,\n",
       " (1, 1, 0): 0.660377358490566,\n",
       " (1, 1, 1): 0.7446808510638298}"
      ]
     },
     "execution_count": 111,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expected_counts2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0]"
      ]
     },
     "execution_count": 105,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mb2.parents"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the EM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Steps:\n",
    "1. Initialize parameters with near uniform values.\n",
    "2. Build the model with the initialized parameters. \n",
    "3. Expectation: Calculate the expected counts for the hidden variables. Update the data for the hidden variables. \n",
    "4. Maximization: Calculate MLE for all parameters by summing over expected counts. \n",
    "5. Update the model with the new parameter values.\n",
    "6. Iterate 3-5 until converges.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def search_hidden(data):\n",
    "    \"\"\"Returns the column index of the hidden node if only one column of NaN.\n",
    "    Only works if data is numeric.\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    data : An ndarray (n_sample, n_nodes)\n",
    "    \n",
    "    Returns\n",
    "    -------\n",
    "    ind_h : the index of the hidden node column\n",
    "    \"\"\"\n",
    "    is_col_nan = np.all(np.isnan(data), axis=0)\n",
    "    ind = np.where(is_col_nan)\n",
    "    if np.size(ind)==1:\n",
    "        ind_h = ind[0][0]\n",
    "    else:\n",
    "        raise ValueError('Data contains more than one hidden nodes or no hidden node')\n",
    "    return ind_h\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class MarkovBlanket():\n",
    "    \"\"\"\n",
    "    An object for storing info on nodes within the markov blanket of the hidden node\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    ind_h : int\n",
    "        index of the hidden node within the model\n",
    "        \n",
    "    Attributes\n",
    "    ----------\n",
    "    hidden : int\n",
    "        index of the hidden node\n",
    "        \n",
    "    parents : list of int\n",
    "        a list of indices of the parent nodes \n",
    "        \n",
    "    children : list of int\n",
    "        a list of indices of the children nodes\n",
    "        \n",
    "    coparents : list of int\n",
    "        a list of indices of the coparent nodes\n",
    "        \n",
    "    prob_table : dict\n",
    "        a dict of probabilities table of nodes within the Markov blanket\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, ind_h):\n",
    "        self.hidden = ind_h\n",
    "        self.parents = []\n",
    "        self.children = []\n",
    "        self.coparents = []\n",
    "        self.prob_table = {}\n",
    "        \n",
    "    def populate(self, model):\n",
    "        \"\"\"populate the parents, children, and coparents nodes\n",
    "        \"\"\"\n",
    "        state_indices = {state.name : i for i, state in enumerate(model.states)}\n",
    "        \n",
    "        edges_list = [(parent.name, child.name) for parent, child in model.edges]\n",
    "        edges_list = [(state_indices[parent],state_indices[child]) \n",
    "                  for parent, child in edges_list]\n",
    "        \n",
    "        self.children = list(set([child for parent, child in edges_list if parent==self.hidden]))\n",
    "        self.parents = list(set([parent for parent, child in edges_list if child==self.hidden]))\n",
    "        self.coparents = list(set([parent for parent, child in edges_list if child in self.children]))\n",
    "        try:\n",
    "            self.coparents.remove(self.hidden)\n",
    "        except ValueError:\n",
    "            pass\n",
    "            \n",
    "    def calculate_prob(self, model):\n",
    "        \"\"\"Create the probability table from nodes\n",
    "        \"\"\"\n",
    "        for ind_state in [self.hidden]+self.children:\n",
    "            distribution = model.states[ind_state].distribution\n",
    "            \n",
    "            if isinstance(distribution, ConditionalProbabilityTable):\n",
    "                table = list(distribution.parameters[0]) # make a copy\n",
    "                self.prob_table[ind_state] = {\n",
    "                    tuple(row[:-1]) : row[-1] for row in table}\n",
    "            else:\n",
    "                self.prob_table[ind_state] = dict(distribution.parameters[0]) # make a copy\n",
    "                \n",
    "    def update_prob(self, model, expected_counts, ct):\n",
    "        \"\"\"Update the probability table using expected counts\n",
    "        \"\"\"\n",
    "        ind = {x : i for i, x in enumerate([self.hidden] + self.parents + self.children + self.coparents)}\n",
    "        mb_keys = expected_counts.counts.keys()\n",
    "        \n",
    "        for ind_state in [self.hidden] + self.children:\n",
    "            distribution = model.states[ind_state].distribution\n",
    "            \n",
    "            if isinstance(distribution, ConditionalProbabilityTable):\n",
    "                idxs = distribution.column_idxs\n",
    "                table = self.prob_table[ind_state] # dict\n",
    "                \n",
    "                # calculate the new parameter for this key\n",
    "                for key in table.keys():\n",
    "                    num = 0\n",
    "                    denom = 0\n",
    "                    \n",
    "                    # marginal counts\n",
    "                    for mb_key in mb_keys:\n",
    "                        # marginal counts of node + parents\n",
    "                        if tuple([mb_key[ind[x]] for x in idxs]) == key:\n",
    "                            num += ct.table[mb_key[1:]]*expected_counts.counts[mb_key] \n",
    "                            \n",
    "                        # marginal counts of parents\n",
    "                        if tuple([mb_key[ind[x]] for x in idxs[:-1]]) == key[:-1]:\n",
    "                            denom += ct.table[mb_key[1:]]*expected_counts.counts[mb_key]\n",
    "                            \n",
    "                    try:\n",
    "                        prob = num/denom\n",
    "                    except ZeroDivisionError:\n",
    "                        prob = 0\n",
    "                        \n",
    "                    # update the parameter\n",
    "                    table[key] = prob\n",
    "                    \n",
    "            else: # DiscreteProb\n",
    "                table = self.prob_table[ind_state] # dict \n",
    "                \n",
    "                # calculate the new parameter for this key\n",
    "                for key in table.keys():\n",
    "                    prob = 0\n",
    "                    for mb_key in mb_keys:\n",
    "                        if mb_key[ind[ind_state]] == key:\n",
    "                            prob += ct.table[mb_key[1:]]*expected_counts.counts[mb_key]\n",
    "                    \n",
    "                    # update the parameter\n",
    "                    table[key] = prob/ct.size\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ExpectedCounts():\n",
    "    \"\"\"Calculate the expected counts using the model parameters\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    model : a BayesianNetwork object\n",
    "    \n",
    "    mb : a MarkovBlanket object\n",
    "    \n",
    "    Attributes\n",
    "    ----------\n",
    "    counts : dict\n",
    "        a dict of expected counts for nodes in the Markov blanket\n",
    "    \"\"\"\n",
    "    \n",
    "    def __init__(self, model, mb):\n",
    "        self.counts = {}\n",
    "        \n",
    "        self.populate(model, mb)\n",
    "        \n",
    "    def populate(self, model, mb):\n",
    "        #create combinations of keys\n",
    "        keys_list = [model.states[mb.hidden].distribution.keys()]\n",
    "        for ind in mb.parents + mb.children + mb.coparents:\n",
    "            keys_list.append(model.states[ind].distribution.keys())\n",
    "        \n",
    "        self.counts = {p:0 for p in itertools.product(*keys_list)}\n",
    "        \n",
    "    def update(self, model, mb):\n",
    "        ind = {x : i for i, x in enumerate([mb.hidden] + mb.parents + mb.children + mb.coparents)}\n",
    "    \n",
    "        marginal_prob = {}\n",
    "    \n",
    "        # calculate joint probability and marginal probability\n",
    "        for i, key in enumerate(self.counts.keys()):\n",
    "            prob = 1\n",
    "        \n",
    "            for j, ind_state in enumerate([mb.hidden] + mb.children):\n",
    "                distribution = model.states[ind_state].distribution\n",
    "            \n",
    "                if isinstance(distribution, ConditionalProbabilityTable):\n",
    "                    idxs = distribution.column_idxs\n",
    "                    state_key = tuple([key[ind[x]] for x in idxs])\n",
    "                else:\n",
    "                    state_key = key[ind[ind_state]]\n",
    "                \n",
    "                prob = prob*mb.prob_table[ind_state][state_key]         \n",
    "                self.counts[key] = prob\n",
    "            try:\n",
    "                marginal_prob[key[1:]] += prob\n",
    "            except KeyError:\n",
    "                marginal_prob[key[1:]] = prob\n",
    "                 \n",
    "        # divide the joint prob by the marginal prob to get the conditional\n",
    "        for i, key in enumerate(self.counts.keys()):\n",
    "            try:\n",
    "                self.counts[key] = self.counts[key]/marginal_prob[key[1:]]\n",
    "            except ZeroDivisionError:\n",
    "                self.counts[key] = 0\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class CountTable():\n",
    "    \"\"\"Counting the data\"\"\"\n",
    "    \n",
    "    def __init__(self, model, mb, items):\n",
    "        \"\"\"\n",
    "        Parameters\n",
    "        ----------\n",
    "        model : BayesianNetwork object\n",
    "        \n",
    "        mb : MarkovBlanket object\n",
    "        \n",
    "        items : ndarray\n",
    "            columns are data for parents, children, coparents\n",
    "        \n",
    "        \"\"\"\n",
    "        self.table ={}\n",
    "        self.ind = {}\n",
    "        self.size = items.shape[0]\n",
    "        \n",
    "        self.populate(model, mb, items)\n",
    "        \n",
    "    def populate(self, model, mb, items):\n",
    "        keys_list = []\n",
    "        for ind in mb.parents + mb.children + mb.coparents:\n",
    "            keys_list.append(model.states[ind].distribution.keys())\n",
    "        \n",
    "        # init\n",
    "        self.table = {p:0 for p in itertools.product(*keys_list)}\n",
    "        self.ind = {p:[] for p in itertools.product(*keys_list)}\n",
    "        \n",
    "        # count\n",
    "        for i, row in enumerate(items):\n",
    "            try:\n",
    "                self.table[tuple(row)] += 1\n",
    "                self.ind[tuple(row)].append(i)\n",
    "            except KeyError:\n",
    "                print 'Items in row', i, 'does not match the set of keys.'\n",
    "                raise KeyError\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def em_bayesnet(model, data, ind_h, max_iter = 50, criteria = 0.005, verbose=False):\n",
    "    \"\"\"Returns the data array with the hidden node filled in.\n",
    "    (model is not modified.)\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    model : a BayesianNetwork object\n",
    "        an already baked BayesianNetwork object with initialized parameters\n",
    "        \n",
    "    data : an ndarray\n",
    "        each column is the data for the node in the same order as the nodes in the model\n",
    "        the hidden node should be a column of NaNs\n",
    "        \n",
    "    ind_h : int\n",
    "        index of the hidden node\n",
    "        \n",
    "    max_iter : int\n",
    "        maximum number of iterations\n",
    "        \n",
    "    criteria : float between 0 and 1\n",
    "        the change in probability in consecutive iterations, below this value counts as convergence \n",
    "        \n",
    "    verbose : boolean\n",
    "        if True then the function prints an update with each iteration\n",
    "        \n",
    "    Returns\n",
    "    -------\n",
    "    data : an ndarray\n",
    "        the same data arary with the hidden node column filled in\n",
    "    \"\"\"\n",
    "    \n",
    "    # create the Markov blanket object for the hidden node\n",
    "    mb = MarkovBlanket(ind_h)\n",
    "    mb.populate(model)\n",
    "    mb.calculate_prob(model)\n",
    "    \n",
    "    # create the count table from data\n",
    "    items = data[:, mb.parents + mb.children + mb.coparents]\n",
    "    ct = CountTable(model, mb, items)\n",
    "    \n",
    "    # create expected counts\n",
    "    expected_counts = ExpectedCounts(model, mb)\n",
    "    expected_counts.update(model, mb)\n",
    "    \n",
    "    # ---- iterate over the E-M steps\n",
    "    i = 0\n",
    "    previous_params = np.array(mb.prob_table[mb.hidden].values())\n",
    "    convergence = False\n",
    "    \n",
    "    while (not convergence) and (i < max_iter):\n",
    "        mb.update_prob(model, expected_counts, ct)\n",
    "        expected_counts.update(model, mb)\n",
    "        if verbose : print 'Iteration',i,mb.prob_table,'\\n'\n",
    "        \n",
    "        # convergence criteria\n",
    "        hidden_params = np.array(mb.prob_table[mb.hidden].values())\n",
    "        change = abs(hidden_params - previous_params)\n",
    "        convergence = max(change) < criteria\n",
    "        previous_params = np.array(mb.prob_table[mb.hidden].values())\n",
    "        \n",
    "        i += 1\n",
    "        \n",
    "    if i == max_iter:\n",
    "        print 'Maximum iterations reached.'\n",
    "    \n",
    "    # ---- fill in the hidden node data by sampling the distribution\n",
    "    labels = {}\n",
    "    for key, prob in expected_counts.counts.items():\n",
    "        try:\n",
    "            labels[key[1:]].append((key[0], prob))\n",
    "        except:\n",
    "            labels[key[1:]] = [(key[0], prob)]\n",
    "            \n",
    "    new_data = data.copy()       \n",
    "    for key, counts in ct.table.items():\n",
    "        label, prob = zip(*labels[key])\n",
    "        prob = tuple(round(p,5) for p in prob)\n",
    "        if not all(p == 0 for p in prob):\n",
    "            samples = np.random.choice(label, size=counts, p=prob)\n",
    "            new_data[ct.ind[key], ind_h] = samples\n",
    "        \n",
    "    return new_data\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "a=np.array([[1,2,3],[4,5,6],[7,8,9]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 1, 33,  3],\n",
       "       [ 4, 44,  6],\n",
       "       [ 7, 55,  9]])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "a[:,1] = [33,44,55]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "b = a.copy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[99, 99, 99],\n",
       "       [ 4,  5,  6],\n",
       "       [ 7,  8,  9]])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b[0] = 99\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "a[[1,2]][:,1] = [99,99]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "a[[0,1],1] = [98,99]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([['0', '1', '0'],\n",
       "       ['0', '1', '1'],\n",
       "       ['1', '1', '0'],\n",
       "       ['1', '1', '1'],\n",
       "       ['1', '1', '0'],\n",
       "       ['0', '1', '1'],\n",
       "       ['1', '1', '1'],\n",
       "       ['1', '1', '0']], \n",
       "      dtype='|S3')"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "em_bayesnet(test_model2, data2, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[  0.,  nan,   0.],\n",
       "       [  0.,  nan,   1.],\n",
       "       [  1.,  nan,   0.],\n",
       "       [  1.,  nan,   1.],\n",
       "       [  1.,  nan,   0.],\n",
       "       [  0.,  nan,   1.],\n",
       "       [  1.,  nan,   1.],\n",
       "       [  1.,  nan,   0.]])"
      ]
     },
     "execution_count": 78,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "mb2 = MarkovBlanket(1)\n",
    "mb2.populate(test_model2)\n",
    "mb2.calculate_prob(test_model2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ct = CountTable(test_model2, mb2, data2[:,[0,2]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('0', '0'): 1, ('0', '1'): 2, ('1', '0'): 3, ('1', '1'): 2}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ct.table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "expected_counts = ExpectedCounts(test_model2, mb2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "expected_counts.update(test_model2, mb2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('0', '0', '0'): 0.4444444444444444,\n",
       " ('0', '0', '1'): 0.3478260869565218,\n",
       " ('0', '1', '0'): 0.33962264150943394,\n",
       " ('0', '1', '1'): 0.2553191489361702,\n",
       " ('1', '0', '0'): 0.5555555555555555,\n",
       " ('1', '0', '1'): 0.6521739130434782,\n",
       " ('1', '1', '0'): 0.660377358490566,\n",
       " ('1', '1', '1'): 0.7446808510638298}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expected_counts.counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{1: {('0', '0'): 0.3800322061191627,\n",
       "  ('0', '1'): 0.6199677938808373,\n",
       "  ('1', '0'): 0.30590124448012845,\n",
       "  ('1', '1'): 0.6940987555198715},\n",
       " 2: {('0', '0'): 0.5481386019791565,\n",
       "  ('0', '1'): 0.4518613980208435,\n",
       "  ('1', '0'): 0.4758909243055429,\n",
       "  ('1', '1'): 0.5241090756944571}}"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mb2.prob_table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('f', 'long'): 0.5, ('f', 'sh'): 0.5, ('m', 'long'): 0.5, ('m', 'sh'): 0.5}"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mb.prob_table[3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{('f', 'short', 'high', 'long'): 0.25,\n",
       " ('f', 'short', 'high', 'sh'): 0.25,\n",
       " ('f', 'short', 'low', 'long'): 0.25,\n",
       " ('f', 'short', 'low', 'sh'): 0.25,\n",
       " ('f', 'tall', 'high', 'long'): 0.25,\n",
       " ('f', 'tall', 'high', 'sh'): 0.25,\n",
       " ('f', 'tall', 'low', 'long'): 0.25,\n",
       " ('f', 'tall', 'low', 'sh'): 0.25,\n",
       " ('m', 'short', 'high', 'long'): 0.25,\n",
       " ('m', 'short', 'high', 'sh'): 0.25,\n",
       " ('m', 'short', 'low', 'long'): 0.25,\n",
       " ('m', 'short', 'low', 'sh'): 0.25,\n",
       " ('m', 'tall', 'high', 'long'): 0.25,\n",
       " ('m', 'tall', 'high', 'sh'): 0.25,\n",
       " ('m', 'tall', 'low', 'long'): 0.25,\n",
       " ('m', 'tall', 'low', 'sh'): 0.25}"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "expected_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[['short', 'low', 'm', 0.5],\n",
       " ['short', 'low', 'f', 0.5],\n",
       " ['short', 'high', 'm', 0.5],\n",
       " ['short', 'high', 'f', 0.5],\n",
       " ['tall', 'low', 'm', 0.5],\n",
       " ['tall', 'low', 'f', 0.5],\n",
       " ['tall', 'high', 'm', 0.5],\n",
       " ['tall', 'high', 'f', 0.5]]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.states[2].distribution.parameters[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Funtional tests\n",
    "Requirements:\n",
    "1. data with NaN outside of the hidden node should raise exception\n",
    "2. correct output for input data of type str, float/int"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import unittest"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class TestDataType(unittest.TestCase):\n",
    "    def test_int_input(self):\n",
    "        data = np.array([[0,np.nan,0,0],\n",
    "               [0,np.nan,0,0],\n",
    "               [0,np.nan,0,1],\n",
    "               [0,np.nan,1,1],\n",
    "               [1,np.nan,0,0],\n",
    "               [1,np.nan,1,1],\n",
    "               [1,np.nan,1,1],\n",
    "               [1,np.nan,1,1]])\n",
    "        cloudy = DiscreteDistribution({0:0.4,1:0.6})\n",
    "        sprinkler = ConditionalProbabilityTable([\n",
    "            [1, 0, 0.7],\n",
    "            [1, 1, 0.23],\n",
    "            [0, 0, 0.4],\n",
    "            [0, 1, 0.6]], [cloudy] )\n",
    "        rain = ConditionalProbabilityTable([\n",
    "            [1, 0, 0.2],\n",
    "            [1, 1, 0.8],\n",
    "            [0, 0, 0.7],\n",
    "            [0, 1, 0.3]], [cloudy] )\n",
    "        wetgrass = ConditionalProbabilityTable([\n",
    "            [0, 0, 0, 0.8],\n",
    "            [0, 0, 1, 0.2],\n",
    "            [0, 1, 0, 0.1],\n",
    "            [0, 1, 1, 0.9],\n",
    "            [1, 0, 0, 0.1],\n",
    "            [1, 0, 1, 0.9],\n",
    "            [1, 1, 0, 0.1],\n",
    "            [1, 1, 1, 0.9]], [sprinkler, rain] )\n",
    "        \n",
    "        s_c = State(cloudy, 'cloudy')\n",
    "        s_s = State(sprinkler, 'sprinkler')\n",
    "        s_r = State(rain, 'rain')\n",
    "        s_w = State(wetgrass, 'wetgrass')\n",
    "        \n",
    "        model = BayesianNetwork('wet')\n",
    "        model.add_states(s_c, s_s, s_r, s_w)\n",
    "        model.add_transition(s_c, s_s)\n",
    "        model.add_transition(s_c, s_r)\n",
    "        model.add_transition(s_s, s_w)\n",
    "        model.add_transition(s_r, s_w)\n",
    "        model.bake()\n",
    "        \n",
    "        new_data = em_bayesnet(model, data, 1)\n",
    "        \n",
    "        print new_data\n",
    "\n",
    "#if __name__ == \"__main__\":\n",
    "#    unittest.main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      ".."
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "em_bayesnet data =  [[  0.  nan]\n",
      " [  0.  nan]\n",
      " [  1.  nan]\n",
      " [  1.  nan]]\n",
      "MarkovBlanket prob_table =  {1: {(0, 1): 0.0, (1, 0): 0.0, (0, 0): 1.0, (1, 1): 1.0}}\n",
      "Iteration 0 {1: {(0, 1): 0.0, (1, 0): 0.0, (0, 0): 1.0, (1, 1): 1.0}}\n",
      "   [ 0.  0.  0.  0.]\n",
      "Labels =  {(0,): [(1, 0.0), (0, 1.0)], (1,): [(0, 0.0), (1, 1.0)]} \n",
      "\n",
      "Key  (0,) (1, 0) (0.0, 1.0) [0 0]\n",
      "Key  (1,) (0, 1) (0.0, 1.0) [1 1]\n",
      "[[ 0.  0.]\n",
      " [ 0.  0.]\n",
      " [ 1.  1.]\n",
      " [ 1.  1.]] [[0 0]\n",
      " [0 0]\n",
      " [1 1]\n",
      " [1 1]]\n",
      "em_bayesnet data =  [['0' 'nan']\n",
      " ['0' 'nan']\n",
      " ['1' 'nan']\n",
      " ['1' 'nan']]\n",
      "MarkovBlanket prob_table =  {1: {('1', '0'): 0.0, ('0', '1'): 0.0, ('1', '1'): 1.0, ('0', '0'): 1.0}}\n",
      "Iteration 0 {1: {('1', '0'): 0.0, ('0', '1'): 0.0, ('1', '1'): 1.0, ('0', '0'): 1.0}}\n",
      "   [ 0.  0.  0.  0.]\n",
      "Labels =  {('1',): [('0', 0.0), ('1', 1.0)], ('0',): [('1', 0.0), ('0', 1.0)]} \n",
      "\n",
      "Key  ('1',) ('0', '1') (0.0, 1.0) ['1' '1']\n",
      "Key  ('0',) ('1', '0') (0.0, 1.0) ['0' '0']\n",
      "[['0' '0']\n",
      " ['0' '0']\n",
      " ['1' '1']\n",
      " ['1' '1']] [['0' '0']\n",
      " ['0' '0']\n",
      " ['1' '1']\n",
      " ['1' '1']]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "----------------------------------------------------------------------\n",
      "Ran 2 tests in 0.028s\n",
      "\n",
      "OK\n"
     ]
    }
   ],
   "source": [
    "class TestDataType(unittest.TestCase):\n",
    "    \n",
    "    def test_int_input(self):\n",
    "        data = np.array([[0,np.nan],\n",
    "               [0,np.nan],\n",
    "               [1,np.nan],\n",
    "               [1,np.nan]])\n",
    "        A = DiscreteDistribution({0:0.4, 1:0.6})\n",
    "        B = ConditionalProbabilityTable([\n",
    "            [0, 0, 1.0],\n",
    "            [0, 1, 0.0],\n",
    "            [1, 0, 0.0],\n",
    "            [1, 1, 1.0]], [A] )\n",
    "        \n",
    "        s_A = State(A, 'A')\n",
    "        s_B = State(B, 'B')\n",
    "        \n",
    "        model = BayesianNetwork('copy')\n",
    "        model.add_states(s_A, s_B)\n",
    "        model.add_transition(s_A, s_B)\n",
    "        model.bake()\n",
    "        \n",
    "        new_data = em_bayesnet(model, data, 1)\n",
    "        \n",
    "        correct_output = np.array([\n",
    "            [0, 0],\n",
    "            [0, 0],\n",
    "            [1, 1],\n",
    "            [1, 1]])\n",
    "        \n",
    "        self.assertEqual(new_data.tolist(), correct_output.tolist())\n",
    "        \n",
    "    def test_str_input(self):\n",
    "        data = np.array([['0',np.nan],\n",
    "               ['0',np.nan],\n",
    "               ['1',np.nan],\n",
    "               ['1',np.nan]])\n",
    "        A = DiscreteDistribution({'0':0.4, '1':0.6})\n",
    "        B = ConditionalProbabilityTable([\n",
    "            ['0', '0', 1.0],\n",
    "            ['0', '1', 0.0],\n",
    "            ['1', '0', 0.0],\n",
    "            ['1', '1', 1.0]], [A] )\n",
    "        \n",
    "        s_A = State(A, 'A')\n",
    "        s_B = State(B, 'B')\n",
    "        \n",
    "        model = BayesianNetwork('copy')\n",
    "        model.add_states(s_A, s_B)\n",
    "        model.add_transition(s_A, s_B)\n",
    "        model.bake()\n",
    "        \n",
    "        new_data = em_bayesnet(model, data, 1)\n",
    "        \n",
    "        correct_output = np.array([\n",
    "            ['0', '0'],\n",
    "            ['0', '0'],\n",
    "            ['1', '1'],\n",
    "            ['1', '1']])\n",
    "        \n",
    "        self.assertEqual(new_data.tolist(), correct_output.tolist())\n",
    "        \n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    #unittest.main()\n",
    "    unittest.main(argv=['first-arg-is-ignored'], exit=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "test = RainTest()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "em_bayesnet data =  [[  0.  nan   0.   0.]\n",
      " [  0.  nan   0.   0.]\n",
      " [  0.  nan   0.   1.]\n",
      " [  0.  nan   1.   1.]\n",
      " [  1.  nan   0.   0.]\n",
      " [  1.  nan   1.   1.]\n",
      " [  1.  nan   1.   1.]\n",
      " [  1.  nan   1.   1.]]\n",
      "MarkovBlanket prob_table =  {1: {(0, 1): 0.6, (1, 0): 0.7, (0, 0): 0.4, (1, 1): 0.23}, 3: {(0, 1, 1): 0.9, (1, 1, 0): 0.1, (1, 0, 0): 0.1, (0, 0, 1): 0.2, (1, 0, 1): 0.9, (0, 0, 0): 0.8, (0, 1, 0): 0.1, (1, 1, 1): 0.9}}\n",
      "CountTable table =  {(0, 1, 1): 0, (1, 1, 0): 0, (1, 0, 0): 0, (0, 0, 1): 0, (1, 0, 1): 0, (0, 0, 0): 0, (0, 1, 0): 0, (1, 1, 1): 0}\n",
      "CountTable row : (0.0, 0.0, 0.0)\n",
      "CountTable row : (0.0, 0.0, 0.0)\n",
      "CountTable row : (0.0, 1.0, 0.0)\n",
      "CountTable row : (0.0, 1.0, 1.0)\n",
      "CountTable row : (1.0, 0.0, 0.0)\n",
      "CountTable row : (1.0, 1.0, 1.0)\n",
      "CountTable row : (1.0, 1.0, 1.0)\n",
      "CountTable row : (1.0, 1.0, 1.0)\n",
      "[ 0.1533107   0.10465335  0.1533107   0.03465335]\n",
      "[ 0.00528861  0.03166255  0.00528861  0.03166255]\n",
      "[ 0.01270834  0.02518209  0.01270834  0.02518209]\n",
      "[ 0.01268642  0.02139549  0.01268642  0.02139549]\n",
      "[ 0.01045257  0.01827844  0.01045257  0.01827844]\n",
      "[ 0.00831566  0.01555799  0.00831566  0.01555799]\n",
      "[ 0.00657567  0.01318875  0.00657567  0.01318875]\n",
      "[ 0.00521374  0.01114513  0.00521374  0.01114513]\n",
      "[ 0.00416004  0.0093973   0.00416004  0.0093973 ]\n",
      "[ 0.0033449   0.00791184  0.0033449   0.00791184]\n",
      "[ 0.00271042  0.00665483  0.00271042  0.00665483]\n",
      "[ 0.00221192  0.00559421  0.00221192  0.00559421]\n",
      "[ 0.00181614  0.00470091  0.00181614  0.00470091]\n",
      "Labels =  {(1, 1, 0): [(0, 0.013307664977571299), (1, 0.9866923350224287)], (0, 1, 1): [(1, 0.511596506089719), (0, 0.4884034939102811)], (1, 0, 0): [(0, 0.9911205386719177), (1, 0.008879461328082336)], (0, 0, 1): [(0, 0), (1, 0)], (1, 0, 1): [(1, 0), (0, 0)], (0, 0, 0): [(1, 0.27055494212692804), (0, 0.7294450578730719)], (0, 1, 0): [(1, 0.9996743313872091), (0, 0.00032566861279093365)], (1, 1, 1): [(1, 0.024677120336477166), (0, 0.9753228796635228)]} \n",
      "\n",
      "Key  (0, 1, 1) (1, 0) (0.5116, 0.4884) [0]\n",
      "Key  (1, 1, 0) (0, 1) (0.01331, 0.98669) []\n",
      "Key  (1, 0, 0) (0, 1) (0.99112, 0.00888) [0]\n",
      "Key  (0, 0, 0) (1, 0) (0.27055, 0.72945) [0 1]\n",
      "Key  (0, 1, 0) (1, 0) (0.99967, 0.00033) [1]\n",
      "Key  (1, 1, 1) (1, 0) (0.02468, 0.97532) [0 0 0]\n",
      "[[ 0.  0.  0.  0.]\n",
      " [ 0.  1.  0.  0.]\n",
      " [ 0.  1.  0.  1.]\n",
      " [ 0.  0.  1.  1.]\n",
      " [ 1.  0.  0.  0.]\n",
      " [ 1.  0.  1.  1.]\n",
      " [ 1.  0.  1.  1.]\n",
      " [ 1.  0.  1.  1.]]\n"
     ]
    }
   ],
   "source": [
    "test.test_sprinkler()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.random.choice([1,0],size=10,p=(0.024,0.976))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<generator object <genexpr> at 0x119498190>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x for x in [1,2,3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "em_bayesnet data =  [['0' 'nan' '0']\n",
      " ['0' 'nan' '1']\n",
      " ['1' 'nan' '0']\n",
      " ['1' 'nan' '1']\n",
      " ['1' 'nan' '0']\n",
      " ['0' 'nan' '1']\n",
      " ['1' 'nan' '1']\n",
      " ['1' 'nan' '0']]\n",
      "MarkovBlanket prob_table =  {1: {('1', '0'): 0.3, ('0', '1'): 0.6, ('1', '1'): 0.7, ('0', '0'): 0.4}, 2: {('1', '0'): 0.5, ('0', '1'): 0.4, ('1', '1'): 0.5, ('0', '0'): 0.6}}\n",
      "CountTable table =  {('1', '0'): 0, ('0', '1'): 0, ('1', '1'): 0, ('0', '0'): 0}\n",
      "CountTable row : ('0', '0')\n",
      "CountTable row : ('0', '1')\n",
      "CountTable row : ('1', '0')\n",
      "CountTable row : ('1', '1')\n",
      "CountTable row : ('1', '0')\n",
      "CountTable row : ('0', '1')\n",
      "CountTable row : ('1', '1')\n",
      "CountTable row : ('1', '0')\n",
      "[ 0.00590124  0.01996779  0.00590124  0.01996779]\n",
      "[ 0.00625935  0.01157651  0.00625935  0.01157651]\n",
      "[ 0.00443835  0.00801674  0.00443835  0.00801674]\n",
      "[ 0.00306137  0.00543619  0.00306137  0.00543619]\n",
      "[ 0.00198772  0.00348802  0.00198772  0.00348802]\n",
      "Labels =  {('1', '0'): [('1', 0.6727601222522195), ('0', 0.3272398777477805)], ('0', '1'): [('0', 0.34567161839155214), ('1', 0.6543283816084479)], ('1', '1'): [('1', 0.6839405982978158), ('0', 0.31605940170218433)], ('0', '0'): [('1', 0.6426477043479893), ('0', 0.3573522956520107)]} \n",
      "\n",
      "Key  ('1', '0') ('1', '0') (0.67276, 0.32724) ['1' '1' '0']\n",
      "Key  ('0', '1') ('0', '1') (0.34567, 0.65433) ['1' '0']\n",
      "Key  ('1', '1') ('1', '0') (0.68394, 0.31606) ['0' '1']\n",
      "Key  ('0', '0') ('1', '0') (0.64265, 0.35735) ['1']\n"
     ]
    }
   ],
   "source": [
    "new_data = em_bayesnet(test_model2, data2, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dtype('S3')"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_data[:,1].dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dtype('S3')"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_data.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dtype('int64')"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array([1,0,1]).dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "data type not understood",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-32-d53d7f9cba36>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0missubdtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minteger\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m/Users/pear/anaconda/envs/env_gdal/lib/python2.7/site-packages/numpy/core/numerictypes.pyc\u001b[0m in \u001b[0;36missubdtype\u001b[0;34m(arg1, arg2)\u001b[0m\n\u001b[1;32m    753\u001b[0m     \"\"\"\n\u001b[1;32m    754\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0missubclass_\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgeneric\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 755\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0missubclass\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marg2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    756\u001b[0m     \u001b[0mmro\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtype\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmro\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    757\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mlen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmro\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m>\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mTypeError\u001b[0m: data type not understood"
     ]
    }
   ],
   "source": [
    "np.issubdtype(np.array([1]),np.integer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'bBhHiIlLqQpP'"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.typecodes['AllInteger']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'i'"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array([1,0,1]).dtype.kind"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "data = np.array([[np.nan, 'yellow', 'sweet', 'long'],\n",
    "                [np.nan, 'green', 'sour', 'round'],\n",
    "                [np.nan, 'green', 'sour', 'round'],\n",
    "                [np.nan, 'yellow', 'sweet', 'long'],\n",
    "                [np.nan, 'yellow', 'sweet', 'long'],\n",
    "                [np.nan, 'green', 'sour', 'round'],\n",
    "                [np.nan, 'green', 'sweet', 'long'],\n",
    "                [np.nan, 'green', 'sweet', 'round']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "Fruit = DiscreteDistribution({'banana':0.4, 'apple':0.6})\n",
    "Color = ConditionalProbabilityTable([['banana', 'yellow', 0.6],\n",
    "                                 ['banana', 'green', 0.4],\n",
    "                                 ['apple', 'yellow', 0.6],\n",
    "                                 ['apple', 'green', 0.4]], [Fruit] ) \n",
    "Taste = ConditionalProbabilityTable([['banana', 'sweet', 0.6],\n",
    "                                    ['banana', 'sour', 0.4],\n",
    "                                    ['apple', 'sweet', 0.4],\n",
    "                                    ['apple', 'sour', 0.6]], [Fruit])\n",
    "Shape = ConditionalProbabilityTable([['banana', 'long', 0.6],\n",
    "                                    ['banana', 'round', 0.4],\n",
    "                                    ['apple', 'long', 0.4],\n",
    "                                    ['apple', 'round', 0.6]], [Fruit])\n",
    "s_fruit = State(Fruit, 'fruit')\n",
    "s_color = State(Color, 'color')\n",
    "s_taste = State(Taste, 'taste')\n",
    "s_shape = State(Shape, 'shape')\n",
    "model = BayesianNetwork('fruit')\n",
    "model.add_states(s_fruit, s_color, s_taste, s_shape)\n",
    "model.add_transition(s_fruit, s_color)\n",
    "model.add_transition(s_fruit, s_taste)\n",
    "model.add_transition(s_fruit, s_shape)\n",
    "model.bake()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "new_data = em_bayesnet(model, data, 0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([['banana', 'yellow', 'sweet', 'long'],\n",
       "       ['apple', 'green', 'sour', 'round'],\n",
       "       ['apple', 'green', 'sour', 'round'],\n",
       "       ['banana', 'yellow', 'sweet', 'long'],\n",
       "       ['banana', 'yellow', 'sweet', 'long'],\n",
       "       ['apple', 'green', 'sour', 'round'],\n",
       "       ['banana', 'green', 'sweet', 'long'],\n",
       "       ['apple', 'green', 'sweet', 'round']], \n",
       "      dtype='|S32')"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'apple': 0.6, 'banana': 0.4}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.states[0].distribution.parameters[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "new_model = model.fit(new_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### README"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Bayesian Network with Hidden Variables\n",
    "\n",
    "bayesnet_em predicts values for a hidden variable in a Bayesian network by implementing the expectation maximization algorithm. It works as an extension to the Beysian network implementation in Pomegranade.\n",
    "\n",
    "## Installation\n",
    "\n",
    "### Dependencies\n",
    "> - numpy \n",
    "> - pomegranade \n",
    "\n",
    "### Installing\n",
    "If you have Git installed:\n",
    "> `pip install git+https://github.com/nnvutisa/EM_BayesNet.git`\n",
    "\n",
    "## Usage\n",
    "\n",
    "A Bayesian network is a probabilistic graphical model that represents relationships between random variables as a direct acyclic graph. Each node in the network represents a random variable whearas each edge represents a conditional dependency. Bayesian networks provides an efficient way to construct a full joint probability distribution over the variables. The random varibles can either be observed variables or unobserved variables, in which case they are called hidden (or latent) variables. \n",
    "\n",
    "Pomegranade currently supports a discrete Baysian network. Each node represents a categorical variable, which means it can take on a discrete number of values. The model parameters can be learned from data. However (at least as of now), it does not support a network with hidden variables. The purpose of bayesnet_em is to work with the pomegranade Bayesian network model to predict values of the hidden variables.  \n",
    "\n",
    "bayesnet_em takes an already constructed and initialized `BayesianNetwork` object, a data array, and the index of the hidden node, and returns a complete data set.\n",
    "\n",
    "### Example\n",
    "Let's use a simple example. Suppose there is a bag of fruits that contains apples and bananas. The observed variables for each sample taken from the bag are color, taste, and shape, while the label of the type of fruits was not recorded. We would like to predict the type of fruits for each sample along with the full probability distribution. This can be represented by a Bayesian network as:\n",
    "\n",
    "This simple relationship describes a naive Bayes model. The full joint probability distribution is\n",
    "> P(F,C,T,S) = P(F)\\*P(C|F)\\*P(T|F)\\*P(S|F)\n",
    "\n",
    "Start with building a Bayesian network model using pomegranade.\n",
    "\n",
    "```\n",
    "import numpy as np\n",
    "from pomegranate import *\n",
    "```\n",
    "\n",
    "Here's our data. bayesnet_em suppost data types of int and string.\n",
    "\n",
    "```\n",
    "data = np.array([[np.nan, 'yellow', 'sweet', 'long'],\n",
    "                [np.nan, 'green', 'sour', 'round'],\n",
    "                [np.nan, 'green', 'sour', 'round'],\n",
    "                [np.nan, 'yellow', 'sweet', 'long'],\n",
    "                [np.nan, 'yellow', 'sweet', 'long'],\n",
    "                [np.nan, 'green', 'sour', 'round'],\n",
    "                [np.nan, 'green', 'sweet', 'long'],\n",
    "                [np.nan, 'green', 'sweet', 'round']])\n",
    "```\n",
    "\n",
    "The columns represent the nodes in a specified order (fruit, color, taste, shape). The order of the columns have to match the order of the nodes (states) when constructing the Bayesian network. The first column with the `nan` values is the hidden node. Next, create the distributions of all the nodes and initialize the probabilities to some non-uniform values. The first node is just P(F). The other three nodes are described by conditional probabilities.\n",
    "\n",
    "```\n",
    "Fruit = DiscreteDistribution({'banana':0.4, 'apple':0.6})\n",
    "Color = ConditionalProbabilityTable([['banana', 'yellow', 0.6],\n",
    "                                 ['banana', 'green', 0.4],\n",
    "                                 ['apple', 'yellow', 0.6],\n",
    "                                 ['apple', 'green', 0.4]], [Fruit] ) \n",
    "Taste = ConditionalProbabilityTable([['banana', 'sweet', 0.6],\n",
    "                                    ['banana', 'sour', 0.4],\n",
    "                                    ['apple', 'sweet', 0.4],\n",
    "                                    ['apple', 'sour', 0.6]], [Fruit])\n",
    "Shape = ConditionalProbabilityTable([['banana', 'long', 0.6],\n",
    "                                    ['banana', 'round', 0.4],\n",
    "                                    ['apple', 'long', 0.4],\n",
    "                                    ['apple', 'round', 0.6]], [Fruit])\n",
    "```\n",
    "Now, create the state (node) objects\n",
    "\n",
    "```\n",
    "s_fruit = State(Fruit, 'fruit')\n",
    "s_color = State(Color, 'color')\n",
    "s_taste = State(Taste, 'taste')\n",
    "s_shape = State(Shape, 'shape')\n",
    "```\n",
    "\n",
    "and the `BayesianNetwork` object.\n",
    "\n",
    "```\n",
    "model = BayesianNetwork('fruit')\n",
    "```\n",
    "\n",
    "Add states and edges to the network.\n",
    "\n",
    "```\n",
    "model.add_states(s_fruit, s_color, s_taste, s_shape)\n",
    "model.add_transition(s_fruit, s_color)\n",
    "model.add_transition(s_fruit, s_taste)\n",
    "model.add_transition(s_fruit, s_shape)\n",
    "model.bake()\n",
    "```\n",
    "\n",
    "Now that we have the initialized model, we want to fill in the first column.\n",
    "\n",
    "```\n",
    "from bayesnet_em import *\n",
    "```\n",
    "\n",
    "call the `em_bayesnet` function. Our hidden node index is 0 since it is the first node.\n",
    "\n",
    "```\n",
    "hidden_node_index = 0\n",
    "new_data = em_bayesnet(model, data, hidden_node_index)\n",
    "```\n",
    "\n",
    "The returned array is the filled in data. Note that the function does not modify the model or the input data; it only returns a new data array.\n",
    "\n",
    "```\n",
    ">>> new_data\n",
    "array([['banana', 'yellow', 'sweet', 'long'],\n",
    "       ['apple', 'green', 'sour', 'round'],\n",
    "       ['apple', 'green', 'sour', 'round'],\n",
    "       ['banana', 'yellow', 'sweet', 'long'],\n",
    "       ['banana', 'yellow', 'sweet', 'long'],\n",
    "       ['apple', 'green', 'sour', 'round'],\n",
    "       ['banana', 'green', 'sweet', 'long'],\n",
    "       ['apple', 'green', 'sweet', 'round']], \n",
    "      dtype='|S32')\n",
    "```\n",
    "\n",
    "Of course the function does not actually know the meaning of the labels. In real uses, it would be more appropriate to use integer labels or something like 'group1', 'group2', etc. For the sake of this example, I just assigned the correct label names.\n",
    "\n",
    "If you want to furthur use the Bayesian network model (to see the distributions or make predictions for example), you can fit the model to the complete data set. \n",
    "\n",
    "```\n",
    "new_model = model.fit(new_data)\n",
    "```\n",
    "\n",
    "The model can be used to answer questions like, \"what is the probability that a banana is yellow?\"\n",
    "\n",
    "```\n",
    ">>> new_model.predict_proba({'fruit':'banana'})[1].parameters\n",
    "[{'green': 0.25000000000000017, 'yellow': 0.74999999999999989}]\n",
    "```\n",
    "\n",
    "What is the probability that a sample would be a sweet, round, and green apple?\n",
    "\n",
    "```\n",
    ">>> new_model.probability(['apple', 'green', 'sweet', 'round'])\n",
    "0.12500000000000003\n",
    "```\n",
    "\n",
    "## License\n",
    "\n",
    "This project is licensed under the MIT License - see the [LICENSE.md](LICENSE.md) file for details\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"banana\" :0.5,\n",
       "            \"apple\" :0.5\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "}"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_model.marginal()[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['apple',\n",
       "       {\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"green\" :1.0,\n",
       "            \"yellow\" :0.0\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "},\n",
       "       {\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"sweet\" :0.25000000000000017,\n",
       "            \"sour\" :0.7499999999999999\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "},\n",
       "       {\n",
       "    \"frozen\" :false,\n",
       "    \"class\" :\"Distribution\",\n",
       "    \"parameters\" :[\n",
       "        {\n",
       "            \"round\" :1.0,\n",
       "            \"long\" :0.0\n",
       "        }\n",
       "    ],\n",
       "    \"name\" :\"DiscreteDistribution\"\n",
       "}], dtype=object)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_model.predict_proba({'fruit':'apple'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.12500000000000003"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_model.probability(['apple', 'green', 'sweet', 'round'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[{'green': 0.25000000000000017, 'yellow': 0.74999999999999989}]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_model.predict_proba({'fruit':'banana'})[1].parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXcAAAD8CAYAAACMwORRAAAABHNCSVQICAgIfAhkiAAAAAlwSFlz\nAAALEgAACxIB0t1+/AAAD1dJREFUeJzt3X+s3Xddx/Hna52VyIYad1HTH9xGO5MGyBbuigbkh2ym\ny2Jr5FcbTVgEGhLKiAixBG20i4ojuhhTEyrMEJPRzYFykasF+eHAMHLv2By2o3JThr1pApcxMIhS\nim//uKfkcLk/vufec3fvPns+kpud7/d87ve+m5099+333u+5qSokSW25bL0HkCQNn3GXpAYZd0lq\nkHGXpAYZd0lqkHGXpAYZd0lqkHGXpAYZd0lq0OXr9YWvuuqqGh0dXa8vL0lPSPfff/9Xq2pkuXXr\nFvfR0VGmpqbW68tL0hNSki91WedlGUlqkHGXpAYZd0lqkHGXpAYZd0lqkHGXpAYZd0lqkHGXpAYZ\nd0lq0LrdoSq1bPTwh9Z7BG1gj7z9pjX/Gp65S1KDjLskNci4S1KDjLskNci4S1KDjLskNci4S1KD\njLskNci4S1KDOsU9yZ4kZ5JMJzm8wPO3J3mw9/EfSb4+/FElSV0t+/YDSTYBx4AbgBlgMsl4VZ2+\ntKaqfqtv/RuAa9dgVklSR13O3HcD01V1tqouACeAfUusPwC8dxjDSZJWpkvctwDn+rZnevt+QJJn\nADuAj61+NEnSSnWJexbYV4us3Q/cU1XfXfBAycEkU0mmZmdnu84oSRpQl7jPANv6trcC5xdZu58l\nLslU1fGqGquqsZGRke5TSpIG0iXuk8DOJDuSbGYu4OPzFyX5OeDHgU8Pd0RJ0qCWjXtVXQQOASeB\nh4G7q+pUkqNJ9vYtPQCcqKrFLtlIkh4nnX4TU1VNABPz9h2Zt/37wxtLkrQa3qEqSQ0y7pLUIOMu\nSQ0y7pLUIOMuSQ0y7pLUIOMuSQ0y7pLUIOMuSQ0y7pLUIOMuSQ0y7pLUIOMuSQ0y7pLUIOMuSQ0y\n7pLUIOMuSQ0y7pLUoE5xT7InyZkk00kOL7LmFUlOJzmV5M7hjilJGsSyv0M1ySbgGHADMANMJhmv\nqtN9a3YCbwWeV1WPJXn6Wg0sSVpelzP33cB0VZ2tqgvACWDfvDWvBY5V1WMAVfWV4Y4pSRpEl7hv\nAc71bc/09vW7Grg6yb8muS/JnmENKEka3LKXZYAssK8WOM5O4EXAVuCTSZ5ZVV//vgMlB4GDANu3\nbx94WElSN13O3GeAbX3bW4HzC6z5QFV9p6q+CJxhLvbfp6qOV9VYVY2NjIysdGZJ0jK6xH0S2Jlk\nR5LNwH5gfN6avwdeDJDkKuYu05wd5qCSpO6WjXtVXQQOASeBh4G7q+pUkqNJ9vaWnQQeTXIa+Djw\nlqp6dK2GliQtrcs1d6pqApiYt+9I3+MC3tT7kCStM+9QlaQGGXdJapBxl6QGGXdJapBxl6QGGXdJ\napBxl6QGGXdJapBxl6QGGXdJapBxl6QGGXdJapBxl6QGGXdJapBxl6QGGXdJapBxl6QGGXdJapBx\nl6QGdYp7kj1JziSZTnJ4gedvTjKb5MHex2uGP6okqatlf0F2kk3AMeAGYAaYTDJeVafnLb2rqg6t\nwYySpAF1OXPfDUxX1dmqugCcAPat7ViSpNXoEvctwLm+7ZnevvlemuShJPck2TaU6SRJK9Il7llg\nX83b/iAwWlXPBv4ZeM+CB0oOJplKMjU7OzvYpJKkzrrEfQboPxPfCpzvX1BVj1bVt3ubfwU8Z6ED\nVdXxqhqrqrGRkZGVzCtJ6qBL3CeBnUl2JNkM7AfG+xck+em+zb3Aw8MbUZI0qGV/WqaqLiY5BJwE\nNgF3VNWpJEeBqaoaB25Jshe4CHwNuHkNZ5YkLWPZuANU1QQwMW/fkb7HbwXeOtzRJEkr5R2qktQg\n4y5JDTLuktQg4y5JDTLuktQg4y5JDTLuktQg4y5JDTLuktQg4y5JDTLuktQg4y5JDTLuktQg4y5J\nDTLuktQg4y5JDTLuktQg4y5JDTLuktSgTnFPsifJmSTTSQ4vse5lSSrJ2PBGlCQNatm4J9kEHANu\nBHYBB5LsWmDdlcAtwGeGPaQkaTBdztx3A9NVdbaqLgAngH0LrLsVuA343yHOJ0lagS5x3wKc69ue\n6e37niTXAtuq6h+WOlCSg0mmkkzNzs4OPKwkqZsucc8C++p7TyaXAbcDv73cgarqeFWNVdXYyMhI\n9yklSQPpEvcZYFvf9lbgfN/2lcAzgU8keQT4eWDcb6pK0vrpEvdJYGeSHUk2A/uB8UtPVtU3quqq\nqhqtqlHgPmBvVU2tycSSpGUtG/equggcAk4CDwN3V9WpJEeT7F3rASVJg7u8y6KqmgAm5u07ssja\nF61+LEnSaniHqiQ1yLhLUoOMuyQ1yLhLUoOMuyQ1yLhLUoOMuyQ1yLhLUoOMuyQ1yLhLUoOMuyQ1\nyLhLUoOMuyQ1yLhLUoOMuyQ1yLhLUoOMuyQ1yLhLUoM6xT3JniRnkkwnObzA869L8rkkDyb5VJJd\nwx9VktTVsnFPsgk4BtwI7AIOLBDvO6vqWVV1DXAb8GdDn1SS1FmXM/fdwHRVna2qC8AJYF//gqr6\nr77NpwI1vBElSYO6vMOaLcC5vu0Z4LnzFyV5PfAmYDPwS0OZTpK0Il3O3LPAvh84M6+qY1X1M8Dv\nAL+74IGSg0mmkkzNzs4ONqkkqbMucZ8BtvVtbwXOL7H+BPCrCz1RVceraqyqxkZGRrpPKUkaSJe4\nTwI7k+xIshnYD4z3L0iys2/zJuALwxtRkjSoZa+5V9XFJIeAk8Am4I6qOpXkKDBVVePAoSTXA98B\nHgNetZZDS5KW1uUbqlTVBDAxb9+RvsdvHPJckqRV8A5VSWqQcZekBhl3SWqQcZekBhl3SWqQcZek\nBhl3SWqQcZekBhl3SWpQpztUN5rRwx9a7xG0gT3y9pvWewRp3XnmLkkNMu6S1CDjLkkNMu6S1CDj\nLkkNMu6S1CDjLkkNMu6S1CDjLkkN6hT3JHuSnEkyneTwAs+/KcnpJA8l+WiSZwx/VElSV8vGPckm\n4BhwI7ALOJBk17xlDwBjVfVs4B7gtmEPKknqrsuZ+25guqrOVtUF4ASwr39BVX28qr7V27wP2Drc\nMSVJg+gS9y3Aub7tmd6+xbwa+MfVDCVJWp0u7wqZBfbVgguT3wDGgBcu8vxB4CDA9u3bO44oSRpU\nlzP3GWBb3/ZW4Pz8RUmuB94G7K2qby90oKo6XlVjVTU2MjKyknklSR10ifsksDPJjiSbgf3AeP+C\nJNcC72Qu7F8Z/piSpEEsG/equggcAk4CDwN3V9WpJEeT7O0tewdwBfC3SR5MMr7I4SRJj4NOv4mp\nqiaAiXn7jvQ9vn7Ic0mSVsE7VCWpQcZdkhpk3CWpQcZdkhpk3CWpQcZdkhpk3CWpQcZdkhpk3CWp\nQcZdkhpk3CWpQcZdkhpk3CWpQcZdkhpk3CWpQcZdkhpk3CWpQcZdkhpk3CWpQZ3inmRPkjNJppMc\nXuD5FyT5bJKLSV42/DElSYNYNu5JNgHHgBuBXcCBJLvmLftP4GbgzmEPKEka3OUd1uwGpqvqLECS\nE8A+4PSlBVX1SO+5/1uDGSVJA+pyWWYLcK5ve6a3b2BJDiaZSjI1Ozu7kkNIkjroEvcssK9W8sWq\n6nhVjVXV2MjIyEoOIUnqoEvcZ4BtfdtbgfNrM44kaRi6xH0S2JlkR5LNwH5gfG3HkiStxrJxr6qL\nwCHgJPAwcHdVnUpyNMlegCTXJZkBXg68M8mptRxakrS0Lj8tQ1VNABPz9h3pezzJ3OUaSdIG4B2q\nktQg4y5JDTLuktQg4y5JDTLuktQg4y5JDTLuktQg4y5JDTLuktQg4y5JDTLuktQg4y5JDTLuktQg\n4y5JDTLuktQg4y5JDTLuktQg4y5JDeoU9yR7kpxJMp3k8ALP/3CSu3rPfybJ6LAHlSR1t2zck2wC\njgE3AruAA0l2zVv2auCxqvpZ4HbgT4Y9qCSpuy5n7ruB6ao6W1UXgBPAvnlr9gHv6T2+B3hJkgxv\nTEnSILrEfQtwrm97prdvwTVVdRH4BvATwxhQkjS4yzusWegMvFawhiQHgYO9zW8mOdPh62t5VwFf\nXe8hNop4UXAj8jXaZ5Wv0Wd0WdQl7jPAtr7trcD5RdbMJLkc+FHga/MPVFXHgeNdBlN3Saaqamy9\n55AW42v08dflsswksDPJjiSbgf3A+Lw148Creo9fBnysqn7gzF2S9PhY9sy9qi4mOQScBDYBd1TV\nqSRHgamqGgfeDfxNkmnmztj3r+XQkqSlxRPsJ74kB3uXvKQNydfo48+4S1KDfPsBSWqQcZc0dEm+\n2fvnaJJ/X+95noyM+xNM70dNJWlJxn2DSfJ7ST6f5CNJ3pvkzUk+keSPkvwL8MYkI0nel2Sy9/G8\n3uc+NckdvX0PJNnX239zkvcn+ackX0hy27r+IfWEk+TWJG/s2/7DJLckeUvv9fZQkj9Y5hhPSfLX\nST7Xe32+uLd/Ismze48fSHKk72u+Zi3/XC0z7htIkjHgpcC1wK8B/Td9/FhVvbCq/hT4c+D2qrqu\nt/5dvTVvY+4eg+uAFwPvSPLU3nPXAK8EngW8Mkn/jWnSct5N716WJJcx9+POXwZ2Mvf+U9cAz0ny\ngiWO8XqAqnoWcAB4T5KnAPcCv5jkacBF4Hm99c8HPjn8P8qTg3/F31ieD3ygqv4HIMkH+567q+/x\n9cCuvvdme1qSK4FfBvYmeXNv/1OA7b3HH62qb/SOe5q5W5j73zNIWlRVPZLk0STXAj8JPABcx9xr\n7oHesiuYi/29ixzm+cBf9I73+SRfAq5mLuC3AF8EPgTckORHgNGq8i1KVsi4byxLvZPmf/c9vgz4\nhUv/E/jeJ8/V/qXz/4NI8lzg2327vov/7jW4dwE3Az8F3AG8BPjjqnpnx89f7PU9ydzfUs8CH2Hu\nfWheC9y/mmGf7Lwss7F8CviV3rXJK4CbFln3YeDQpY0k1/QengTecOntlntnWdKw/B2wh7kz9pO9\nj9/svVZJsiXJ05f4/HuBX++tvZq5v1We6b2V+DngFcB9zJ3JvxkvyayKcd9AqmqSuffp+Tfg/cAU\nc2+fPN8twFjvm1ingdf19t8K/BDwUO/Hz25d+6n1ZNGL8MeBu6vqu1X1YeBO4NNJPsfc73K4colD\n/CWwqbf2LuDmqrr0N8pPAl+uqm/1Hm/FuK+Kd6huMEmuqKpv9q453gscrKrPrvdcUu8bqZ8FXl5V\nX1jvebQ0z9w3nuNJHmTuP6L3GXZtBL1frTnN3DfmDfsTgGfuktQgz9wlqUHGXZIaZNwlqUHGXZIa\nZNwlqUHGXZIa9P+VZbePeGfllgAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<matplotlib.figure.Figure at 0x1147a1ed0>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.bar(range(2), list(a1.values()), align='center')\n",
    "plt.xticks(range(2), list(a1.keys()))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}