Removing topology above local roots

[hyanwong]

Is there a simple way to identify (and then remove/adjust) the edges above local roots in the tree, while keeping all the nodes in the node table?

In fact, is there a way to simplify edges without touching the node table? Should there be?

[molpopgen]

In fact, is there a way to simplify edges without touching the node table? Should there be?

This is something that'll come up when I present. There is at least one use case for not remapping nodes.

[hyanwong]

There is a question about what to do with edges above local roots, I think.

[hyanwong]

Here's a hack:

def simplify_keeping_all_nodes(ts):
    """
    Run the Hudson algorithm to convert from an implicit to an explicit edge encoding,
    but keep all the nodes (like simplify(filter_nodes=False) if that existed)
    """
    # get the edges to keep
    ts2, node_map = ts.simplify(keep_unary=True, map_nodes=True)
    val, inverted_map = np.unique(node_map, return_index=True)
    inverted_map = inverted_map[val != tskit.NULL]
    # only use the edges in the simplified one, but keep the nodes from the original
    tables = ts.dump_tables()
    tables.edges.clear()
    for edge in ts2.tables.edges:
        tables.edges.append(
            edge.replace(
                child=inverted_map[edge.child], parent=inverted_map[edge.parent]
            )
        )
    tables.sort()
    return tables.tree_sequence()

[hyanwong]

Here's some code I've put together to remove the edges above local roots, but leave the nodes untouched. it's pretty intricate though.

def remove_edges_above_local_roots(ts):
    # use simplify() to find the roots
    ts_simp, node_map = ts.simplify(map_nodes=True)
    inverted_node_map = {n: i for i, n in enumerate(node_map)}
    edges_to_prune = {}
    for intvl, orig_tree, tree in ts.coiterate(ts_simp):
        for root in tree.roots:
            orig_node = inverted_node_map[root]
            if orig_tree.parent(orig_node) != tskit.NULL:
                edge_above = orig_tree.edge(orig_node)
                if edge_above in edges_to_prune:
                    intervals = edges_to_prune[edge_above]
                    if intervals[-1].right == intvl.left:
                        intervals[-1] = tskit.Interval(intervals[-1].left, intvl.right)
                    else:
                        intervals.append(intvl)
                else:
                    edges_to_prune[edge_above] = [intvl]
    # Now split the edges to prune by their intervals
    tables = ts.dump_tables()
    edges = tables.edges
    for edge_id, intervals in edges_to_prune.items():
        deleted_edge = tables.edges[edge_id]
        if deleted_edge.left < intervals[0].left:
            edges.append(deleted_edge.replace(right=intervals[0].left))
        for i in range(len(intervals)):
            try:
                edges.append(deleted_edge.replace(
                    left=intervals[i].right, right=intervals[i+1].left))
            except IndexError:
                # last interval to delete
                if deleted_edge.right > intervals[i].right:
                    edges.append(deleted_edge.replace(left=intervals[i].right))
    keep = np.ones(edges.num_rows, dtype=bool)
    keep[list(edges_to_prune.keys())] = False
    tables.edges.replace_with(edges[keep])
    tables.sort()
    return tables.tree_sequence()

[jeromekelleher]

Actually implementing filter_nodes=False probably wouldn't be that hard.

[hyanwong]

A thought I had yesterday: there are really two things we do in simplify: change edges, and change nodes. If we were starting again from scratch, I might suggest just 2 parameters to simplify, e.g. keep_nodes="nonancestral" | "unary" | other_bitfields. (default=0) and keep_edge_spans="for_kept_nodes" | "above_local_roots"| .... Probably too late to change any of this now, but the principle might help to think of how to rationalise the simplify parameter options.

[jeromekelleher]

We went through this recently - we can't change the C API so we're stuck with adding these options in piecewise.

[hyanwong]

Right, and presumably we want to reflect the C API in the Python one, to avoid confusing people. You could imagine a table of parameter combos in the simplification tutorial, with nodes_removed on the X and edges_trimmed on the Y.

[jeromekelleher]

Not just that, it would actually be quite hard to implement a different Python API on top of the C API structured in the way it currently is. It's not worth the time and effort.

[hyanwong]

Yeah, I'm not actively promoting the idea of changing the API, but I think trying to classify the various operations in some logical way in our heads will help us when trying to explain in the docs what the various options do.