diff --git a/infra/graph_generators.py b/infra/graph_generators.py new file mode 100644 index 0000000..4c12dd3 --- /dev/null +++ b/infra/graph_generators.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Sun Mar 17 13:00:33 2019 + +@author: Zargham +""" +import networkx as nx + +def lineGraphGen(N, bidir=False,nodeTypeName='vanilla',edgeTypeName='vanilla'): + + line = nx.path_graph(N, create_using=nx.MultiDiGraph) + if not(bidir): + G = line + else: + edges = line.edges + G = nx.MultiDiGraph() + for e in edges: + G.add_edge(e[0],e[1]) + G.add_edge(e[1],e[0]) + + nx.set_node_attributes(G,nodeTypeName, 'type') + nx.set_edge_attributes(G,edgeTypeName, 'type') + + return G + +def starGraphGen(N, kind='sink',nodeTypeName='vanilla',edgeTypeName='vanilla'): + + star = nx.star_graph(N) + G = nx.MultiDiGraph() + + for e in star.edges: + if (kind == 'source') or (kind == 'bidir'): + G.add_edge(e[0],e[1]) + if (kind == 'sink') or (kind == 'bidir'): + G.add_edge(e[1],e[0]) + + nx.set_node_attributes(G,nodeTypeName, 'type') + nx.set_edge_attributes(G,edgeTypeName, 'type') + + return G + +def circleGraphGen(N, bidir=False,nodeTypeName='vanilla',edgeTypeName='vanilla' ): + + circle = nx.cycle_graph(N, create_using=nx.MultiDiGraph) + if not(bidir): + G = circle + else: + edges = circle.edges + G = nx.MultiDiGraph() + for e in edges: + G.add_edge(e[0],e[1]) + G.add_edge(e[1],e[0]) + + nx.set_node_attributes(G,nodeTypeName, 'type') + nx.set_edge_attributes(G,edgeTypeName, 'type') + + return G + +def treeGraphGen(r,h, kind='sink',nodeTypeName='vanilla',edgeTypeName='vanilla'): + + tree = nx.balanced_tree(r,h, create_using=nx.MultiDiGraph) + + if kind=='source': + G = tree + elif kind =='sink': + G = nx.MultiDiGraph() + for e in tree.edges: + G.add_edge(e[1],e[0]) + elif kind == 'bidir': + G = nx.MultiDiGraph() + for e in tree.edges: + G.add_edge(e[1],e[0]) + G.add_edge(e[0],e[1]) + + nx.set_node_attributes(G,nodeTypeName, 'type') + nx.set_edge_attributes(G,edgeTypeName, 'type') + + return G \ No newline at end of file diff --git a/infra/graph_generators_test.py b/infra/graph_generators_test.py new file mode 100644 index 0000000..2b5f824 --- /dev/null +++ b/infra/graph_generators_test.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Wed Mar 20 17:17:54 2019 + +@author: Zargham +""" +import networkx as nx +import unittest +from graph_generators import lineGraphGen, starGraphGen, treeGraphGen + +class GraphGenerators_LineGraph(unittest.TestCase): + def test_lineGraphGen_sanity(self): + g = lineGraphGen(4) + self.assertEqual(list(g.edges()), [(0, 1), (1, 2), (2, 3)]) + + def test_lineGraphGen_bidir(self): + g = lineGraphGen(4, bidir=True) + # (node, neighbor) + self.assertEqual(list(g.edges()), [(0, 1), (1, 0), (1, 2), (2, 1), (2, 3), (3, 2)]) + +class GraphGenerators_StarGraph(unittest.TestCase): + +# how to validate that 'kind' is either 'source', 'sink', or 'bidir' ?? + + def test_starGraph_sink(self): + g = starGraphGen(4) + self.assertEqual(list(g.edges()), [(1, 0), (2, 0), (3, 0), (4, 0)]) + + def test_starGraph_source(self): + g = starGraphGen(4, kind='source') + self.assertEqual(list(g.edges()), [(0, 1), (0, 2), (0, 3), (0, 4)]) + + def test_starGraph_bidir(self): + g = starGraphGen(4, kind='bidir') + self.assertEqual(list(g.edges()), [(0, 1), (0, 2), (0, 3), (0, 4), (1, 0), (2, 0), (3, 0), (4, 0)]) + +class GraphGenerators_TreeGraph(unittest.TestCase): + #note: our default is diff from networkx's default in this case (sink vs source) + def test_treeGraph_source_sink(self): + g = treeGraphGen(2,2) + self.assertEqual(list(g.edges()), [(1, 0), (2, 0), (3, 1), (4, 1), (5, 2), (6, 2)]) + + def test_treeGraph_source(self): + g = treeGraphGen(2,2, kind='source') + self.assertEqual(list(g.edges()), [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)]) + + def test_treeGraph_bidir(self): + g = treeGraphGen(2,2, kind='bidir') + self.assertEqual(list(g.edges()), [ + (1, 0), (1, 3), (1, 4), (0, 1), (0, 2), (2, 0), + (2, 5), (2, 6), (3, 1), (4, 1), (5, 2), (6, 2) + ]) + +# this is the way I might test that the types assigned in Z's functions are correct +# I think it would be a good idea to abstract out that logic into a helper function +# so that the code isn't duplicated in every generator + +class GraphGenerators_SetType(unittest.TestCase): + def test_lineGraphGen_setType(self): + g = lineGraphGen(4, nodeTypeName='foo',edgeTypeName='bar') + + # each node and edge should have the correct 'type' value + for k, v in nx.get_node_attributes(g, 'type').items(): + self.assertEqual(v, 'foo') + + for k, v in nx.get_edge_attributes(g, 'type').items(): + self.assertEqual(v, 'bar') \ No newline at end of file