Adding phasing errors to a tree sequence

[hyanwong]

I suspect it should be rather easy to detect phasing errors in an inferred tree sequence. As part of that, I thought it would be interesting to impose phase errors onto a simulated tree sequence. Here's (untested) code that I cooked up to do this, using the unsquash function (#2657)

import tskit
import numpy as np
import itertools

def unsquash(edge_table, positions, edges=None):
    """
    For a set of positions and a given set of edges (or all edges if ``edges`` is None),
    create a set of new edges which cover the same span but which are chopped up into
    separate edges at every specified position. This is essentially the opposite of
    EdgeTable.squash()
    """
    new_edges = tskit.EdgeTable()
    positions = np.unique(positions) # sort and uniquify
    skip = []
    if edges is not None:
        skip = np.ones(edge_table.num_rows, dtype=bool)
        skip[edges] = False
    for edge, do_skip in itertools.zip_longest(edge_table, skip, fillvalue=False):
        if do_skip:
            new_edges.append(edge)
            continue
        for l, r in itertools.pairwise(itertools.chain(
            [edge.left],
            positions[np.logical_and(positions > edge.left, positions < edge.right)],
            [edge.right]
        )):
            # This should create the edges in the correct order, so a sort isn't needed
            new_edges.append(edge.replace(left=l, right=r))
    edge_table.replace_with(new_edges)

def rephase(tables, first_sample, second_sample, positions):
    """
    Create a new tree sequence with the samples switching topological positions
    at each position, mimicking the effect of bad phasing
    """
    positions = np.unique(positions) # sort and uniquify
    edges_to_swap = np.isin(tables.edges.child, [first_sample, second_sample]) 
    unsquash(tables.edges, positions, np.where(edges_to_swap)[0])
    positions = set(positions)
    edge_child = tables.edges.child
    swap=False
    for diffs in tables.tree_sequence().edge_diffs():
        if diffs.interval.left in positions:
            swap = not swap
        if swap:
            for e in diffs.edges_in:
                if e.child == first_sample:
                    edge_child[e.id] = second_sample
                elif e.child == second_sample:
                    edge_child[e.id] = first_sample
    tables.edges.child = edge_child
    tables.sort()

Example:

import msprime
ts = msprime.simulate(10, random_seed=1)
tables = ts.dump_tables()
rephase(tables, 0, 1, [0.1, 0.3, 0.9])
tables.tree_sequence().draw_svg(time_scale="rank")