Fixed #872

pgmpy · Apr 24, 2017 · 1eaaeb9 · 1eaaeb9
1 parent c302a5d
commit 1eaaeb9
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/pgmpy/models/BayesianModel.py b/pgmpy/models/BayesianModel.py
@@ -285,7 +285,8 @@ def get_cpds(self, node=None):
             for cpd in self.cpds:
                 if cpd.variable == node:
                     return cpd
-            raise ValueError("CPD not added for the node: {node}".format(node=node))
+            else:
+                return None
         else:
             return self.cpds
 
@@ -346,7 +347,9 @@ def check_model(self):
         for node in self.nodes():
             cpd = self.get_cpds(node=node)
 
-            if isinstance(cpd, TabularCPD):
+            if cpd is None:
+                raise ValueError('No CPD associated with {}'.format(node))
+            elif isinstance(cpd, TabularCPD):
                 evidence = cpd.variables[:0:-1]
                 parents = self.get_parents(node)
                 if set(evidence if evidence else []) != set(parents if parents else []):

diff --git a/pgmpy/tests/test_models/test_BayesianModel.py b/pgmpy/tests/test_models/test_BayesianModel.py
@@ -282,7 +282,7 @@ def test_get_cpds1(self):
         self.assertRaises(ValueError, self.model.get_cpds, 'B')
 
         self.model.add_node('B')
-        self.assertRaises(ValueError, self.model.get_cpds, 'B')
+        self.assertIsNone(self.model.get_cpds('B'))
 
     def test_add_single_cpd(self):
         cpd_s = TabularCPD('s', 2, np.random.rand(2, 2), ['i'], [2])