Opposite of EdgeTable.squash

[hyanwong]

This is something that I would have found useful in two different contexts (one when we split up the root in tsinfer into several nodes, the other when I have been trying to inject phasing error into a tree sequence). I stick it down here for posterity, but I wonder if an opposite of "squash" is a useful thing in the API. I've called it "unsquash", but something like "EdgeTable.slice" would work too. Because of using itertools.pairwise this only works on Python 3.10 or greater.

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]
        )):
            new_edges.append(edge.replace(left=l, right=r))
    edge_table.replace_with(new_edges)

[hyanwong]

This actually turns out to be useful in a number of scenarios. For instance, I just used it in awohns/unified_genealogy_paper#33 for splitting the ultimate ancestor, and it would be useful to create test tree sequences of multiple trees without performing simulations, e.g

import tskit
tables = tskit.Tree.generate_balanced(5, span=10).tree_sequence.dump_tables()
unsquash(tables.edges, np.arange(10), [0]) # split the first edge
ts = tables.tree_sequence()
ts.draw_svg()