Displaying inheritance nested within a pedigree

[hyanwong]

Off the back of https://tskit.dev/msprime/docs/stable/pedigrees.html#visualising-pedigrees, I tweaked the function to allow more customisation, and (more interestingly) to be able to show the tree sequence edges embedded in the pedigree. I'm showing the pairs or node IDs rather than the individuals IDs as node labels here, and using hexagons to avoid confusion with the squares and circles used in the normal tskit visualizations.

There's a slight issue in that you can't tell which of the two nodes within each individual is being targeted by the red lines, but I can't think of an easy way around that.

def draw_pedigree(ped_ts, figsize=None, font_size=None, ts_edge_colour=None, **kwargs):
    # If edge_colour is None, do not overlay tskit edges over the pedigree links
    if figsize is not None:
        plt.figure(figsize=figsize)
    G = nx.DiGraph()
    for ind in ped_ts.individuals():
        time = ped_ts.node(ind.nodes[0]).time
        pop = ped_ts.node(ind.nodes[0]).population
        G.add_node(ind.id, time=time, population=pop)
        for p in ind.parents:
            if p != tskit.NULL:
                G.add_edge(p, ind.id)
    pos = nx.multipartite_layout(G, subset_key="time", align="horizontal")
    colours = plt.rcParams['axes.prop_cycle'].by_key()['color']
    node_colours = [colours[node_attr["population"]] for node_attr in G.nodes.values()]
    nx.draw_networkx(G, pos, with_labels=False, node_color=node_colours, **kwargs)
    labels = {i.id: "\n".join(str(u) for u in i.nodes) for i in ped_ts.individuals()}
    nx.draw_networkx_labels(G, pos, labels, font_size=font_size)
    if ts_edge_colour is not None:
        new_G = nx.create_empty_copy(G)
        for e in ped_ts.edges():
            parent = ped_ts.node(e.parent)
            child = ped_ts.node(e.child)
            if parent.individual != tskit.NULL and child.individual != tskit.NULL:
                new_G.add_edge(parent.individual, child.individual)
        nx.draw_networkx_edges(new_G, pos, edge_color=ts_edge_colour, arrowstyle="-")
    plt.show()

For example, with the non-recombining forward_WF code in the fwd simulations tutorial

from matplotlib import pyplot as plt
import networkx as nx

ts = forward_WF(5, 5e4, generations=7, random_seed=5)
ts = ts.simplify(np.flatnonzero(ts.nodes_time==0), filter_nodes=False, filter_individuals=False, keep_unary=True)
draw_pedigree(ts, figsize=(5, 8), ts_edge_colour="red", font_size=5.5, node_size=250, width=3, node_shape="H", arrowstyle="simple")

[hyanwong]

Possibly nicer using grey and black:

params = {
    "node_color": "skyblue",
    "font_size": 5.8,
    "node_size":350,
    "node_shape": "H",
    "edge_color": "lightgrey",
    "arrowstyle": "simple",
}
draw_pedigree(ts, width=2, ts_edge_colour="black", **params)