Repository: top-tier-ci-coders/pgmpy

Issue:
Move `BayesianModel.active_trail_nodes` method
Move BayesianModel.active_trail_nodes to DAG

Discussion:

Patch:
diff --git a/pgmpy/base/DAG.py b/pgmpy/base/DAG.py
index 50cdb9f..2d9f4c0 100644
--- a/pgmpy/base/DAG.py
+++ b/pgmpy/base/DAG.py
@@ -443,3 +443,100 @@ class DAG(nx.DiGraph):
         blanket_nodes = set(blanket_nodes)
         blanket_nodes.remove(node)
         return list(blanket_nodes)
+
+    def active_trail_nodes(self, variables, observed=None):
+        """
+        Returns a dictionary with the given variables as keys and all the nodes reachable
+        from that respective variable as values.
+        Parameters
+        ----------
+        variables: str or array like
+            variables whose active trails are to be found.
+        observed : List of nodes (optional)
+            If given the active trails would be computed assuming these nodes to be observed.
+        Examples
+        --------
+        >>> from pgmpy.base import DAG
+        >>> student = DAG()
+        >>> student.add_nodes_from(['diff', 'intel', 'grades'])
+        >>> student.add_edges_from([('diff', 'grades'), ('intel', 'grades')])
+        >>> student.active_trail_nodes('diff')
+        {'diff': {'diff', 'grades'}}
+        >>> student.active_trail_nodes(['diff', 'intel'], observed='grades')
+        {'diff': {'diff', 'intel'}, 'intel': {'diff', 'intel'}}
+        References
+        ----------
+        Details of the algorithm can be found in 'Probabilistic Graphical Model
+        Principles and Techniques' - Koller and Friedman
+        Page 75 Algorithm 3.1
+        """
+        if observed:
+            observed_list = observed if isinstance(observed, (list, tuple)) else [observed]
+        else:
+            observed_list = []
+        ancestors_list = self._get_ancestors_of(observed_list)
+
+        # Direction of flow of information
+        # up ->  from parent to child
+        # down -> from child to parent
+
+        active_trails = {}
+        for start in variables if isinstance(variables, (list, tuple)) else [variables]:
+            visit_list = set()
+            visit_list.add((start, 'up'))
+            traversed_list = set()
+            active_nodes = set()
+            while visit_list:
+                node, direction = visit_list.pop()
+                if (node, direction) not in traversed_list:
+                    if node not in observed_list:
+                        active_nodes.add(node)
+                    traversed_list.add((node, direction))
+                    if direction == 'up' and node not in observed_list:
+                        for parent in self.predecessors(node):
+                            visit_list.add((parent, 'up'))
+                        for child in self.successors(node):
+                            visit_list.add((child, 'down'))
+                    elif direction == 'down':
+                        if node not in observed_list:
+                            for child in self.successors(node):
+                                visit_list.add((child, 'down'))
+                        if node in ancestors_list:
+                            for parent in self.predecessors(node):
+                                visit_list.add((parent, 'up'))
+            active_trails[start] = active_nodes
+        return active_trails
+
+    def _get_ancestors_of(self, obs_nodes_list):
+        """
+        Returns a dictionary of all ancestors of all the observed nodes including the
+        node itself.
+        Parameters
+        ----------
+        obs_nodes_list: string, list-type
+            name of all the observed nodes
+        Examples
+        --------
+        >>> from pgmpy.base import DAG
+        >>> model = DAG([('D', 'G'), ('I', 'G'), ('G', 'L'),
+        ...                        ('I', 'L')])
+        >>> model._get_ancestors_of('G')
+        {'D', 'G', 'I'}
+        >>> model._get_ancestors_of(['G', 'I'])
+        {'D', 'G', 'I'}
+        """
+        if not isinstance(obs_nodes_list, (list, tuple)):
+            obs_nodes_list = [obs_nodes_list]
+
+        for node in obs_nodes_list:
+            if node not in self.nodes():
+                raise ValueError('Node {s} not in not in graph'.format(s=node))
+
+        ancestors_list = set()
+        nodes_list = set(obs_nodes_list)
+        while nodes_list:
+            node = nodes_list.pop()
+            if node not in ancestors_list:
+                nodes_list.update(self.predecessors(node))
+            ancestors_list.add(node)
+        return ancestors_list
diff --git a/pgmpy/models/BayesianModel.py b/pgmpy/models/BayesianModel.py
index 22a3e44..b5e4c92 100644
--- a/pgmpy/models/BayesianModel.py
+++ b/pgmpy/models/BayesianModel.py
@@ -390,110 +390,6 @@ class BayesianModel(DAG):
                                      " is not equal to 1.".format(node=node))
         return True
 
-    def _get_ancestors_of(self, obs_nodes_list):
-        """
-        Returns a dictionary of all ancestors of all the observed nodes including the
-        node itself.
-
-        Parameters
-        ----------
-        obs_nodes_list: string, list-type
-            name of all the observed nodes
-
-        Examples
-        --------
-        >>> from pgmpy.models import BayesianModel
-        >>> model = BayesianModel([('D', 'G'), ('I', 'G'), ('G', 'L'),
-        ...                        ('I', 'L')])
-        >>> model._get_ancestors_of('G')
-        {'D', 'G', 'I'}
-        >>> model._get_ancestors_of(['G', 'I'])
-        {'D', 'G', 'I'}
-        """
-        if not isinstance(obs_nodes_list, (list, tuple)):
-            obs_nodes_list = [obs_nodes_list]
-
-        for node in obs_nodes_list:
-            if node not in self.nodes():
-                raise ValueError('Node {s} not in not in graph'.format(s=node))
-
-        ancestors_list = set()
-        nodes_list = set(obs_nodes_list)
-        while nodes_list:
-            node = nodes_list.pop()
-            if node not in ancestors_list:
-                nodes_list.update(self.predecessors(node))
-            ancestors_list.add(node)
-        return ancestors_list
-
-    def active_trail_nodes(self, variables, observed=None):
-        """
-        Returns a dictionary with the given variables as keys and all the nodes reachable
-        from that respective variable as values.
-
-        Parameters
-        ----------
-
-        variables: str or array like
-            variables whose active trails are to be found.
-
-        observed : List of nodes (optional)
-            If given the active trails would be computed assuming these nodes to be observed.
-
-        Examples
-        --------
-        >>> from pgmpy.models import BayesianModel
-        >>> student = BayesianModel()
-        >>> student.add_nodes_from(['diff', 'intel', 'grades'])
-        >>> student.add_edges_from([('diff', 'grades'), ('intel', 'grades')])
-        >>> student.active_trail_nodes('diff')
-        {'diff': {'diff', 'grades'}}
-        >>> student.active_trail_nodes(['diff', 'intel'], observed='grades')
-        {'diff': {'diff', 'intel'}, 'intel': {'diff', 'intel'}}
-
-        References
-        ----------
-        Details of the algorithm can be found in 'Probabilistic Graphical Model
-        Principles and Techniques' - Koller and Friedman
-        Page 75 Algorithm 3.1
-        """
-        if observed:
-            observed_list = observed if isinstance(observed, (list, tuple)) else [observed]
-        else:
-            observed_list = []
-        ancestors_list = self._get_ancestors_of(observed_list)
-
-        # Direction of flow of information
-        # up ->  from parent to child
-        # down -> from child to parent
-
-        active_trails = {}
-        for start in variables if isinstance(variables, (list, tuple)) else [variables]:
-            visit_list = set()
-            visit_list.add((start, 'up'))
-            traversed_list = set()
-            active_nodes = set()
-            while visit_list:
-                node, direction = visit_list.pop()
-                if (node, direction) not in traversed_list:
-                    if node not in observed_list:
-                        active_nodes.add(node)
-                    traversed_list.add((node, direction))
-                    if direction == 'up' and node not in observed_list:
-                        for parent in self.predecessors(node):
-                            visit_list.add((parent, 'up'))
-                        for child in self.successors(node):
-                            visit_list.add((child, 'down'))
-                    elif direction == 'down':
-                        if node not in observed_list:
-                            for child in self.successors(node):
-                                visit_list.add((child, 'down'))
-                        if node in ancestors_list:
-                            for parent in self.predecessors(node):
-                                visit_list.add((parent, 'up'))
-            active_trails[start] = active_nodes
-        return active_trails
-
     def local_independencies(self, variables):
         """
         Returns an instance of Independencies containing the local independencies


Test patch:
diff --git a/pgmpy/tests/test_models/test_BayesianModel.py b/pgmpy/tests/test_models/test_BayesianModel.py
index be8d9c4..45b7c2f 100644
--- a/pgmpy/tests/test_models/test_BayesianModel.py
+++ b/pgmpy/tests/test_models/test_BayesianModel.py
@@ -11,6 +11,7 @@ import pgmpy.tests.help_functions as hf
 from pgmpy.factors.discrete import TabularCPD, JointProbabilityDistribution, DiscreteFactor
 from pgmpy.independencies import Independencies
 from pgmpy.estimators import BayesianEstimator, BaseEstimator, MaximumLikelihoodEstimator
+from pgmpy.base import DAG
 
 
 class TestBaseModelCreation(unittest.TestCase):
@@ -237,17 +238,19 @@ class TestBayesianModelCPD(unittest.TestCase):
     def setUp(self):
         self.G = BayesianModel([('d', 'g'), ('i', 'g'), ('g', 'l'),
                                 ('i', 's')])
+        self.G2 = DAG([('d', 'g'), ('i', 'g'), ('g', 'l'),
+                                ('i', 's')])
 
     def test_active_trail_nodes(self):
-        self.assertEqual(sorted(self.G.active_trail_nodes('d')['d']), ['d', 'g', 'l'])
-        self.assertEqual(sorted(self.G.active_trail_nodes('i')['i']), ['g', 'i', 'l', 's'])
-        self.assertEqual(sorted(self.G.active_trail_nodes(['d', 'i'])['d']), ['d', 'g', 'l'])
+        self.assertEqual(sorted(self.G2.active_trail_nodes('d')['d']), ['d', 'g', 'l'])
+        self.assertEqual(sorted(self.G2.active_trail_nodes('i')['i']), ['g', 'i', 'l', 's'])
+        self.assertEqual(sorted(self.G2.active_trail_nodes(['d', 'i'])['d']), ['d', 'g', 'l'])
 
     def test_active_trail_nodes_args(self):
-        self.assertEqual(sorted(self.G.active_trail_nodes(['d', 'l'], observed='g')['d']), ['d', 'i', 's'])
-        self.assertEqual(sorted(self.G.active_trail_nodes(['d', 'l'], observed='g')['l']), ['l'])
-        self.assertEqual(sorted(self.G.active_trail_nodes('s', observed=['i', 'l'])['s']), ['s'])
-        self.assertEqual(sorted(self.G.active_trail_nodes('s', observed=['d', 'l'])['s']), ['g', 'i', 's'])
+        self.assertEqual(sorted(self.G2.active_trail_nodes(['d', 'l'], observed='g')['d']), ['d', 'i', 's'])
+        self.assertEqual(sorted(self.G2.active_trail_nodes(['d', 'l'], observed='g')['l']), ['l'])
+        self.assertEqual(sorted(self.G2.active_trail_nodes('s', observed=['i', 'l'])['s']), ['s'])
+        self.assertEqual(sorted(self.G2.active_trail_nodes('s', observed=['d', 'l'])['s']), ['g', 'i', 's'])
 
     def test_is_active_trail_triplets(self):
         self.assertTrue(self.G.is_active_trail('d', 'l'))