155 lines
4.9 KiB
Python
155 lines
4.9 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, sys
|
|
|
|
# Change dir to location of this python file
|
|
print(os.getcwd())
|
|
#os.chdir(os.path.dirname(sys.argv[0]))
|
|
#print(os.getcwd())
|
|
|
|
def prep_seirds(weighted):
|
|
print("Prepping the SEIRDS model, using trivial weighting=", weighted)
|
|
# 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 = {} # used for drawing/graphing
|
|
|
|
# Compartments
|
|
S = 0
|
|
I_R = 0
|
|
I_D = 0
|
|
E = 0
|
|
R = 0
|
|
D = 0
|
|
|
|
ep_tmp = 0 # counter for epsilon
|
|
inf_ct = 0 # infection rate counter
|
|
recov_ct = 0 # recov rate counter
|
|
|
|
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))
|
|
out_edges = list(G.out_edges(node))
|
|
|
|
tmp_S = 1
|
|
tmp_recov = 0
|
|
tmp_inf = 0
|
|
for source in in_edges:
|
|
tmp_S = 1
|
|
# If previous node was infected, then we are recovered
|
|
if (color_d[source[0]] == 'red'):
|
|
if(weighted == 'False' or not in_edges):
|
|
tmp_recov = 1
|
|
else:
|
|
recov_ct = recov_ct + 1/len(in_edges) # trivial weighting
|
|
recov_ct = recov_ct + tmp_recov
|
|
for source in out_edges:
|
|
if (color_d[source[0]] == 'red'):
|
|
if(weighted == 'False' or not out_edges):
|
|
tmp_inf = 1
|
|
else:
|
|
inf_ct = inf_ct + 1/len(out_edges) # trivial weighting
|
|
inf_ct = inf_ct + tmp_inf
|
|
S = S + tmp_S
|
|
elif color == 'yellow':
|
|
color_map.append(color)
|
|
color_d[node] = color
|
|
in_edges = list(G.in_edges(node))
|
|
|
|
tmp_E = 1
|
|
tmp_R = 0
|
|
tmp_inf = 0
|
|
tmp_recov = 0
|
|
for source in in_edges:
|
|
# If previous node was infected, then we are recovered
|
|
if (color_d[source[0]] == 'red'):
|
|
tmp_E = 0
|
|
if(weighted == 'False' or not in_edges):
|
|
tmp_recov = 1
|
|
else:
|
|
recov_ct = recov_ct + 1/len(in_edges) # trivial weighting
|
|
if (color_d[source[0]] == '' or color_d[source[0]] == 'white'):
|
|
if(weighted == 'False' or not in_edges):
|
|
tmp_inf = 1 # add 1 for the inf counter
|
|
else:
|
|
inf_ct = inf_ct + 1/len(in_edges) # trivial weighting
|
|
E = E + tmp_E
|
|
recov_ct = recov_ct + tmp_recov
|
|
inf_ct = inf_ct + tmp_inf
|
|
else:
|
|
color_map.append(color)
|
|
color_d[node] = color
|
|
# Check if node dies
|
|
out_edges = list(G.out_edges(node))
|
|
if not out_edges:
|
|
D = D + 1
|
|
I_D = I_D + 1
|
|
else:
|
|
I_R = I_R + 1
|
|
# Check if imported
|
|
in_edges = list(G.in_edges(node))
|
|
if not in_edges:
|
|
ep_tmp = ep_tmp + 1
|
|
else:
|
|
tmp_inf = 0
|
|
for source in in_edges:
|
|
if (color_d[source[0]] == '' or color_d[source[0]] == 'white'):
|
|
if(weighted == 'False' or not in_edges):
|
|
tmp_inf = 1
|
|
else:
|
|
inf_ct = inf_ct + 1/len(in_edges) # trivial weightin
|
|
inf_ct = inf_ct + tmp_inf
|
|
|
|
# Params
|
|
beta = (inf_ct)/len(A) # rate of infec
|
|
delta = 1 # incubation period
|
|
#gamma_r = R/len(A) # recov rate
|
|
gamma_r = recov_ct/len(A)
|
|
gamma_d = D/len(A) # death rate
|
|
mu = D/(I_R+I_D) # fatality ratio
|
|
epsilon = ep_tmp/len(A) # infected import rate
|
|
omega = 1 # waning immunity rate
|
|
|
|
|
|
print("Model Compartments:")
|
|
print("S:", str(S))
|
|
print("E:", str(E))
|
|
print("I_R:", str(I_R))
|
|
print("I_D:", str(I_D))
|
|
print("R:", str(R))
|
|
print("D:", str(D))
|
|
print("\n")
|
|
|
|
print("Model Parameters:")
|
|
print("infect counter:", str(inf_ct))
|
|
print("beta:", str(beta))
|
|
print("delta:", str(delta))
|
|
print("gamma_r:", str(gamma_r))
|
|
print("gamma_d:", str(gamma_d))
|
|
print("mu:", str(mu))
|
|
print("epsilon:", str(epsilon))
|
|
print("omega:", str(omega))
|
|
|
|
return (S, E, I_R+I_D, R, D, beta, delta, gamma_r, gamma_d, mu, epsilon, omega)
|
|
|