def edges_iter(self):
for edge in self.contact_network.edges():
uNum,vNum = edge
attr = self.contact_network[uNum][vNum]['attribute']
u = Node(self, self.num_to_name[uNum], uNum)
v = Node(self, self.num_to_name[vNum], vNum)
yield Edge(u,v,attr)
def __init__(self, edge_list=None):
# if ModuleFactory is just testing, do nothing
if edge_list is None:
return
# set up NetworkX and graph
self.contact_network = DiGraph()
self.name_to_num = {} # map original node names to numbers
self.num_to_name = [] # map numbers to original node names
self.nodes = set() # store all nodes
self.uninfected_nodes = set() # store uninfected nodes
self.infected_nodes = set() # store infected nodes
self.transmissions = [] # store u->v transmission as (u,v,time)
self.directed = False # am I directed (True) or undirected (False)?
# read in Contact Network as node+edge list
for line in edge_list:
line = line.strip()
if len(line) == 0 or line[0] == '#':
continue
parts = line.split('\t')
# add node to contact network
if parts[0] == 'NODE':
name = parts[1]
if name in self.name_to_num:
raise ValueError("Node %r already exists!" % name)
num = len(self.num_to_name)
self.name_to_num[name] = num
self.num_to_name.append(name)
self.contact_network.add_node(num)
if parts[2] == '.':
self.contact_network.node[num]['attribute'] = set()
else:
self.contact_network.node[num]['attribute'] = set([e.strip().upper() for e in parts[2].split(',')])
self.contact_network.node[num]['infections_from'] = []
self.contact_network.node[num]['infections_to'] = []
self.contact_network.node[num]['infection_trees'] = []
# add edge to contact network
elif parts[0] == 'EDGE':
uName = parts[1]
if uName not in self.name_to_num:
raise ValueError("Node %r does not exist!" % uName)
vName = parts[2]
if vName not in self.name_to_num:
raise ValueError("Node %r does not exist!" % vName)
uNum = self.name_to_num[uName]
vNum = self.name_to_num[vName]
self.contact_network.add_edge(uNum,vNum)
if parts[3] == '.':
self.contact_network[uNum][vNum]['attribute'] = set()
else:
self.contact_network[uNum][vNum]['attribute'] = set(parts[3].split(','))
if parts[4] not in {'d','u'}:
raise ValueError("Invalid directionality: %r" % parts[4])
if parts[4] == 'u': # undirected edge, so add v to u too
self.contact_network.add_edge(vNum,uNum)
self.contact_network[vNum][uNum]['attribute'] = self.contact_network[uNum][vNum]['attribute']
else:
self.directed = True
# invalid type
else:
raise ValueError("Invalid type in list: %r" % parts[0])
# create sets of nodes
for node in self.contact_network.nodes():
self.nodes.add(Node(self, self.num_to_name[node], node))
for node in self.nodes:
self.uninfected_nodes.add(node)
def get_edges_to(self,node):
nx_edges = self.contact_network.in_edges(self.name_to_num[node.get_name()])
return [Edge(Node(self, self.num_to_name[uNum], uNum), node, self.contact_network[uNum][vNum]['attribute']) for uNum,vNum in nx_edges]
def get_node(self, name):
if name is None:
return None
return Node(self, name, self.name_to_num[name])