Repository: pgmpy/pgmpy
Pull request #812 (2017-01-03 17:34:12)

Active_trail_nodes accepts iterables
Solving #722 first.

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

Coverage increased (+0.001%) to 96.194% when pulling **6f5e80626b2a7bb63c031c2c0824b33272b63399 on raghavg7796:PR#722.1** into **5d885cb211b5ac9ccfdc6ace639c9ef62c9138df on pgmpy:dev**.

raghavg7796: @ankurankan I think this is ready for merge. Please review.

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

Coverage increased (+0.001%) to 96.194% when pulling **53dde1420ce681caeb489201972eeec5f06aaa22 on raghavg7796:PR#722.1** into **5d885cb211b5ac9ccfdc6ace639c9ef62c9138df on pgmpy:dev**.

ankurankan: @raghavg7796 Let me know when this is ready.

Patch:
diff --git a/pgmpy/models/BayesianModel.py b/pgmpy/models/BayesianModel.py
index 413abe6f..8963358b 100644
--- a/pgmpy/models/BayesianModel.py
+++ b/pgmpy/models/BayesianModel.py
@@ -391,17 +391,19 @@ class BayesianModel(DirectedGraph):
             ancestors_list.add(node)
         return ancestors_list
 
-    def active_trail_nodes(self, start, observed=None):
+    def active_trail_nodes(self, variables, observed=None):
         """
-        Returns all the nodes reachable from start via an active trail.
+        Returns a dictionary with the given variables as keys and all the nodes reachable
+        from that respective variable as values.
 
         Parameters
         ----------
 
-        start: Graph node
+        variables: str or array like
+            variables whose active trails are to be found.
 
         observed : List of nodes (optional)
-            If given the active trail would be computed assuming these nodes to be observed.
+            If given the active trails would be computed assuming these nodes to be observed.
 
         Examples
         --------
@@ -410,9 +412,9 @@ class BayesianModel(DirectedGraph):
         >>> student.add_nodes_from(['diff', 'intel', 'grades'])
         >>> student.add_edges_from([('diff', 'grades'), ('intel', 'grades')])
         >>> student.active_trail_nodes('diff')
-        {'diff', 'grade'}
-        >>> student.active_trail_nodes('diff', observed='grades')
-        {'diff', 'intel'}
+        {'diff': {'diff', 'grades'}}
+        >>> student.active_trail_nodes(['diff', 'intel'], observed='grades')
+        {'diff': {'diff', 'intel'}, 'intel': {'diff', 'intel'}}
 
         References
         ----------
@@ -421,7 +423,7 @@ class BayesianModel(DirectedGraph):
         Page 75 Algorithm 3.1
         """
         if observed:
-            observed_list = [observed] if isinstance(observed, str) else observed
+            observed_list = observed if isinstance(observed, (list, tuple)) else [observed]
         else:
             observed_list = []
         ancestors_list = self._get_ancestors_of(observed_list)
@@ -430,29 +432,32 @@ class BayesianModel(DirectedGraph):
         # up ->  from parent to child
         # down -> from child to parent
 
-        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':
+        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:
-                        for child in self.successors(node):
-                            visit_list.add((child, 'down'))
-                    if node in ancestors_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'))
-        return active_nodes
+                        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):
         """
@@ -491,7 +496,7 @@ class BayesianModel(DirectedGraph):
             return descendents
 
         independencies = Independencies()
-        for variable in [variables] if isinstance(variables, str) else variables:
+        for variable in variables if isinstance(variables, (list, tuple)) else [variables]:
             non_descendents = set(self.nodes()) - {variable} - set(dfs(variable))
             parents = set(self.get_parents(variable))
             if non_descendents - parents:
@@ -527,7 +532,7 @@ class BayesianModel(DirectedGraph):
         >>> student.is_active_trail('grades', 'sat')
         True
         """
-        if end in self.active_trail_nodes(start, observed):
+        if end in self.active_trail_nodes(start, observed)[start]:
             return True
         else:
             return False
@@ -556,7 +561,7 @@ class BayesianModel(DirectedGraph):
             for r in range(len(rest)):
                 for observed in itertools.combinations(rest, r):
                     d_seperated_variables = rest - set(observed) - set(
-                        self.active_trail_nodes(start, observed=observed))
+                        self.active_trail_nodes(start, observed=observed)[start])
                     if d_seperated_variables:
                         independencies.add_assertions([start, d_seperated_variables, observed])
 


Test patch:
diff --git a/pgmpy/tests/test_models/test_BayesianModel.py b/pgmpy/tests/test_models/test_BayesianModel.py
index cc66420c..72531785 100644
--- a/pgmpy/tests/test_models/test_BayesianModel.py
+++ b/pgmpy/tests/test_models/test_BayesianModel.py
@@ -213,14 +213,15 @@ class TestBayesianModelCPD(unittest.TestCase):
                                 ('i', 's')])
 
     def test_active_trail_nodes(self):
-        self.assertEqual(sorted(self.G.active_trail_nodes('d')), ['d', 'g', 'l'])
-        self.assertEqual(sorted(self.G.active_trail_nodes('i')), ['g', 'i', 'l', 's'])
+        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'])
 
     def test_active_trail_nodes_args(self):
-        self.assertEqual(sorted(self.G.active_trail_nodes('d', observed='g')), ['d', 'i', 's'])
-        self.assertEqual(sorted(self.G.active_trail_nodes('l', observed='g')), ['l'])
-        self.assertEqual(sorted(self.G.active_trail_nodes('s', observed=['i', 'l'])), ['s'])
-        self.assertEqual(sorted(self.G.active_trail_nodes('s', observed=['d', 'l'])), ['g', 'i', 's'])
+        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'])
 
     def test_is_active_trail_triplets(self):
         self.assertTrue(self.G.is_active_trail('d', 'l'))