From eb1b4e06bbee479ff0937cc9920dd26a9f725f37 Mon Sep 17 00:00:00 2001
From: Cathy Yeh <cathy@driver.xyz>
Date: Mon, 13 Nov 2017 13:23:18 -0800
Subject: beliefs package

---
 beliefs/__init__.py             |  0
 beliefs/models/DirectedGraph.py | 35 +++++++++++++++++++++++++++++++++++
 beliefs/models/__init__.py      |  0
 3 files changed, 35 insertions(+)
 create mode 100644 beliefs/__init__.py
 create mode 100644 beliefs/models/DirectedGraph.py
 create mode 100644 beliefs/models/__init__.py

diff --git a/beliefs/__init__.py b/beliefs/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/beliefs/models/DirectedGraph.py b/beliefs/models/DirectedGraph.py
new file mode 100644
index 0000000..8dfb9bd
--- /dev/null
+++ b/beliefs/models/DirectedGraph.py
@@ -0,0 +1,35 @@
+import networkx as nx
+
+
+class DirectedGraph(nx.DiGraph):
+    """
+    Base class for all directed graphical models.
+    """
+    def __init__(self, edges, node_labels):
+        """
+        Input:
+            edges: an edge list, e.g. [(parent1, child1), (parent1, child2)]
+            node_labels: a list of strings of node labels
+        """
+        super().__init__()
+        if edges is not None:
+            self.add_edges_from(edges)
+        if node_labels is not None:
+            self.add_nodes_from(node_labels)
+
+    def get_leaves(self):
+        """
+        Returns a list of leaves of the graph.
+        """
+        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.
+        """
+        return [node for node, in_degree in self.in_degree().items() if
+                in_degree == 0]
+
+    def get_topologically_sorted_nodes(self, reverse=False):
+        return nx.topological_sort(self, reverse=reverse)
diff --git a/beliefs/models/__init__.py b/beliefs/models/__init__.py
new file mode 100644
index 0000000..e69de29
-- 
cgit v1.2.3