241 lines
7.7 KiB
Python
Executable File
241 lines
7.7 KiB
Python
Executable File
#! /usr/bin/env python3
|
|
|
|
import matplotlib.pyplot as plt
|
|
import networkx as nx
|
|
import random, math
|
|
import json
|
|
import sys, os
|
|
import string
|
|
import typer
|
|
from enum import Enum
|
|
|
|
|
|
# Enums & Consts
|
|
|
|
# To add a new node type, add appropriate entries to the nodeType and nodeTypeSwitch
|
|
class nodeType(Enum):
|
|
DESKTOP = "desktop" # waku desktop config
|
|
MOBILE = "mobile" # waku mobile config
|
|
|
|
nodeTypeSwitch = {
|
|
nodeType.DESKTOP : "rpc-admin = true\nkeep-alive = true\n",
|
|
nodeType.MOBILE : "rpc-admin = true\nkeep-alive = true\n"
|
|
}
|
|
|
|
|
|
# To add a new network type, add appropriate entries to the networkType and networkTypeSwitch
|
|
# the networkTypeSwitch is placed before generate_network(): fwd declaration mismatch with typer/python :/
|
|
class networkType(Enum):
|
|
CONFIGMODEL = "configmodel"
|
|
SCALEFREE = "scalefree" # power law
|
|
NEWMANWATTSSTROGATZ = "newmanwattsstrogatz" # mesh, smallworld
|
|
BARBELL = "barbell" # partition
|
|
BALANCEDTREE = "balancedtree" # committees?
|
|
STAR = "star" # spof
|
|
|
|
|
|
NW_DATA_FNAME = "network_data.json"
|
|
NODE_PREFIX = "waku"
|
|
SUBNET_PREFIX = "subnetwork"
|
|
|
|
|
|
### I/O related fns ##############################################################
|
|
|
|
# Dump to a json file
|
|
def write_json(dirname, json_dump):
|
|
fname = os.path.join(dirname, NW_DATA_FNAME)
|
|
with open(fname, "w") as f:
|
|
json.dump(json_dump, f, indent=2)
|
|
|
|
|
|
def write_toml(dirname, node_name, toml):
|
|
fname = os.path.join(dirname, f"{node_name}.toml")
|
|
with open(fname, "w") as f:
|
|
f.write(toml)
|
|
|
|
|
|
# Draw the network and output the image to a file; does not account for subnets yet
|
|
def draw(dirname, H):
|
|
nx.draw(H, pos=nx.kamada_kawai_layout(H), with_labels=True)
|
|
fname = os.path.join(dirname, NW_DATA_FNAME)
|
|
plt.savefig(f"{os.path.splitext(fname)[0]}.png", format="png")
|
|
plt.show()
|
|
|
|
|
|
# Has trouble with non-integer/non-hashable keys
|
|
def read_json(fname):
|
|
with open(fname) as f:
|
|
jdata = json.load(f)
|
|
return nx.node_link_graph(jdata)
|
|
|
|
|
|
# check if the required dir can be created
|
|
def exists_or_nonempty(dirname):
|
|
if not os.path.exists(dirname):
|
|
return False
|
|
elif not os.path.isfile(dirname) and os.listdir(dirname):
|
|
print(f"{dirname}: exists and not empty")
|
|
return True
|
|
elif os.path.isfile(dirname):
|
|
print(f"{dirname}: exists but not a directory")
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
|
|
### topics related fns #############################################################
|
|
|
|
# Generate a random string of upper case chars
|
|
def generate_random_string(n):
|
|
return "".join(random.choice(string.ascii_uppercase) for _ in range(n))
|
|
|
|
|
|
# Generate the topics - topic followed by random UC chars - Eg, topic_XY"
|
|
def generate_topics(num_topics):
|
|
topic_len = int(math.log(num_topics)/math.log(26)) + 1 # base is 26 - upper case letters
|
|
topics = {i: f"topic_{generate_random_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)
|
|
hi = random.randint(lo + 1, n)
|
|
sublist = []
|
|
for i in range(lo, hi):
|
|
sublist.append(topics[i])
|
|
return sublist
|
|
|
|
|
|
### network processing related fns #################################################
|
|
|
|
# Network Types
|
|
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 be 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=D=3
|
|
def generate_newmanwattsstrogatz_graph(n):
|
|
return nx.newman_watts_strogatz_graph(n, 3, 0.5)
|
|
|
|
|
|
def generate_barbell_graph(n):
|
|
return nx.barbell_graph(int(n/2), 1)
|
|
|
|
|
|
def generate_balanced_tree(n, fanout=3):
|
|
height = int(math.log(n)/math.log(fanout))
|
|
return nx.balanced_tree(fanout, height)
|
|
|
|
|
|
def generate_star_graph(n):
|
|
return nx.star_graph(n)
|
|
|
|
|
|
networkTypeSwitch = {
|
|
networkType.CONFIGMODEL : generate_config_model,
|
|
networkType.SCALEFREE : generate_scalefree_graph,
|
|
networkType.NEWMANWATTSSTROGATZ : generate_newmanwattsstrogatz_graph,
|
|
networkType.BARBELL : generate_barbell_graph,
|
|
networkType.BALANCEDTREE: generate_balanced_tree,
|
|
networkType.STAR : generate_star_graph
|
|
}
|
|
|
|
|
|
# Generate the network from nw type
|
|
def generate_network(n, network_type):
|
|
return postprocess_network(networkTypeSwitch.get(network_type)(n))
|
|
|
|
|
|
# Label the generated network with prefix
|
|
def postprocess_network(G):
|
|
G = nx.Graph(G) # prune out parallel/multi edges
|
|
G.remove_edges_from(nx.selfloop_edges(G)) # remove the self-loops
|
|
mapping = {i: f"{NODE_PREFIX}_{i}" for i in range(len(G))}
|
|
return nx.relabel_nodes(G, mapping) # label the nodes
|
|
|
|
|
|
def generate_subnets(G, num_subnets):
|
|
n = len(G.nodes)
|
|
if num_subnets == n: # if num_subnets == size of the network
|
|
return {f"{NODE_PREFIX}_{i}": f"{SUBNET_PREFIX}_{i}" for i in range(n)}
|
|
|
|
lst = list(range(n))
|
|
random.shuffle(lst)
|
|
offsets = sorted(random.sample(range(0, n), num_subnets - 1))
|
|
offsets.append(n-1)
|
|
|
|
start = 0
|
|
subnets = {}
|
|
subnet_id = 0
|
|
for end in offsets:
|
|
for i in range(start, end+1):
|
|
subnets[f"{NODE_PREFIX}_{lst[i]}"] = f"{SUBNET_PREFIX}_{subnet_id}"
|
|
start = end
|
|
subnet_id += 1
|
|
return subnets
|
|
|
|
|
|
### file format related fns ###########################################################
|
|
#Generate per node toml configs
|
|
def generate_toml(topics, node_type=nodeType.DESKTOP):
|
|
topic_str = " ".join(get_random_sublist(topics)) # space separated topics
|
|
return f"{nodeTypeSwitch.get(node_type)}topics = \"{topic_str}\"\n"
|
|
|
|
|
|
# Generates network-wide json and per-node toml and writes them
|
|
def generate_and_write_files(dirname, num_topics, num_subnets, G):
|
|
topics = generate_topics(num_topics)
|
|
subnets = generate_subnets(G, num_subnets)
|
|
json_dump = {}
|
|
for node in G.nodes:
|
|
write_toml(dirname, node, generate_toml(topics)) # per node toml
|
|
json_dump[node] = {}
|
|
json_dump[node]["static-nodes"] = []
|
|
for edge in G.edges(node):
|
|
json_dump[node]["static-nodes"].append(edge[1])
|
|
json_dump[node][SUBNET_PREFIX] = subnets[node]
|
|
write_json(dirname, json_dump) # network wide json
|
|
|
|
|
|
### the main ##########################################################################
|
|
def main(
|
|
dirname: str = "WakuNetwork", num_nodes: int = 4, num_topics: int = 1,
|
|
network_type: networkType = networkType.NEWMANWATTSSTROGATZ.value,
|
|
node_type: nodeType = nodeType.DESKTOP.value,
|
|
num_subnets: int = -1,
|
|
num_partitions: int = 1):
|
|
|
|
# sanity checks
|
|
if num_partitions > 1:
|
|
raise ValueError(f"--num-partitions {num_partitions}, Sorry, we do not yet support partitions")
|
|
if num_subnets > num_nodes:
|
|
raise ValueError(f"num_subnets must be <= num_nodes: num_subnets={num_subnets}, num_nodes={num_nodes}")
|
|
if num_subnets == -1:
|
|
num_subnets = num_nodes
|
|
|
|
# Generate the network
|
|
G = generate_network(num_nodes, network_type)
|
|
|
|
# Refuse to overwrite non-empty dirs
|
|
if exists_or_nonempty(dirname) :
|
|
sys.exit(1)
|
|
os.makedirs(dirname, exist_ok=True)
|
|
|
|
# Generate file format specific data structs and write the files; optionally, draw the network
|
|
generate_and_write_files(dirname, num_topics, num_subnets, G)
|
|
#draw(dirname, G)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
typer.run(main)
|