def construct_graph(args):
# Load pre-processed datasets
v = pd.read_csv(args.v_input, index_col=False, delim_whitespace=True)
e = pd.read_csv(args.e_input, index_col=False, delim_whitespace=True)
logging.info(
"Setting up graph data with {} nodes, {} edges and {} clusters".format(
v.shape[0], e.shape[0], len(v["cluster"].unique())))
# Coin dataframes into SQL for graphframes
v_schema = StructType([
StructField("id", IntegerType(), True),
StructField("cluster", IntegerType(), True)
])
v = sql_context.createDataFrame(v, schema=v_schema).dropDuplicates(['id'])
e_schema = StructType([
StructField("src", IntegerType(), True),
StructField("dst", IntegerType(), True)
])
e = sql_context.createDataFrame(e, schema=e_schema)
# Generate graph before simulation starts
g = GF.GraphFrame(v, e)
return g
#Step 8-6-2 : Creating DataFrames of edge of given graph.
edgeDataList = [('A', 'C'), ('A', 'B'), ('B', 'A'), ('B', 'C'), ('B', 'G'),
('B', 'F'), ('C', 'A'), ('C', 'B'), ('C', 'F'), ('C', 'D'),
('D', 'C'), ('D', 'F'), ('D', 'E'), ('E', 'D'), ('E', 'F'),
('F', 'B'), ('F', 'C'), ('F', 'D'), ('F', 'E'), ('F', 'G'),
('G', 'B'), ('G', 'F')]
edgeRDD = sc.parallelize(edgeDataList, 4)
edgeRDD.take(4)
edgeRDDRows = edgeRDD.map(lambda data: Row(data[0], data[1]))
edgeRDDRows.take(4)
sourceColumn = StructField('src', StringType(), True)
destinationColumn = StructField('dst', StringType(), True)
edgeSchema = StructType([sourceColumn, destinationColumn])
edgeDataFrame = sqlContext.createDataFrame(edgeRDDRows, edgeSchema)
edgeDataFrame.show(5)
#Step 8-6-3 : Creating GraphFrame object.
import graphframes.graphframe as gfm
ourGraph = gfm.GraphFrame(verticesDataFrame, edgeDataFrame)
ourGraph.vertices.show(5)
ourGraph.edges.show(5)
#Step 8-6-3 : Running Breath First Search Algorithm.
bfsPath = ourGraph.bfs(fromExpr="id='D'", toExpr="id='G'")
bfsPath.show()
def simulate(g, p_is, p_id, p_ih, p_ir, p_hr, p_hd, t_latent, t_infectious,
num_i_seeds, num_time_steps):
# select the vertices as a list
nodes = list(g.vertices.select("id").toPandas()["id"])
i_nodes = list(random.sample(nodes, num_i_seeds))
s_nodes = list(set(nodes) - set(i_nodes))
h_nodes = []
e_nodes = []
r_nodes = []
d_nodes = []
duration = 0
nodes_counter = pd.DataFrame(
columns=["n_{}".format(x) for x in ['s', 'e', 'i', 'r', 'h', 'd']],
index=[i for i in range(num_time_steps)])
nodes_counter.loc[0] = [
len(s_nodes),
len(e_nodes),
len(i_nodes),
len(r_nodes),
len(h_nodes),
len(d_nodes)
]
# ADD "neighbors" column: select neighbors of a node based on a graph (used in S_flow step)
neighbor = g.find("(a)-[e]->(b)").drop("e").groupBy('a.id').agg(
collect_list('b.id').alias('neighbors'))
g_neighbor = neighbor.join(g.vertices, ['id'], "right_outer")
g = GF.GraphFrame(g_neighbor, e)
# ADD "state" column
# At t0: number of I nodes and H nodes are based on user-defined functions. ALL OTHER NODES are assumed to be S
g_temp = g.vertices.withColumn(
"state",
when(g.vertices.id.isin(i_nodes), "I").otherwise(
when(g.vertices.id.isin(h_nodes), "H").otherwise("S")))
# ADD "i_days" and "e_days" column
# At t0: 1 for all I_nodes and 0 for all others
g_temp = g_temp.withColumn("e_days", lit(0))
g_temp = g_temp.withColumn(
"i_days",
when(g.vertices.id.isin(i_nodes), lit(1)).otherwise(lit(0)))
g = GF.GraphFrame(g_temp, e)
# TO DO: allow p_is to be a vector of rates (time-dependent)
for step in range(1, num_time_steps + 1):
H_flow(h_nodes, r_nodes, d_nodes, p_hr, p_hd)
I_flow(i_nodes, r_nodes, h_nodes, d_nodes, p_id, p_ih, p_ir)
new_I_nodes = E_flow(g, e_nodes, i_nodes, t_latent)
new_E_nodes = S_flow(g, s_nodes, e_nodes, i_nodes, r_nodes, h_nodes,
d_nodes, p_is, p_id, p_ih, p_ir, t_infectious)
# update the state column using the new lists ("x_nodes")
g_temp = g.vertices.withColumn(
"state",
when(g.vertices.id.isin(s_nodes), "S").otherwise(
when(g.vertices.id.isin(e_nodes), "E").otherwise(
when(g.vertices.id.isin(i_nodes), "I").otherwise(
when(g.vertices.id.isin(r_nodes), "R").otherwise(
when(g.vertices.id.isin(h_nodes),
"H").otherwise("D"))))))
# update i_days and e_days (1. initialize to zero the newly turned I nodes; 2. add one to previously exposed or infectious nodes)
old_I_nodes = list(set(i_nodes) - set(new_I_nodes))
g_temp = g_temp.withColumn(
"i_days",
when(g_temp.id.isin(new_I_nodes), 1).otherwise(
when(g_temp.id.isin(old_I_nodes),
g_temp.i_days + 1).otherwise(0)))
old_E_nodes = list(set(e_nodes) - set(new_E_nodes))
g_temp = g_temp.withColumn(
"e_days",
when(g_temp.id.isin(new_E_nodes), 1).otherwise(
when(g_temp.id.isin(old_E_nodes),
g_temp.e_days + 1).otherwise(0)))
# finish upating the vertices --> let's update the graph!
g = GF.GraphFrame(g_temp, e)
duration += 1
if len(i_nodes) == 0:
print("TERMINATED: No more infectious nodes left to update")
break
nodes_counter.loc[step] = [
len(s_nodes),
len(e_nodes),
len(i_nodes),
len(r_nodes),
len(h_nodes),
len(d_nodes)
]
return nodes_counter, duration
StructField("sex", IntegerType(), True),
StructField("cluster", IntegerType(), True)
])
#v = sql_context.createDataFrame(v, schema = v_schema).dropDuplicates(['id'])
# ^ the above line leads to a random number of nodes in the graph.
v = sql_context.createDataFrame(v, schema=v_schema)
e_schema = StructType([
StructField("src", IntegerType(), True),
StructField("dst", IntegerType(), True)
])
e = sql_context.createDataFrame(e, schema=e_schema)
# Generate graph before simulation starts
g = GF.GraphFrame(v, e)
# Define parameters
p_is = 0.5
p_id = 0.0134
p_ih = 0.0678
p_ir = 0.3945
p_hr = 0.3945 / 2
p_hd = 0.0419
t_latent = 5
t_infectious = 5
num_i_seeds = 20
num_time_steps = 3