Turning a tree sequence into a networkx graph

[hyanwong]

There's information in the tutorials about turning a tree into a networkx graph, but some complications with doing it for an entire tree sequence (see #1296). But in case anyone else wants to do it for a tree sequence rather than a tree, here the most efficient way I've found to do this, via a pandas DataFrame. This adds the time and flags as a node attribute. We could add other attributes too, even metadata.

Note that we need nx.MultiDiGraph because we could have 2 edges connecting the same nodes (at different genomic locations)

import msprime
import networkx as nx
import pandas as pd

test_ts = msprime.simulate(10, recombination_rate=1)

def to_networkx_graph(ts):
    edges=ts.tables.edges
    G = nx.from_pandas_edgelist(
        pd.DataFrame({'source': edges.parent, 'target': edges.child, 'left': edges.left, 'right': edges.right}),
        edge_attr=True,
        create_using=nx.MultiDiGraph
    )
    nx.set_node_attributes(G, {n.id: {'flags':n.flags, 'time': n.time} for n in ts.nodes()})
    return G

G = to_networkx_graph(test_ts)