Repository: pgmpy/pgmpy
Pull request #899 (2017-09-28 04:54:54)

BayesianModel.py: Updated get_cardinality method and added test
In reference to issue https://github.com/pgmpy/pgmpy/issues/722 and https://github.com/pgmpy/pgmpy/issues/893 . As per the discussion now get_cardinality of BayesianModel accepts node as an optional argument and return the cardinality of the node. If node is not specified returns a dictionary with the given factors as keys and their respective cardinality as values.

coveralls: 
[![Coverage Status](https://coveralls.io/builds/13469859/badge)](https://coveralls.io/builds/13469859)

Coverage increased (+0.005%) to 94.77% when pulling **51e0681cc8d2bfa24420948e3f591cf2e4adf94c on lohani2280:branch8** into **dd00b1e5aa059bccf0c0acf35c9fda8c5f1ae371 on pgmpy:dev**.

coveralls: 
[![Coverage Status](https://coveralls.io/builds/13484201/badge)](https://coveralls.io/builds/13484201)

Coverage increased (+0.02%) to 94.77% when pulling **cf9c8416d0d44f177d5023b8ba4ecb672b62d122 on lohani2280:branch8** into **cf6641c13defd8a04336f4e539f8c7a53de9e195 on pgmpy:dev**.

coveralls: 
[![Coverage Status](https://coveralls.io/builds/13484201/badge)](https://coveralls.io/builds/13484201)

Coverage increased (+0.02%) to 94.77% when pulling **cf9c8416d0d44f177d5023b8ba4ecb672b62d122 on lohani2280:branch8** into **cf6641c13defd8a04336f4e539f8c7a53de9e195 on pgmpy:dev**.

coveralls: 
[![Coverage Status](https://coveralls.io/builds/13484466/badge)](https://coveralls.io/builds/13484466)

Coverage increased (+0.02%) to 94.77% when pulling **27982ca3a9578bb578d193f967301a80640ffe3e on lohani2280:branch8** into **cf6641c13defd8a04336f4e539f8c7a53de9e195 on pgmpy:dev**.

ankurankan: Also when you are done with the changes, please squash all your commits into a single commit and have a look at this https://robots.thoughtbot.com/5-useful-tips-for-a-better-commit-message for writing a good commit message.

Patch:
diff --git a/pgmpy/models/BayesianModel.py b/pgmpy/models/BayesianModel.py
index 1823c05c..a729b9d5 100644
--- a/pgmpy/models/BayesianModel.py
+++ b/pgmpy/models/BayesianModel.py
@@ -316,20 +316,49 @@ class BayesianModel(DirectedGraph):
                 cpd = self.get_cpds(cpd)
             self.cpds.remove(cpd)
 
-    def get_cardinality(self, node):
+    def get_cardinality(self, node=None):
         """
         Returns the cardinality of the node. Throws an error if the CPD for the
         queried node hasn't been added to the network.
 
         Parameters
         ----------
-        node: Any hashable python object.
+        node: Any hashable python object(optional).
+              The node whose cardinality we want. If node is not specified returns a
+              dictionary with the given variable as keys and their respective cardinality
+              as values.
 
         Returns
         -------
-        int: The cardinality of the node.
+        int or dict : If node is specified returns the cardinality of the node.
+                      If node is not specified returns a dictionary with the given
+                      variable as keys and their respective cardinality as values.
+
+        Examples
+        --------
+        >>> from pgmpy.models import BayesianModel
+        >>> from pgmpy.factors.discrete import TabularCPD
+        >>> student = BayesianModel([('diff', 'grade'), ('intel', 'grade')])
+        >>> cpd_diff = TabularCPD('diff',2,[[0.6,0.4]]);
+        >>> cpd_intel = TabularCPD('intel',2,[[0.7,0.3]]);
+        >>> cpd_grade = TabularCPD('grade', 2, [[0.1, 0.9, 0.2, 0.7],
+        ...                                     [0.9, 0.1, 0.8, 0.3]],
+        ...                                 ['intel', 'diff'], [2, 2])
+        >>> student.add_cpds(cpd_diff,cpd_intel,cpd_grade)
+        >>> student.get_cardinality()
+        defaultdict(int, {'diff': 2, 'grade': 2, 'intel': 2})
+
+        >>> student.get_cardinality('intel')
+        2
         """
-        return self.get_cpds(node).cardinality[0]
+
+        if node:
+            return self.get_cpds(node).cardinality[0]
+        else:
+            cardinalities = defaultdict(int)
+            for cpd in self.cpds:
+                cardinalities[cpd.variable] = cpd.cardinality[0]
+            return cardinalities
 
     def check_model(self):
         """


Test patch:
diff --git a/pgmpy/tests/test_models/test_BayesianModel.py b/pgmpy/tests/test_models/test_BayesianModel.py
index 8d5a304e..357d4795 100644
--- a/pgmpy/tests/test_models/test_BayesianModel.py
+++ b/pgmpy/tests/test_models/test_BayesianModel.py
@@ -144,6 +144,14 @@ class TestBayesianModelMethods(unittest.TestCase):
     def test_get_ancestors_of_failure(self):
         self.assertRaises(ValueError, self.G2._get_ancestors_of, 'h')
 
+    def test_get_cardinality(self):
+        self.assertDictEqual(self.G1.get_cardinality(), {'diff': 2, 'intel': 3, 'grade': 3})
+
+    def test_get_cardinality_with_node(self):
+        self.assertEqual(self.G1.get_cardinality('diff'), 2)
+        self.assertEqual(self.G1.get_cardinality('intel'), 3)
+        self.assertEqual(self.G1.get_cardinality('grade'), 3)
+
     def test_local_independencies(self):
         self.assertEqual(self.G.local_independencies('a'), Independencies(['a', ['b', 'c']]))
         self.assertEqual(self.G.local_independencies('c'), Independencies(['c', ['a', 'd', 'e'], 'b']))