Removing mutations above the root, and changing ancestral state accordingly

[hyanwong]

Here's a useful recipe, in case there are mutations above the root, which can happen when you simplify to a subset of samples (see #260). @savitakartik and I hit this problem independently recently, when testing tsdate. I have tested with some msprime-generated tree sequences and it seems to work:

import numpy as np

def remove_root_mutations(ts):
    # Removes mutations above the root in a tree, and changes the ancestral state as required
    # assumes all trees have only one root.
    tables = ts.dump_tables()
    keep_muts = np.ones(tables.mutations.num_rows, dtype=bool)
    deleted_mut_parents = set()
    for tree in ts.trees():
        for site in tree.sites():
            root = tree.root
            root_muts = [mut for mut in site.mutations if mut.node == root]
            if len(root_muts) > 0:
                # Take the last mutation above the root
                anc_state = root_muts[-1].derived_state
                deleted_mut_parents.add(root_muts[-1].id)
                keep_muts[[mut.id for mut in root_muts]] = False
                tables.sites[site.id] = tables.sites[site.id].replace(ancestral_state=anc_state)
    # Some non-root mutations may have a parent that we intend to delete
    # We can safely set the parent of these mutations to NULL, as we have
    # now set the ancestral state to the state of that parent
    deleted_mut_parents = np.array(list(deleted_mut_parents))
    mut_parents = tables.mutations.parent
    mut_parents[np.isin(mut_parents, deleted_mut_parents)] = tskit.NULL
    tables.mutations.parent = mut_parents
    tables.mutations.keep_rows(keep_muts)
    return tables.tree_sequence()

E.g. test:

import msprime
ts = msprime.simulate(100, mutation_rate=4, length=100, recombination_rate=0.1, random_seed=2).simplify(range(5))
print(ts.num_mutations - remove_root_mutations(ts).num_mutations, "root mutations removed")
for v1, v2 in zip(ts.variants(), remove_root_mutations(ts).variants()):
    assert np.all(np.array(v1.alleles)[v1.genotypes] == np.array(v2.alleles)[v2.genotypes])