57 lines
1.3 KiB
Python
57 lines
1.3 KiB
Python
#!/usr/bin/python3
|
|
|
|
import networkx as nx
|
|
import matplotlib.pyplot as plt
|
|
from collections import OrderedDict
|
|
from operator import getitem
|
|
import itertools, os
|
|
|
|
# AGraph preserves attributes, networkx Graph does not.
|
|
# Many of the desired functions are in networkx.
|
|
# So import AGraph to keep attributes, then convert to Networkx.
|
|
A = nx.drawing.nx_agraph.to_agraph(nx.drawing.nx_pydot.read_dot("./1_mo_color_DOTFILE.dot"))
|
|
A.layout('dot')
|
|
A.draw('tree.png')
|
|
A.remove_node('\\n') # Remove "newline" node from newline end of dot file
|
|
G=nx.DiGraph(A)
|
|
|
|
color_map = []
|
|
color_d = {}
|
|
node_pos = {}
|
|
|
|
# Compartments
|
|
S = 0
|
|
I = 0
|
|
E = 0
|
|
R = 0
|
|
D = 0
|
|
|
|
# Params
|
|
beta = 0
|
|
delta = 0
|
|
gamma_r = 0
|
|
gamma_d = 0
|
|
mu = 0
|
|
epsilon = 0
|
|
omega = 0
|
|
|
|
for node in A:
|
|
color = A.get_node(node).attr.to_dict()['fillcolor']
|
|
str_pos = A.get_node(node).attr.to_dict()['pos']
|
|
coords = str_pos.split(',')
|
|
x = coords[0] # layout for draw function
|
|
y = coords[1]
|
|
node_pos[node] = float(x), float(y)
|
|
if color is None or color == '':
|
|
color_map.append("white")
|
|
color_d[node] = color
|
|
in_edges = list(G.in_edges(node))
|
|
elif color == 'yellow':
|
|
color_map.append(color)
|
|
color_d[node] = color
|
|
in_edges = list(G.in_edges(node))
|
|
else:
|
|
color_map.append(color)
|
|
color_d[node] = color
|
|
in_edges = list(G.in_edges(node))
|