Tweaks (#2)

* added scalefree/newman_watts_strogatz models, fixed a bug

* Readme

* Readme

* added PR2 suggestions

Readme.md

@@ -27,7 +27,7 @@ gen_jsons.sh can generate given number of Waku networs and outputs them to a dir
 > usage: ./gen_jsons.sh <output_dir> <#json files needed> </br>
 
 ## generate_network.py
-generate_network.py can generate networks with specified number of nodes and topics. the network types currently supported is "configuration_model" and more are on the way. Use with Python3.
+generate_network.py can generate networks with specified number of nodes and topics. the network types currently supported is "configuration_model" and more are on the way. Use with Python3. Comment out the `#draw(fname, H)` line to visualise the generated graph.
 
 > usage: generate_network [-h] [-o <file_name>] [-n <#nodes>] [-t <#topics>]
                         [-T <type>] <br>

generate_network.py

@@ -1,34 +1,40 @@
 #! /usr/bin/env python3
 
-import matplotlib.pyplot as mp
+import matplotlib.pyplot as plt
 import networkx as nx
-import networkx.readwrite.json_graph
 import random, math
 import json
-import argparse,sys
+import argparse, os, sys
 
+# Dump to a json file
 def write_json(filename, data_2_dump):
     json.dump(data_2_dump, open(filename,'w'), indent=2)
 
+
 # has trouble with non-integer/non-hashable keys 
 def read_json(filename):
     with open(filename) as f:
         jdata = json.load(f)
     return nx.node_link_graph(jdata)
 
-def draw(H):
-    nx.draw(H, pos=nx.kamada_kawai_layout(H), with_labels=True)
-    mp.show()
-    mp.savefig("topology.png", format="PNG")
 
-def init_arg_parser() :
-    # Initialize parser, add arguments and set the defaults
+# Draw the network and output the image to a file
+def draw(fname, H):
+    nx.draw(H, pos=nx.kamada_kawai_layout(H), with_labels=True)
+    plt.savefig(os.path.splitext(fname)[0] + ".png", format="png")
+    plt.show()
+
+
+# Initialize parser, set the defaults, and extract the options
+def get_options():
     parser = argparse.ArgumentParser(
             prog = 'generate_network',
             description = '''Generates and outputs 
             the Waku network conforming to input parameters''',
-            epilog = '''The defaults are: -o "Topology.json"; 
-            -n 1;  -t 1;  -p 1; -T "configuration_model"''')
+            epilog = '''Defaults: -o "Topology.json"; 
+            -n 1; -t 1; -p 1; -T "configuration_model"
+            Supported nw types "configuration_model", "scalefree",
+            "newman_watts_strogatz"''')
     parser.add_argument("-o", "--output", 
             default='Topology.json', dest='fname', 
             help='output json filename for the Waku network', 
@@ -53,34 +59,28 @@ def init_arg_parser() :
 #            default=1, dest='num_edges', 
 #            help='The number of edges in the Waku network', 
 #            type=int, metavar='#edges>')
-    return parser
+    return parser.parse_args()
 
-# https://networkx.org/documentation/stable/reference/generated/networkx.generators.degree_seq.configuration_model.html
-def generate_config_model(n):
-    #degrees = nx.random_powerlaw_tree_sequence(n, tries=10000)
-    degrees = [random.randint(1, n) for i in range(n)]
-    if (sum(degrees)) % 2 != 0:         # adjust the degree sum to be even
-        degrees[-1] += 1
-    G = nx.configuration_model(degrees) # generate the graph
-    return G
 
+# Generate a random string (UC chars) of len n
 def generate_topic_string(n):
     rs = ""
     for _ in range(n):
-        r = random.randint(65, 65 + 26 - 1) # only letters
+        r = random.randint(65, 65 + 26 - 1) # generate a random UC char
         rs += (chr(r))                      # append the char generated
     return rs
 
+
+# Generate the topics - UC chars prefixed by "topic"
 def generate_topics(num_topics):
-    # generate the topics - uppercase chars prefixed by "topic"
     topics = []
     base = 26
     topic_len = int(math.log(num_topics)/math.log(base)) + 1
-    topics = {}
-    for i in range(num_topics):
-        topics[i] = "topic_" + generate_topic_string(topic_len)
+    topics = {i: f"topic_{generate_topic_string(topic_len)}" for i in range(num_topics)}
     return topics
 
+
+# Get a random sub-list of topics
 def get_random_sublist(topics):
     n = len(topics)
     lo = random.randint(0, n - 1)
@@ -90,28 +90,53 @@ def get_random_sublist(topics):
         sublist.append(topics[i])
     return sublist
 
-def generate_network(num_nodes, prefix):
+
+# Network Types
+# https://networkx.org/documentation/stable/reference/generated/networkx.generators.degree_seq.configuration_model.html
+def generate_config_model(n):
+    #degrees = nx.random_powerlaw_tree_sequence(n, tries=10000)
+    degrees = [random.randint(1, n) for i in range(n)]
+    if (sum(degrees)) % 2 != 0:         # adjust the degree to even
+        degrees[-1] += 1
+    return nx.configuration_model(degrees) # generate the graph
+
+
+def generate_scalefree_graph(n):
+    return  nx.scale_free_graph(n)
+
+
+# n must be larger than k
+def generate_newman_watts_strogatz_graph(n):
+    return nx.newman_watts_strogatz_graph(n, 12, 0.5)
+
+
+# Generate the network from nw type
+def generate_network(num_nodes, nw_type, prefix):
     G = nx.empty_graph()
     if nw_type == "configuration_model":
         G = generate_config_model(num_nodes)
+    elif nw_type == "scalefree":
+        G = generate_scalefree_graph(num_nodes)
+    elif nw_type == "newman_watts_strogatz":
+        G = generate_newman_watts_strogatz_graph(num_nodes) 
     else: 
         print(nw_type +": Unsupported network type")
         sys.exit(1)
     H = postprocess_network(G, prefix)
     return H
 
-# used by generate_dump_data - *ought* to be global for handling partitions
-ports_shifted = 0
 
+# used by generate_dump_data, *ought* to be global to handle partitions
+ports_shifted = 0
 def postprocess_network(G, prefix):
     G = nx.Graph(G)         # prune out parallel/multi edges
     G.remove_edges_from(nx.selfloop_edges(G))   # Removing self-loops
     # Labeling nodes to match waku containers
-    mapping = {}
-    for i in range(num_nodes):
-        mapping[i] = prefix + str(i)
+    mapping = {i: f"{prefix}{i}" for i in range(len(G))}
     return nx.relabel_nodes(G, mapping)
 
+
+# Generate dump data from the network and topics
 def generate_dump_data(H, topics):
     data_to_dump = {}
     global ports_shifted
@@ -125,31 +150,34 @@ def generate_dump_data(H, topics):
             data_to_dump[node]["static-nodes"].append(edge[1])
     return data_to_dump
 
-#extract the CLI arguments
-args = init_arg_parser().parse_args()
 
-#parameters to generate the network
-fname = args.fname
-num_nodes = args.num_nodes
-num_topics = args.num_topics
-nw_type = args.nw_type
-prefix = "waku_"     
-num_partitions = args.num_partitions
-#num_edges = args.num_edges     ## do we need to control #edges?
+def main():
+    #extract the CLI arguments and assign params
+    options = get_options()
+    fname = options.fname
+    num_nodes = options.num_nodes
+    num_topics = options.num_topics
+    nw_type = options.nw_type
+    prefix = "waku_"     
+    num_partitions = options.num_partitions
+    #num_edges = options.num_edges    ## need to control num_edges?
 
-if num_partitions > 1 :
-    print("-p",num_partitions, ": Sorry, we do not yet support partitions")
-    sys.exit(1)
+    if num_partitions > 1:
+        print("-p",num_partitions,
+                "Sorry, we do not yet support partitions")
+        sys.exit(1)
 
-# Generate the network and postprocess it
-H = generate_network(num_nodes, prefix)
+    # Generate the network and postprocess it
+    H = generate_network(num_nodes, nw_type, prefix)
+    # Generate the topics
+    topics = generate_topics(num_topics)
+    # Generate the dump data
+    dump_data = generate_dump_data(H, topics)
+    # Dump the network in a json file
+    write_json(fname, dump_data)
+    # Display the graph
+    draw(fname, H)
 
-#generate the topics
-topics = generate_topics(num_topics)
 
-# Generate the dump data
-dump_data = generate_dump_data(H, topics)
-
-# dump the network to the json file
-write_json(fname, dump_data)
-#draw(H)
+if __name__ == "__main__":
+    main()