def build_graph(df, cluster_id, case_id, date_col, color, gen_mean, gen_sd):
"""
Generate a directed graph from data on transmission tree.
Node color is determined by node attributes, e.g. case severity or gender.
df = pandas dataframe
"""
clusters = basics.cluster_builder(df=df, cluster_id=cluster_id, \
case_id=case_id, date_col=date_col, attr_col=color, \
gen_mean=gen_mean, gen_sd=gen_sd)
G = nx.DiGraph()
G.add_nodes_from(clusters['case_id'])
edgelist = [pair for pair in clusters[['source_node']].dropna().itertuples()]
G.add_edges_from(edgelist)
nx.set_node_attributes(G, 'date', clusters['time'].to_dict())
nx.set_node_attributes(G, 'pltdate', clusters['pltdate'].to_dict())
nx.set_node_attributes(G, 'source_node', clusters['source_node'].to_dict())
nx.set_node_attributes(G, color, clusters[color].to_dict())
nx.set_node_attributes(G, 'index_node', clusters['index_node'].to_dict())
G = nx.DiGraph.reverse(G)
for i in G.nodes():
G.node[i]['generation'] = _generations(G, i)
return G
def build_graph(df, cluster_id, case_id, date_col, color, gen_mean, gen_sd):
"""
Generate a directed graph from data on transmission tree.
Node color is determined by node attributes, e.g. case severity or gender.
df = pandas dataframe
"""
clusters = basics.cluster_builder(df=df, cluster_id=cluster_id, \
case_id=case_id, date_col=date_col, attr_col=color, \
gen_mean=gen_mean, gen_sd=gen_sd)
G = nx.DiGraph()
G.add_nodes_from(clusters['case_id'])
edgelist = [
pair for pair in clusters[['source_node']].dropna().itertuples()
]
G.add_edges_from(edgelist)
nx.set_node_attributes(G, 'date', clusters['time'].to_dict())
nx.set_node_attributes(G, 'pltdate', clusters['pltdate'].to_dict())
nx.set_node_attributes(G, 'source_node', clusters['source_node'].to_dict())
nx.set_node_attributes(G, color, clusters[color].to_dict())
nx.set_node_attributes(G, 'index_node', clusters['index_node'].to_dict())
G = nx.DiGraph.reverse(G)
for i in G.nodes():
G.node[i]['generation'] = _generations(G, i)
return G
def test_cluster_to_tuple():
df = _test_data()
df['datetime'] = df['date'].map(basics.date_convert)
df_out = basics.cluster_builder(df, 'cluster', 'id', 'datetime', \
'sex', 2, 1)
df_out = df_out.sort('case_id')
#sanity check
assert df_out.ix[0]['case_id'] == 0
assert df_out.ix[3]['case_id'] == 3
#index nodes
assert df_out.ix[0]['index_node'] == 0
assert df_out.ix[4]['index_node'] == 1
#source nodes
assert df_out.ix[5]['source_node'] == 4
