Improved make_ancestors_ts function

[hyanwong]

In tests/tsutil.py, we have a function called "make_ancestors_ts". This is one way in which we can take an existing tree sequence (say from the unified genealogy) and match a new sample into it, without having the old ancestors_ts lying around. Using a (simplified) final tree sequence isn't quite as good as using the original, as various non-coalescent (unary or even non-ancestral) portions of haplotypes will have been lost, but it's probably not too bad. Also note that Jerome is also working on methods to expose the LS algorithm more generally, so the approach below could be obsoleted by his improvements.

The current method allows matching against a tree sequence in which (optionally) all the nodes at time 0 have been removed: these are assumed to be the samples. Generalise this, and allow matching a new sample (possibly at a non-zero time), we could do with 2 things:

1. A method to remove all nodes younger than time t (i.e. not just t=0)
2. A node mapping between the nodes in the original TS and the node IDs in the new ancestors TS.

Below is a function that will do this:

def make_ancestors_ts(ts, remove_recent=None, epsilon=None, remove_sites=None):
    """
    Return a tree sequence suitable for use as an ancestors tree sequence from the
    specified source tree sequence, together with a node map that for each node in the
    input tree sequence, specifies a node ID in the output tree sequence (or tskit.NULL
    if a node has been removed from the output)
    
    If a non-negative remove_recent value is given, remove any nodes that are that age
    or younger (e.g., to remove only the contemporary nodes, specify remove_recent=0).

    Epsilon is an amount added to times in the resulting tree sequence so that
    new samples at time 0 can be added to the resulting ancestors_ts even if nodes at
    time 0 were retained. As such it is not usually needed unless remove_recent
    is None, in which case a positive value is required.

    ``remove_sites`` is a list specifying any extra sites to remove as inference
    sites from the ancestors_ts. Additional sites will be added to this list if e.g. they 
    correspond to multi-allelic sites or singletons.

    We generally assume that this is a standard tree sequence, as output by
    msprime.simulate. We remove populations, as normally ancestors tree sequences
    do not have populations defined.
    """
    if epsilon is None:
        epsilon = 0
    if epsilon < 0:
        raise ValueError("Epsilon cannot be negative")
    
    # Get the non-singleton sites and those with > 1 mutation
    remove_sites = [] if remove_sites is None else list(remove_sites)
    for tree in ts.trees():
        for site in tree.sites():
            if len(site.mutations) != 1:
                remove_sites.append(site.id)
            else:
                if tree.num_samples(site.mutations[0].node) < 2:
                    remove_sites.append(site.id)

    reduced = ts.delete_sites(np.unique(remove_sites))
    minimised, node_map = reduced.simplify(
        reduce_to_site_topology=True,
        filter_sites=False,
        filter_individuals=False,
        filter_populations=False,
        keep_unary=True,
        map_nodes=True,
    )

    tables = minimised.dump_tables()
    # Rewrite the nodes so that 0 is one older than all the other nodes.
    nodes = tables.nodes.copy()
    tables.populations.clear()
    tables.nodes.clear()
    tables.nodes.add_row(flags=1, time=np.max(nodes.time) + 1)
    tables.nodes.append_columns(
        flags=np.ones_like(nodes.flags),  # Everything is a sample
        time=nodes.time,
        individual=nodes.individual,
        metadata=nodes.metadata,
        metadata_offset=nodes.metadata_offset,
    )
    # Add one to all node references to account for this.
    tables.edges.set_columns(
        left=tables.edges.left,
        right=tables.edges.right,
        parent=tables.edges.parent + 1,
        child=tables.edges.child + 1,
    )
    tables.mutations.node += 1
    tables.mutations.time = np.full_like(tables.mutations.time, tskit.UNKNOWN_TIME)

    trees = minimised.trees()
    tree = next(trees)
    left = 0
    # To simplify things a bit we assume that there's one root. This can
    # violated if we've got no sites at the end of the sequence and get
    # n roots instead.
    root = tree.root
    for tree in trees:
        if tree.root != root:
            tables.edges.add_row(left, tree.interval[0], 0, root + 1)
            root = tree.root
            left = tree.interval[0]
    tables.edges.add_row(left, ts.sequence_length, 0, root + 1)
    tables.sort()

    # After we inserted a new node at index 0, we must shift the original
    # node_map to match the tables numbering.  BUT only shift non-removed nodes.
    # node_map came from minimised.simplify(original_ts, map_nodes=True).
    # It maps original node ids -> minimised node ids (or -1 if removed).
    node_map = np.where(node_map == -1, -1, node_map + 1).astype(np.int32)
    
    if remove_recent is None:
        if epsilon <= 0 and np.any(ts.nodes_time <= 0):
            raise ValueError("If any nodes are at time 0, you must specify a positive epsilon value")
    else:
        if remove_recent < 0:
            raise ValueError("Cannot set a negative remove_recent value")
        samples = np.where(tables.nodes.time > remove_recent)[0].astype(np.int32)
        new_node_map = tables.simplify(samples=samples, filter_sites=False, keep_unary=True)
        # Build final node_map: for entries that were removed earlier (node_map == -1)
        # keep -1; for others, map through new_node_map.
        final_map = np.full_like(node_map, -1, dtype=np.int32)
        alive = node_map != -1
        final_map[alive] = new_node_map[node_map[alive]]
        node_map = final_map
    return tables.tree_sequence(), node_map

Here's an outline for a demo matching an ancient sample into an inferred and dated ts. The last step will need fixing, which I'm working on now:

# First simulate a TS with some ancient samples
import stdpopsim

species = stdpopsim.get_species('HomSap')
model = species.get_demographic_model('PapuansOutOfAfrica_10J19')
contig = species.get_contig('1', right=1000000, mutation_rate=model.mutation_rate)

engine = stdpopsim.get_engine('msprime')
samples = {"YRI": 200, "DenA": 2, "NeaA": 2}

sim_ts = engine.simulate(model, contig, samples=samples, seed=123)

# Now infer just with the recents

import tsinfer
import tsdate
import numpy as np

# Input from the old SampleData class, until https://github.com/tskit-dev/tsinfer/issues/991 is fixed
sd = tsinfer.SampleData.from_tree_sequence(sim_ts)
sample_times = sim_ts.nodes_time[sim_ts.samples()]

ancestors = tsinfer.generate_ancestors(sd, progress_monitor=True)
a_ts = tsinfer.match_ancestors(sd, ancestors, progress_monitor=True)
inferred_ts = tsinfer.match_samples(sd, a_ts, indexes = np.where(sample_times == 0)[0], progress_monitor=True)
# Annoyingly, to get better dates we probably want to simplify, which will remove useful non-coalescent regions
post_ts = tsdate.preprocess_ts(inferred_ts.simplify(filter_populations=False))
dated_ts = tsdate.date(post_ts, mutation_rate=model.mutation_rate, progress=True)

# Now make an ancestors_ts from the result, and match the NeaA back in:
# Some sites in the new_ats have more than 2 alleles in the sd file (e.g. if they have a new
# mutation in one of the historical samples. We need to remove these
exclude_positions = sd.sites_position[:][sd.num_alleles() > 2]
new_ats, node_map = make_ancestors_ts(
    dated_ts,
    remove_recent=0,
    remove_sites=np.where(np.isin(dated_ts.sites_position, exclude_positions))[0],
)
rev_map = -np.ones_like(node_map, shape=new_ats.num_nodes)
kept = node_map != tskit.NULL
rev_map[node_map[kept]] = np.arange(len(node_map))[kept]
valid = rev_map >= 0  # there is an extra node in new_ats not present in the original
assert np.all(new_ats.nodes_time[valid] == dated_ts.nodes_time[rev_map][valid])

# Get the mapping of the new samples
tmp_ts = tsinfer.match_samples(
    sd,
    new_ats,
    indexes = np.where(sample_times != 0)[0],
    progress_monitor=True,
    path_compression=False,  # Make sure that we only have edges from the new samples to existing nodes
    post_process=False,
)

# check that match_samples has just added nodes to the end of new_ats
assert np.all(tmp_ts.nodes_time[np.arange(new_ats.num_nodes)] == new_ats.nodes_time)

# add all the historical samples back into the dated tree sequence
# by using the edges from the newly matched tmp_ts
tables = dated_ts.dump_tables()
new_samples = tmp_ts.samples()
individuals_map = {}
new_node_map = {}
for i in np.unique(tmp_ts.nodes_individual[new_samples]):
    individuals_map[i] = tables.individuals.append(tmp_ts.individual(i))
for u in new_samples:
    sample_node = tmp_ts.node(u)
    new_node_map[u] = tables.nodes.append(sample_node.replace(individual=individuals_map[sample_node.individual]))
for e in np.where(np.isin(tmp_ts.edges_child, new_samples))[0]:
    edge = tmp_ts.edge(e)                      
    tables.edges.add_row(
        child=new_node_map[edge.child],
        parent=rev_map[edge.parent],
        right=edge.right,
        left=edge.left,
    )
tables.sort()
final_ts = tables.tree_sequence()

[gregorgorjanc]

@hyanwong the demo code is duplicated;) Thanks for sharing though;)

[hyanwong]

Oops, I was editing it and obviously messed up somehow. Now corrected

[hyanwong]

Here's a test displaying the topologies in the original simulation, and then when the historical samples are added in again:

from tqdm.auto import tqdm

from tqdm.auto import tqdm

def plot_topologies(ts, sample_set_dict=None, title=None, max_trees=None):
    # Print the proportion of different topologies, most common first
    node_labels = {}
    if sample_set_dict is None:
        population_map = {p.metadata["id"]: p.id for p in ts.populations()}
        sample_populations = list(sorted({ts.node(u).population for u in ts.samples()}))
        names = [ts.population(p).metadata["id"] for p in sample_populations]
        sample_sets = [ts.samples(population=p) for p in sample_populations]
    else:
        names = list(sample_set_dict.keys())
        sample_sets = list(sample_set_dict.values())
    
    topology_span = {tree.rank(): 0 for tree in tskit.all_trees(len(sample_sets))}
    for i, name in enumerate(names):
        node_labels[i] = name

    for topology_counter, tree in tqdm(
        zip(ts.count_topologies(sample_sets=sample_sets), ts.trees()),
        total=ts.num_trees,
    ):
        if tree.num_edges > 0:
            embedded_topologies = topology_counter[np.arange(len(sample_sets))]
            for rank, count in embedded_topologies.items():
                topology_span[rank] += count * tree.span
        if max_tree is not None and tree.index >= max_trees:
            break
            
    styles = ".sample text.lab {font-size: 0.7em;}"
    
    total = sum(topology_span.values())
    w = 160
    h = 150
    s = f'<svg version="1.1" width="{w*len(topology_span)}" height="{h + 20}" xmlns="http://www.w3.org/2000/svg">'
    if title is not None:
        s += f'<text x="{w*len(topology_span)/2}" y="{h}" text-anchor="middle">{title}</text>'

    # Order topologies by count
    for i, (rank, weight) in enumerate(sorted(topology_span.items(), key=lambda x: -x[1])):
        label = f"{weight/total *100:.1f}% of genome"
        embedded_tree = tskit.Tree.unrank(len(sample_sets), rank)
        s += embedded_tree.draw_svg(
            size=(w, h), root_svg_attributes={"x": w*i}, style="".join(styles), node_labels=node_labels, title=label
        )
    s += "</svg>"
    return tskit.drawing.SVGString(s)

display(plot_topologies(sim_ts, title="Simulated tree sequence topologies"))
display(plot_topologies(final_ts, title="Inferred tree sequence topologies, with historical samples added subsequently"))

[hyanwong]

Here's an example of remapping the Denisovan samples onto chr20 of the unified genealogy. The code is basically the same as above. Note that the existing unified genealogies can be improved upon (especially the dates and the number of recurrent sites), but we appear to get reasonable results. We have to stop the topology counting early because sometimes the neanderthals are matches as ancestors of the denisovan, and topology inference doesn't make sense with internal samples.

import tszip
import tsdate
import json
# Download hgdp_tgp_sgdp_high_cov_ancients_chr20_p.dated.trees.tsz from https://zenodo.org/records/5512994
ts = tszip.load("hgdp_tgp_sgdp_high_cov_ancients_chr20_p.dated.trees.tsz")

sd = tsinfer.SampleData.from_tree_sequence(ts)
sample_times = ts.nodes_time[ts.samples()]

den_time = 2556
den_samples = []
for u in ts.samples():
    if ts.nodes_time[u] == den_time:  # denisovan time
        den_samples.append(u)
        assert json.loads(ts.individual(ts.node(u).individual).metadata.decode())['name'] == "Denisova"

# Make a base_ts that has the Denisovan samples removed
base_ts = ts.simplify(np.setdiff1d(ts.samples(), den_samples))

# Now make an ancestors_ts from the base_ts, so we can try to match the Denisovan back in:
# Some sites in the have more than 2 alleles in the sd file (e.g. if they have a new
# mutation in one of the historical samples). We need to remove these
exclude_positions = sd.sites_position[:][sd.num_alleles() > 2]
new_ats, node_map = make_ancestors_ts(
    base_ts,
    remove_recent=0,
    remove_sites=np.where(np.isin(base_ts.sites_position, exclude_positions))[0],
)
rev_map = -np.ones_like(node_map, shape=new_ats.num_nodes)
kept = node_map != tskit.NULL
rev_map[node_map[kept]] = np.arange(len(node_map))[kept]
valid = rev_map >= 0  # there is an extra node in tmp_ats not present in the original
assert np.all(new_ats.nodes_time[valid] == base_ts.nodes_time[rev_map][valid])


# Get the mapping of the new samples
tmp_ts = tsinfer.match_samples(
    sd,
    new_ats,
    indexes = den_samples,
    progress_monitor=True,
    path_compression=False,  # Make sure that we only have edges from the new samples to existing nodes
    post_process=False,
)


# check that match_samples has just added nodes to the end of tmp_ats
assert np.all(tmp_ts.nodes_time[np.arange(new_ats.num_nodes)] == new_ats.nodes_time)

# add all the historical samples back into the base tree sequence
# by using the edges from the newly matched tmp_ts
tables = base_ts.dump_tables()
new_samples = tmp_ts.samples()
individuals_map = {}
new_node_map = {}
for i in np.unique(tmp_ts.nodes_individual[new_samples]):
    individuals_map[i] = tables.individuals.append(tmp_ts.individual(i))
for u in new_samples:
    sample_node = tmp_ts.node(u)
    new_node_map[u] = tables.nodes.append(sample_node.replace(individual=individuals_map[sample_node.individual]))
for e in np.where(np.isin(tmp_ts.edges_child, new_samples))[0]:
    edge = tmp_ts.edge(e)                      
    tables.edges.add_row(
        child=new_node_map[edge.child],
        parent=rev_map[edge.parent],
        right=edge.right,
        left=edge.left,
    )
tables.sort()
final_ts = tables.tree_sequence()

And here are some topology visualisation (second row has had the Denisovans added by the process above)

def make_afr_eur_nea_den_sample_set_dict(ts):
    # map "AFR" => [all, african, sample, ids] etc. - this is specific to the metadata stored in 
    # the unified genealogies, which differs depending on whether the samples are from 1KGP, HGDP, etc.
    superpop = collections.defaultdict(list)
    for pop in ts.populations():
        md = json.loads(pop.metadata.decode())
        if md.get('region') == "AFRICA" or md.get('region') == "Africa" or md.get('super_population') == "AFR":
            superpop["AFR"] += list(ts.samples(pop.id))
        elif md.get('super_population') == "Max Planck":
            if md.get("name") == "Denisovan":
                superpop["DEN"] += list(ts.samples(pop.id))
            else:
                superpop["NEA"] += list(ts.samples(pop.id))
        elif md.get('region') == "EUROPE" or md.get('region') == "WestEurasia" or md.get('super_population') == "EUR":
            superpop["EUR"] += list(ts.samples(pop.id))

    return superpop


display(plot_topologies(
    ts,
    sample_set_dict=make_afr_eur_nea_den_sample_set_dict(ts),
    title="Unified genealogy: original tree seq",
    max_trees=10_000)
)

display(plot_topologies(
    final_ts,
    sample_set_dict=make_afr_eur_nea_den_sample_set_dict(final_ts),
    title="Unified genealogy: original tree seq",
    max_trees=10_000)
)