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