Repository: pgmpy/pgmpy
Pull request #777 (2016-12-08 10:26:47)

Adds a few methods that will be used for modified Variable Elimination


Patch:
diff --git a/pgmpy/base/DirectedGraph.py b/pgmpy/base/DirectedGraph.py
index 61c75750..39f9b800 100644
--- a/pgmpy/base/DirectedGraph.py
+++ b/pgmpy/base/DirectedGraph.py
@@ -288,3 +288,30 @@ class DirectedGraph(nx.DiGraph):
                 itertools.combinations(self.get_parents(node), 2))
 
         return moral_graph
+
+    def get_leaves(self):
+        """
+        Returns a list of leaves of the graph.
+
+        Examples
+        --------
+        >>> from pgmpy.base import DirectedGraph
+        >>> graph = DirectedGraph([('A', 'B'), ('B', 'C'), ('B', 'D')])
+        >>> graph.get_leaves()
+        ['C', 'D']
+        """
+        return [node for node, out_degree in self.out_degree_iter() if
+                out_degree == 0]
+
+    def get_roots(self):
+        """
+        Returns a list of roots of the graph.
+
+        Examples
+        --------
+        >>> from pgmpy.base import DirectedGraph
+        >>> graph = DirectedGraph([('A', 'B'), ('B', 'C'), ('B', 'D'), ('E', 'B')])
+        >>> graph.get_roots()
+        ['A', 'E']
+        """
+        return [node for node, in_degree in self.in_degree().items() if in_degree == 0]
diff --git a/pgmpy/models/BayesianModel.py b/pgmpy/models/BayesianModel.py
index 04b2b6d2..6e6923ef 100644
--- a/pgmpy/models/BayesianModel.py
+++ b/pgmpy/models/BayesianModel.py
@@ -103,8 +103,8 @@ class BayesianModel(DirectedGraph):
         u,v : nodes
               Nodes can be any hashable python object.
 
-        EXAMPLE
-        -------
+        Examples
+        --------
         >>> from pgmpy.models import BayesianModel/home/abinash/software_packages/numpy-1.7.1
         >>> G = BayesianModel()
         >>> G.add_nodes_from(['grade', 'intel'])
@@ -118,6 +118,93 @@ class BayesianModel(DirectedGraph):
         else:
             super(BayesianModel, self).add_edge(u, v, **kwargs)
 
+    def remove_node(self, node):
+        """
+        Remove node from the model.
+
+        Removing a node also removes all the associated edges, removes the CPD
+        of the node and marginalizes the CPDs of it's children.
+
+        Parameters
+        ----------
+        node : node
+            Node which is to be removed from the model.
+
+        Returns
+        -------
+        None
+
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> import numpy as np
+        >>> from pgmpy.models import BayesianModel
+        >>> model = BayesianModel([('A', 'B'), ('B', 'C'),
+        ...                        ('A', 'D'), ('D', 'C')])
+        >>> values = pd.DataFrame(np.random.randint(low=0, high=2, size=(1000, 4)),
+        ...                       columns=['A', 'B', 'C', 'D'])
+        >>> model.fit(values)
+        >>> model.get_cpds()
+        [<TabularCPD representing P(A:2) at 0x7f28248e2438>,
+         <TabularCPD representing P(B:2 | A:2) at 0x7f28248e23c8>,
+         <TabularCPD representing P(C:2 | B:2, D:2) at 0x7f28248e2748>,
+         <TabularCPD representing P(D:2 | A:2) at 0x7f28248e26a0>]
+        >>> model.remove_node('A')
+        >>> model.get_cpds()
+        [<TabularCPD representing P(B:2) at 0x7f28248e23c8>,
+         <TabularCPD representing P(C:2 | B:2, D:2) at 0x7f28248e2748>,
+         <TabularCPD representing P(D:2) at 0x7f28248e26a0>]
+        """
+        affected_nodes = [v for u, v in self.edges() if u == node]
+
+        for affected_node in affected_nodes:
+            node_cpd = self.get_cpds(node=affected_node)
+            if node_cpd:
+                node_cpd.marginalize([node], inplace=True)
+
+        if self.get_cpds(node=node):
+            self.remove_cpds(node)
+        super(BayesianModel, self).remove_node(node)
+
+    def remove_nodes_from(self, nodes):
+        """
+        Remove multiple nodes from the model.
+
+        Removing a node also removes all the associated edges, removes the CPD
+        of the node and marginalizes the CPDs of it's children.
+
+        Parameters
+        ----------
+        nodes : list, set (iterable)
+            Nodes which are to be removed from the model.
+
+        Returns
+        -------
+        None
+
+        Examples
+        --------
+        >>> import pandas as pd
+        >>> import numpy as np
+        >>> from pgmpy.models import BayesianModel
+        >>> model = BayesianModel([('A', 'B'), ('B', 'C'),
+        ...                        ('A', 'D'), ('D', 'C')])
+        >>> values = pd.DataFrame(np.random.randint(low=0, high=2, size=(1000, 4)),
+        ...                       columns=['A', 'B', 'C', 'D'])
+        >>> model.fit(values)
+        >>> model.get_cpds()
+        [<TabularCPD representing P(A:2) at 0x7f28248e2438>,
+         <TabularCPD representing P(B:2 | A:2) at 0x7f28248e23c8>,
+         <TabularCPD representing P(C:2 | B:2, D:2) at 0x7f28248e2748>,
+         <TabularCPD representing P(D:2 | A:2) at 0x7f28248e26a0>]
+        >>> model.remove_nodes_from(['A', 'B'])
+        >>> model.get_cpds()
+        [<TabularCPD representing P(C:2 | D:2) at 0x7f28248e2a58>,
+         <TabularCPD representing P(D:2) at 0x7f28248e26d8>]
+        """
+        for node in nodes:
+            self.remove_node(node)
+
     def add_cpds(self, *cpds):
         """
         Add CPD (Conditional Probability Distribution) to the Bayesian Model.
@@ -761,3 +848,41 @@ class BayesianModel(DirectedGraph):
             return True
         else:
             return False
+
+    def copy(self):
+        """
+        Returns a copy of the model.
+
+        Returns
+        -------
+        BayesianModel: Copy of the model on which the method was called.
+
+        Examples
+        --------
+        >>> from pgmpy.models import BayesianModel
+        >>> from pgmpy.factors.discrete import TabularCPD
+        >>> model = BayesianModel([('A', 'B'), ('B', 'C')])
+        >>> cpd_a = TabularCPD('A', 2, [[0.2], [0.8]])
+        >>> cpd_b = TabularCPD('B', 2, [[0.3, 0.7], [0.7, 0.3]],
+                               evidence=['A'],
+                               evidence_card=[2])
+        >>> cpd_c = TabularCPD('C', 2, [[0.1, 0.9], [0.9, 0.1]],
+                               evidence=['B'],
+                               evidence_card=[2])
+        >>> model.add_cpds(cpd_a, cpd_b, cpd_c)
+        >>> copy_model = model.copy()
+        >>> copy_model.nodes()
+        ['C', 'A', 'B']
+        >>> copy_model.edges()
+        [('A', 'B'), ('B', 'C')]
+        >>> copy_model.get_cpds()
+        [<TabularCPD representing P(A:2) at 0x7f2824930a58>,
+         <TabularCPD representing P(B:2 | A:2) at 0x7f2824930a90>,
+         <TabularCPD representing P(C:2 | B:2) at 0x7f2824944240>]
+        """
+        model_copy = BayesianModel()
+        model_copy.add_nodes_from(self.nodes())
+        model_copy.add_edges_from(self.edges())
+        if self.cpds:
+            model_copy.add_cpds(*[cpd.copy() for cpd in self.cpds])
+        return model_copy


Test patch:
diff --git a/pgmpy/tests/test_base/test_DirectedGraph.py b/pgmpy/tests/test_base/test_DirectedGraph.py
index 6db9177e..4db1bd07 100644
--- a/pgmpy/tests/test_base/test_DirectedGraph.py
+++ b/pgmpy/tests/test_base/test_DirectedGraph.py
@@ -99,6 +99,19 @@ class TestDirectedGraphCreation(unittest.TestCase):
         self.assertListEqual(self.graph.predecessors('b'), ['a'])
         self.assertListEqual(self.graph.predecessors('c'), ['b'])
 
+    def test_get_leaves(self):
+        self.graph.add_edges_from([('A', 'B'), ('B', 'C'), ('B', 'D'),
+                                   ('D', 'E'), ('D', 'F'), ('A', 'G')])
+        self.assertEqual(sorted(self.graph.get_leaves()),
+                         sorted(['C', 'G', 'E', 'F']))
+
+    def test_get_roots(self):
+        self.graph.add_edges_from([('A', 'B'), ('B', 'C'), ('B', 'D'),
+                                   ('D', 'E'), ('D', 'F'), ('A', 'G')])
+        self.assertEqual(['A'], self.graph.get_roots())
+        self.graph.add_edge('H', 'G')
+        self.assertEqual(sorted(['A', 'H']), sorted(self.graph.get_roots()))
+
     def tearDown(self):
         del self.graph
 
diff --git a/pgmpy/tests/test_models/test_BayesianModel.py b/pgmpy/tests/test_models/test_BayesianModel.py
index b587eada..42104838 100644
--- a/pgmpy/tests/test_models/test_BayesianModel.py
+++ b/pgmpy/tests/test_models/test_BayesianModel.py
@@ -173,6 +173,34 @@ class TestBayesianModelMethods(unittest.TestCase):
         G3 = BayesianModel([('W', 'V'), ('W', 'X'), ('Y', 'X'), ('Z', 'Y')])
         self.assertFalse(G3.is_iequivalent(G2))
 
+    def test_copy(self):
+        model_copy = self.G1.copy()
+        self.assertEqual(sorted(self.G1.nodes()), sorted(model_copy.nodes()))
+        self.assertEqual(sorted(self.G1.edges()), sorted(model_copy.edges()))
+        self.assertNotEqual(id(self.G1.get_cpds('diff')),
+                            id(model_copy.get_cpds('diff')))
+
+        self.G1.remove_cpds('diff')
+        diff_cpd = TabularCPD('diff', 2, values=[[0.3], [0.7]])
+        self.G1.add_cpds(diff_cpd)
+        self.assertNotEqual(self.G1.get_cpds('diff'),
+                            model_copy.get_cpds('diff'))
+
+        self.G1.remove_node('intel')
+        self.assertNotEqual(sorted(self.G1.nodes()), sorted(model_copy.nodes()))
+        self.assertNotEqual(sorted(self.G1.edges()), sorted(model_copy.edges()))
+
+    def test_remove_node(self):
+        self.G1.remove_node('diff')
+        self.assertEqual(sorted(self.G1.nodes()), sorted(['grade', 'intel']))
+        self.assertRaises(ValueError, self.G1.get_cpds, 'diff')
+
+    def test_remove_nodes_from(self):
+        self.G1.remove_nodes_from(['diff', 'grade'])
+        self.assertEqual(sorted(self.G1.nodes()), sorted(['intel']))
+        self.assertRaises(ValueError, self.G1.get_cpds, 'diff')
+        self.assertRaises(ValueError, self.G1.get_cpds, 'grade')
+
     def tearDown(self):
         del self.G
         del self.G1