Simplifying when keeping unary nodes but excluding input roots

[hyanwong]

The simplify docs say:

keep_unary (bool) – If True, preserve unary nodes (i.e., nodes with exactly one child) that exist on the path from samples to root. (Default: False)

and also

keep_input_roots (bool) – Whether to retain history ancestral to the MRCA of the samples. If False, no topology older than the MRCAs of the samples will be included. If True the roots of all trees in the returned tree sequence will be the same roots as in the original tree sequence. (Default: False)

However, on testing the interacting between these, it seems as if, contrary to the docs, even when keep_input_roots=False, topology older than the MRCAs of the samples is still included if keep_unary=True:

ts = tskit.Tree.generate_comb (10).tree_sequence
ts = ts.simplify([6, 7, 8, 9], keep_input_roots=True)
ts = ts.simplify(keep_unary=True, keep_input_roots=False)  # I think this should remove node 7, but it doesn't
ts.draw_svg()

This is either a bug, or a documentation issue. If we keep the current behaviour and document it, then I am looking for a way to keep unary nodes except those above the local roots, and I can't think of an easy way to do that, without going through the trees one-by-one.

[petrelharp]

This isn't a bug, I'd say, because both of these options say "please keep these certain other things that would usually be removed". So I guess it's a doc bug for keep_input_roots; this perhaps should say "If False, no topology older than the MRCAs of the samples will be included, unless some other option overrides this."

[hyanwong]

OK, happy to alter the doc string. However, I am not sure what approach I can use to keep unary nodes except those above the root, which is required to use simplify() to subset a full ARG (see https://tskit.dev/tutorials/advanced_simplification.html#subsetting-a-full-arg)

[hyanwong]

Here's another interesting example that I just found: I want to "part simplify" a simulation, keeping haplotypes "extended", as in #2127. I though the code below might do it (keeping only the coalescent nodes, then resimplifying using keep_unary), but this can increase the number of edges over the standard simplify() because of keeping edges above the roots:

def simplify_keeping_part_coalescent(ts):
    simp, node_map = ts.simplify(focal, map_nodes=True)
    keep_nodes = np.where(node_map != tskit.NULL)[0]  #includes sample and coalescent nodes

    # Retain lots of nodes by treating them as focal, but don't change sample flags
    tmp_ts = ts.simplify(keep_nodes, update_sample_flags=False)
    # Now re-simplify, in case any coalescent nodes have unwanted ancestry above
    return tmp_ts.simplify(keep_unary=True)

Example output (note trees 5, 8, 10, and 12)