diff options
author | Cathy Yeh <cathy@driver.xyz> | 2017-11-13 13:23:18 -0800 |
---|---|---|
committer | Cathy Yeh <cathy@driver.xyz> | 2017-11-13 13:23:18 -0800 |
commit | eb1b4e06bbee479ff0937cc9920dd26a9f725f37 (patch) | |
tree | 9d3ddc2ad513a7ba5712e53a3ffd58268ff91182 | |
parent | 60c2c9ad85fd8685a69f21325c93479b2b5c59f9 (diff) | |
download | beliefs-eb1b4e06bbee479ff0937cc9920dd26a9f725f37.tar.gz beliefs-eb1b4e06bbee479ff0937cc9920dd26a9f725f37.tar.bz2 beliefs-eb1b4e06bbee479ff0937cc9920dd26a9f725f37.zip |
beliefs package
-rw-r--r-- | beliefs/__init__.py | 0 | ||||
-rw-r--r-- | beliefs/models/DirectedGraph.py | 35 | ||||
-rw-r--r-- | beliefs/models/__init__.py | 0 |
3 files changed, 35 insertions, 0 deletions
diff --git a/beliefs/__init__.py b/beliefs/__init__.py new file mode 100644 index 0000000..e69de29 --- /dev/null +++ b/beliefs/__init__.py 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 --- /dev/null +++ b/beliefs/models/__init__.py |