Repository: pgmpy/pgmpy
Issue: Missing Values Crash ML

I'm trying to get estimates for conditional probabilities for a completely observed graphical model with some missing data. I attempted to simply add the edges to the model and fit to my data in a pandas data frame, but it looks like the base estimator class attempts to coerce the columns to type int which crashes. Since in pandas integers with missing data (np.nan) are of type float and can't be cast to int I'm curious if this was not intended to work with missing data, but the MLE code seems to have checks for missing columns so I'm not sure if this handled missing data at some point or if I'm missing some way to handle the missing values differently.

Removing the cast to int in estimators/base.py allows the code to run but I'm not sure if this actually handling things properly in MLE or if there are other changes to make to get proper estimates.


Changed files:
pgmpy/tests/test_estimators/test_MaximumLikelihoodEstimator.py, pgmpy/estimators/MLE.py, pgmpy/estimators/base.py

diff --git a/pgmpy/estimators/MLE.py b/pgmpy/estimators/MLE.py
index 0b6b2929..be0e7c2c 100644
--- a/pgmpy/estimators/MLE.py
+++ b/pgmpy/estimators/MLE.py
@@ -1,3 +1,5 @@
+# coding:utf-8
+
 from pgmpy.estimators import BaseEstimator
 from pgmpy.factors import TabularCPD
 from pgmpy.models import BayesianModel
@@ -7,14 +9,19 @@ import pandas as pd
 
 class MaximumLikelihoodEstimator(BaseEstimator):
     """
-    Class used to compute parameters for a model using Maximum Likelihood Estimate.
+    Class used to compute parameters for a model using Maximum Likelihood Estimation.
 
     Parameters
     ----------
     model: A pgmpy.models.BayesianModel instance
 
     data: pandas DataFrame object
-        datafame object with column names same as the variable names of the network
+        DataFrame object with column names identical to the variable names of the network
+
+    state_names: dict (optional)
+        A dict indicating, for each variable, the discrete set of states
+        that the variable can take. If unspecified, the observed values
+        in the data set are taken to be the only possible states.
 
     Examples
     --------
@@ -22,20 +29,20 @@ class MaximumLikelihoodEstimator(BaseEstimator):
     >>> import pandas as pd
     >>> from pgmpy.models import BayesianModel
     >>> from pgmpy.estimators import MaximumLikelihoodEstimator
-    >>> values = pd.DataFrame(np.random.randint(low=0, high=2, size=(1000, 5)),
+    >>> data = pd.DataFrame(np.random.randint(low=0, high=2, size=(1000, 5)),
     ...                       columns=['A', 'B', 'C', 'D', 'E'])
     >>> model = BayesianModel([('A', 'B'), ('C', 'B'), ('C', 'D'), ('B', 'E')])
-    >>> estimator = MaximumLikelihoodEstimator(model, values)
+    >>> estimator = MaximumLikelihoodEstimator(model, data)
     """
-    def __init__(self, model, data):
+    def __init__(self, model, data, state_names=None):
         if not isinstance(model, BayesianModel):
             raise NotImplementedError("Maximum Likelihood Estimate is only implemented for BayesianModel")
 
-        super(MaximumLikelihoodEstimator, self).__init__(model, data)
+        super(MaximumLikelihoodEstimator, self).__init__(model, data, state_names)
 
     def get_parameters(self):
         """
-        Method used to get parameters.
+        Method to estimate the model parameters (CPDs).
 
         Returns
         -------
@@ -48,38 +55,79 @@ class MaximumLikelihoodEstimator(BaseEstimator):
         >>> import pandas as pd
         >>> from pgmpy.models import BayesianModel
         >>> from pgmpy.estimators import MaximumLikelihoodEstimator
-        >>> values = pd.DataFrame(np.random.randint(low=0, high=2, size=(1000, 5)),
-        ...                       columns=['A', 'B', 'C', 'D', 'E'])
-        >>> model = BayesianModel([('A', 'B'), ('C', 'B'), ('C', 'D'), ('B', 'E')])
+        >>> values = pd.DataFrame(np.random.randint(low=0, high=2, size=(1000, 4)),
+        ...                       columns=['A', 'B', 'C', 'D'])
+        >>> model = BayesianModel([('A', 'B'), ('C', 'B'), ('C', 'D'))
         >>> estimator = MaximumLikelihoodEstimator(model, values)
         >>> estimator.get_parameters()
+        [<TabularCPD representing P(C:2) at 0x7f7b534251d0>,
+        <TabularCPD representing P(B:2 | C:2, A:2) at 0x7f7b4dfd4da0>,
+        <TabularCPD representing P(A:2) at 0x7f7b4dfd4fd0>,
+        <TabularCPD representing P(D:2 | C:2) at 0x7f7b4df822b0>]
         """
         parameters = []
 
         for node in self.model.nodes():
-            parents = self.model.get_parents(node)
-            if not parents:
-                state_counts = self.data.ix[:, node].value_counts()
-                state_counts = state_counts.reindex(sorted(state_counts.index))
-                cpd = TabularCPD(node, self.node_card[node],
-                                 state_counts.values[:, np.newaxis])
-                cpd.normalize()
-                parameters.append(cpd)
-            else:
-                parent_card = np.array([self.node_card[parent] for parent in parents])
-                var_card = self.node_card[node]
-
-                values = self.data.groupby([node] + parents).size().unstack(parents).fillna(0)
-                if not len(values.columns) == np.prod(parent_card):
-                    # some columns are missing if for some states of the parents no data was observed.
-                    # reindex to add missing columns and fill in uniform (conditional) probabilities:
-                    full_index = pd.MultiIndex.from_product([range(card) for card in parent_card], names=parents)
-                    values = values.reindex(columns=full_index).fillna(1.0/var_card)
-
-                cpd = TabularCPD(node, var_card, np.array(values),
-                                 evidence=parents,
-                                 evidence_card=parent_card.astype('int'))
-                cpd.normalize()
-                parameters.append(cpd)
+            cpd = self._estimate_cpd(node)
+            parameters.append(cpd)
 
         return parameters
+
+    def _estimate_cpd(self, node):
+        """
+        Method to estimate the CPD for a given variable.
+
+        Parameters
+        ----------
+        node: int, string (any hashable python object)
+            The name of the variable for which the CPD is to be estimated.
+
+        Returns
+        -------
+        CPD: TabularCPD
+
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> from pgmpy.models import BayesianModel
+        >>> from pgmpy.estimators import MaximumLikelihoodEstimator
+        >>> data = pd.DataFrame(data={'A': [0, 0, 1], 'B': [0, 1, 0], 'C': [1, 1, 0]})
+        >>> model = BayesianModel([('A', 'C'), ('B', 'C')])
+        >>> cpd_A = MaximumLikelihoodEstimator(model, data)._estimate_cpd('A')
+        >>> print(cpd_A)
+        ╒══════╤══════════╕
+        │ A(0) │ 0.666667 │
+        ├──────┼──────────┤
+        │ A(1) │ 0.333333 │
+        ╘══════╧══════════╛
+        """
+
+        parents = sorted(self.model.get_parents(node))
+        node_cardinality = len(self.state_names[node])
+        parents_cardinalities = np.array([len(self.state_names[parent]) for parent in parents])
+
+        if not parents:
+            state_count_data = self.data.ix[:, node].value_counts()
+            state_counts = state_count_data.reindex(self.state_names[node]).fillna(0).values[:, np.newaxis]
+
+        else:
+            state_count_data = self.data.groupby([node] + parents).size()
+            state_counts = state_count_data.unstack(parents).reindex(self.state_names[node]).fillna(0)
+            if isinstance(state_counts.index, pd.MultiIndex):
+                state_counts = state_counts.sortlevel(axis=1)
+            else:
+                state_counts = state_counts.sort_index(axis=1)
+
+            # some columns might be missing if for some states of the parents no data was observed:
+            if not len(state_counts.columns) == np.prod(parents_cardinalities):
+                possible_parents_states = [self.state_names[parent] for parent in parents]
+                # reindex to add missing columns and fill in uniform (conditional) probabilities:
+                full_index = pd.MultiIndex.from_product(possible_parents_states, names=parents)
+                state_counts = state_counts.reindex(columns=full_index).fillna(1.0 / node_cardinality)
+
+        cpd = TabularCPD(node, node_cardinality, np.array(state_counts),
+                         evidence=parents,
+                         evidence_card=parents_cardinalities,
+                         state_names=self.state_names)
+        cpd.normalize()
+        return cpd
diff --git a/pgmpy/estimators/base.py b/pgmpy/estimators/base.py
index d74775e0..f67f66c8 100644
--- a/pgmpy/estimators/base.py
+++ b/pgmpy/estimators/base.py
@@ -4,8 +4,7 @@ import numpy as np
 
 class BaseEstimator(object):
     """
-    Base class for estimator class in pgmpy. Estimator class is used for parameter estimation as well
-    as structure estimation
+    Base class for parameter estimators in pgmpy.
 
     Parameters
     ----------
@@ -13,11 +12,28 @@ class BaseEstimator(object):
         model for which parameter estimation is to be done
 
     data: pandas DataFrame object
-        datafame object with column names same as the variable names of the network
+        datafame object with column names identical to the variable names of the model
+
+    state_names: dict (optional)
+        A dict indicating, for each variable, the discrete set of states (or values)
+        that the variable can take. If unspecified, the observed values in the data set
+        are taken to be the only possible states.
     """
-    def __init__(self, model, data):
+    def __init__(self, model, data, state_names=None):
         self.model = model
-        self.data = data.astype(np.int)
+        self.data = data
+        if not isinstance(state_names, dict):
+            self.state_names = {node: self._get_state_names(node) for node in model.nodes()}
+        else:
+            self.state_names = dict()
+            for node in model.nodes():
+                if node in state_names:
+                    if not set(self._get_state_names(node)) <= set(state_names[node]):
+                        raise ValueError("Data contains unexpected states for variable '{0}'.".format(str(node)))
+                    self.state_names[node] = sorted(state_names[node])
+                else:
+                    self.state_names[node] = self._get_state_names(node)
 
-        get_node_card = lambda _node, _data: _data.ix[:, _node].value_counts().shape[0]
-        self.node_card = {_node: get_node_card(_node, data) for _node in self.model.nodes()}
+    def _get_state_names(self, variable):
+        states = sorted(list(self.data.ix[:, variable].unique()))
+        return states
diff --git a/pgmpy/tests/test_estimators/test_MaximumLikelihoodEstimator.py b/pgmpy/tests/test_estimators/test_MaximumLikelihoodEstimator.py
index 6d7f8749..72e4277e 100644
--- a/pgmpy/tests/test_estimators/test_MaximumLikelihoodEstimator.py
+++ b/pgmpy/tests/test_estimators/test_MaximumLikelihoodEstimator.py
@@ -10,15 +10,54 @@ class TestMLE(unittest.TestCase):
     def setUp(self):
         self.m1 = BayesianModel([('A', 'C'), ('B', 'C')])
         self.d1 = pd.DataFrame(data={'A': [0, 0, 1], 'B': [0, 1, 0], 'C': [1, 1, 0]})
+        self.cpds = [TabularCPD('A', 2, [[2.0/3], [1.0/3]]),
+                     TabularCPD('B', 2, [[2.0/3], [1.0/3]]),
+                     TabularCPD('C', 2, [[0.0, 0.0, 1.0, 0.5],
+                                         [1.0, 1.0, 0.0, 0.5]],
+                                evidence=['A', 'B'], evidence_card=[2, 2])]
+        self.mle1 = MaximumLikelihoodEstimator(self.m1, self.d1)
 
     def test_get_parameters_missing_data(self):
-        mle = MaximumLikelihoodEstimator(self.m1, self.d1)
-        cpds = [TabularCPD('A', 2, [[2.0/3], [1.0/3]]),
-                TabularCPD('C', 2, [[0.0, 0.0, 1.0, 0.5],
-                                    [1.0, 1.0, 0.0, 0.5]],
-                           evidence=['A', 'B'], evidence_card=[2, 2]),
-                TabularCPD('B', 2, [[2.0/3], [1.0/3]])]
+        self.assertSetEqual(set(self.mle1.get_parameters()), set(self.cpds))
 
+    def test_estimate_cpd(self):
+        self.assertEqual(self.mle1._estimate_cpd('A'), self.cpds[0])
+        self.assertEqual(self.mle1._estimate_cpd('B'), self.cpds[1])
+        self.assertEqual(self.mle1._estimate_cpd('C'), self.cpds[2])
+
+    def test_state_names1(self):
+        m = BayesianModel([('A', 'B')])
+        d = pd.DataFrame(data={'A': [2, 3, 8, 8, 8], 'B': ['X', 'O', 'X', 'O', 'X']})
+        cpd_b = TabularCPD('B', 2, [[0, 1, 1.0 / 3], [1, 0, 2.0 / 3]],
+                           evidence=['A'], evidence_card=[3])
+        mle2 = MaximumLikelihoodEstimator(m, d)
+        self.assertEqual(mle2._estimate_cpd('B'), cpd_b)
+
+    def test_state_names2(self):
+        m = BayesianModel([('Light?', 'Color'), ('Fruit', 'Color')])
+        d = pd.DataFrame(data={'Fruit': ['Apple', 'Apple', 'Apple', 'Banana', 'Banana'],
+                               'Light?': [True,   True,   False,   False,    True],
+                               'Color': ['red',   'green', 'black', 'black',  'yellow']})
+        color_cpd = TabularCPD('Color', 4, [[1, 0, 1, 0], [0, 0.5, 0, 0],
+                                            [0, 0.5, 0, 0], [0, 0, 0, 1]],
+                               evidence=['Fruit', 'Light?'], evidence_card=[2, 2])
+        mle2 = MaximumLikelihoodEstimator(m, d)
+        self.assertEqual(mle2._estimate_cpd('Color'), color_cpd)
+
+    def test_class_init(self):
+        mle = MaximumLikelihoodEstimator(self.m1, self.d1,
+                                         state_names={'A': [0, 1], 'B': [0, 1], 'C': [0, 1]})
+        self.assertSetEqual(set(mle.get_parameters()), set(self.cpds))
+
+    def test_nonoccurring_values(self):
+        mle = MaximumLikelihoodEstimator(self.m1, self.d1,
+                                         state_names={'A': [0, 1, 23], 'B': [0, 1], 'C': [0, 42, 1], 1: [2]})
+        cpds = [TabularCPD('A', 3, [[2.0/3], [1.0/3], [0]]),
+                TabularCPD('B', 2, [[2.0/3], [1.0/3]]),
+                TabularCPD('C', 3, [[0.0, 0.0, 1.0, 1.0/3, 1.0/3, 1.0/3],
+                                    [1.0, 1.0, 0.0, 1.0/3, 1.0/3, 1.0/3],
+                                    [0.0, 0.0, 0.0, 1.0/3, 1.0/3, 1.0/3]],
+                           evidence=['A', 'B'], evidence_card=[3, 2])]
         self.assertSetEqual(set(mle.get_parameters()), set(cpds))
 
     def tearDown(self):