How to add individuals from multiple generations to an inference?

[hyanwong]

Several people, e.g. @minesrebollo and @lakishadavid have sampled genomes which can potentially descend from other samples. In some cases, we might even have a pedigree for these individuals. There is currently no recommended way to incorporate these "non-contemporaneous" samples into tsinfer. However, if we can assume that the multigenerations in the pedigree are all more recent than the age of any of the MRCAs associated a derived allele, so we can probably insert the historical samples between time 0 and time 1/num_samples, ordered by their date of birth.

It would be really interesting to see if this approach ends up producing a relatively accurate summary of the known pedigree. I guess that’s something we can test relatively easily both by simulation and, more excitingly, from real data.

I've opened up this discussion so that we can thrash out the details here, and then (hopefully) turn the recommendations into some sort of tutorial.

[hyanwong]

The method I was suggesting would involve using the "time" argument to the SampleData.add_individual function when constructing the sample data file, then do ancestors = tsinfer.generate_ancestors(), then do new_ancestors = ancestors.insert_proxy_samples(), and use the new_ancestors for inference. This inserts one proxy sample into the ancestors file for each sample which is associated with a non-zero-time individual.

Note that the insert_proxy_samples method was intended for cases where we were relatively sure that these "proxy samples" were not actually ancestors, and we would then match the true samples (including the non-contemporaneous ones) later against the tree sequence using tsinfer.match_samples(..., force_sample_times=True). In the case of known ancestors we probably want to say that the proxy samples actually are the true samples, so we'd need to think a little how to do this.

[hyanwong]

I was going to test this with data simulated using the (new, undocumented, due-to-be-changed) msprime pedigree simulator (e.g. see a demo at tskit-dev/msprime#1865), but currently this doesn't incorporate internal (non-contemporaneous) samples, see tskit-dev/msprime#1855. It would be nice to address that issue so that we can test this all easily. I don't think there's another way to use msprime to produce internal samples, so for the moment we could try to use SLiM to create some test cases.

[hyanwong]

Here's some code to create a TS with samples at generations 0, 1, 2, and 3 into the past, using the most recent (GitHub) version of SLiM, which saves pedigree info in the tree sequence. We have 5 individuals, and for one of the individuals, mark all its (pedigree) ancestors as samples for the previous 3 generations. It also uses msprime to recapitate the TS. We should then be able to throw down mutations to create some samples for inference, while still having the ground-truth TS:

import collections
import os
from string import Template
import subprocess
import tempfile

import msprime
import numpy as np
import pyslim
import tskit

# Define an eidos script that outputs a 100Mb tree sequence to ${ts_file}
eidos_script = Template("""
// Keywords: tree-sequence recording, tree sequence recording

initialize() {
    initializeTreeSeq();
    initializeMutationRate(0);
    initializeMutationType("m1", 0.5, "f", 0.0);
    initializeGenomicElementType("g1", m1, 1.0);
    initializeGenomicElement(g1, 0, 1e8-1);
    initializeRecombinationRate(1e-8);
}
1 {
    sim.addSubpop("p1", 50);
}
1:3 early() {
    sim.treeSeqRememberIndividuals(p1.individuals);
}
3 late() {
    sim.treeSeqOutput("${ts_file}");
}
""")

random_seed = 12345

# Save and load the tree sequence using a temporary directory
with tempfile.TemporaryDirectory() as dir:
    file = os.path.join(dir, "slim_output.trees")
    script = eidos_script.substitute(ts_file=file).encode()
    # Run SLiM, then load the generated tree sequence
    subprocess.run(["slim", "-s", str(random_seed)], input=script)
    ts = pyslim.load(file)

# mark 5 individuals (10 nodes) in the current generation as samples, and mark
# all the ancestral nodes of the last individual as samples
num_nodes = 10

flags = np.zeros_like(ts.tables.nodes.flags)
for i in ts.individuals():
    for n in i.nodes:
        if ts.node(n).time == 0:
            last_individual = i
            flags[n] |= tskit.NODE_IS_SAMPLE
    if np.sum(flags != 0) == num_nodes:
        print()
        break

def mark_parents_as_samples(individual):
    for parent_ind_id in individual.parents:
        if parent_ind_id >= 0:
            parent_ind = ts.individual(parent_ind_id)
            for n in parent_ind.nodes:
                flags[n] |= tskit.NODE_IS_SAMPLE
            mark_parents_as_samples(parent_ind)

mark_parents_as_samples(last_individual)

tables = ts.dump_tables()
tables.nodes.flags = flags
node_map = tables.simplify()
tables.sort_individuals()
simplified_ts = tables.tree_sequence()

sample_nodes_at_time = collections.defaultdict(int)
for u in simplified_ts.samples():
    sample_nodes_at_time[simplified_ts.node(u).time] += 1
# Check there are 2 parents (4 nodes) at time 1.0, 4 parents (8 nodes) at time 2.0, etc
print(sample_nodes_at_time)

# recapitate
N = 1000  # use a larger population for recapitating
final_ts = msprime.sim_ancestry(
    initial_state=simplified_ts, population_size=N, model="dtwf", random_seed=123456, 
)

If in a Jupyter notebook, you can print out e.g. the first tree to check:

from IPython.display import display, SVG
display(SVG(final_ts.draw_svg(
    size=(800, 800),
    x_lim=[0, 10],
    time_scale="log_time",
    y_axis=True,
    style=",".join([f".n{node_map[n]} .sym" for n in last_individual.nodes]) + "{fill: red}"
)))

[hyanwong]

Here's some demo code for the inference, using the final_ts constructed above, although it would be good to solve #605 to do this properly, so we can test how well the pedigree links have been inferred.

import json

import msprime
import tsinfer


mut_ts = msprime.sim_mutations(final_ts, rate=1e-8)

# Make the sampledata file

max_sample_time = max(mut_ts.node(u).time for u in mut_ts.samples())
min_freq = 1/mut_ts.num_samples # we never build an actual ancestor at min_freq, because it would be a singleton

ploidy = 2

with tsinfer.SampleData(
    path="test.samples", sequence_length=mut_ts.sequence_length
) as samples:
    for p in mut_ts.populations():
        samples.add_population(metadata=p.metadata)
    sampled_individuals = np.unique(mut_ts.tables.nodes.individual[mut_ts.samples()])
    for i in sampled_individuals:
        indiv = mut_ts.individual(i)
        assert len(indiv.nodes) == 2
        nd1, nd2 = (mut_ts.node(u) for u in indiv.nodes)
        assert nd1.time == nd2.time
        assert nd1.population == nd2.population

        # Here's the hack to place recent pedigree ancestors individualsPlace the time of the individual at a fraction of the smallest time
        t = nd1.time / max_sample_time  # 0 <= t <= 1
        t *= min_freq  # 0 <= t <= min_freq
        samples.add_individual(
            ploidy=ploidy,
            population=nd1.population,
            metadata=indiv.metadata,
            time=t,
        )

    # the order of the sample nodes in the TS may not match the order of the individuals
    # so we must ensure the right order
    node_order = np.array([u for i in sampled_individuals for u in mut_ts.individual(i).nodes])
    for v in mut_ts.variants():
        samples.add_site(
            v.position,
            genotypes=v.genotypes[node_order],
            alleles=v.alleles,
        )
        
assert samples.num_samples == mut_ts.num_samples
assert np.any(samples.individuals_time[:] > 0)

# Make the ancestors for inference, adding in the historical samples as proxy ancestors

ancestors = tsinfer.generate_ancestors(samples)
ancestors_with_proxies = ancestors.insert_proxy_samples(samples, epsilon=0)
# Hack to identify the sample ancestors until https://github.com/tskit-dev/tsinfer/issues/605 is solved
proxy_ancestors = np.where(ancestors_with_proxies.ancestors_time[:] <= min_freq)[0]
num_proxy_ancestors = ancestors_with_proxies.num_ancestors - ancestors.num_ancestors
assert num_proxy_ancestors == len(proxy_ancestors)
assert num_proxy_ancestors == np.sum(samples.individuals_time[:] > 0) * ploidy

ancestor_ts = tsinfer.match_ancestors(samples, ancestors_with_proxies)

# only match the contemporaneous samples, and don't simplify, so we don't delete the historical nodes
contempory_samples = samples.subset(individuals=np.where(samples.individuals_time[:] == 0)[0])
inferred_ts = tsinfer.match_samples(contempory_samples, ancestor_ts, simplify=False)

# Flag up the proxy nodes as samples. Note we have lost the individuals information here, though
tables = inferred_ts.dump_tables()
flags = tables.nodes.flags
for node in inferred_ts.nodes():
    if node.metadata:
        md = json.loads(node.metadata)
        if 'ancestor_data_id' in md:
            if md['ancestor_data_id'] in proxy_ancestors:
                flags[node.id] |= tskit.NODE_IS_SAMPLE
tables.nodes.flags = flags
final_inferred_ts = tables.tree_sequence().simplify()
print(
    "Final inferred ts has",
    final_inferred_ts.num_samples,
    "samples,",
    inferred_ts.num_samples,
    "from match_samples and",
    num_proxy_ancestors,
    "from historical samples added as ancestors."
)

[gregorgorjanc]

Thanks for this @hyanwong !!!

[AprilYUZhang]

Hi @hyanwong, I am building ARG based on the inferred tree (based on founders) and segregation probability (progeny). However, I have some challenges in dating the nodes, as they are on different timescales.

I notice there is the following formula to rescale time in your code.

t = nd1.time / max_sample_time
t *= min_freq

I was wondering if there is a specific principle or rationale behind this approach?

Additionally, if I infer the ARG using founders and aim to integrate the edges and nodes of progenies based on segregation probabilities, would you have any suggestions on how to appropriately rescale the time to ensure consistency?
I built a tree like this:

In this tree, Nodes 4 and 5 are ancestor nodes, while Nodes 0, 1, 2, and 3 represent the founder nodes (I inferred ancestral ARG based on founder nodes). The remaining nodes are progeny.

Any insights or guidance you could provide would be greatly appreciated. Thank you!

[gregorgorjanc]

@AprilYUZhang infer tree sequence for pedigree founders and date it (tsinfer&tsdate). Then build tree sequence within pedigree. Having dated tree sequence for pedigree founders, the founders will be considered as samples and, I think, have time 0. Then count time for pedigree non-founders from founders forward-in-time (but decreasing the values since ts ecosystem assumes that time increases backward-in-time): -1, -2, …

We need to be mindful of time units though. I am not sure what time units do we get from, say, tsdate. Within pedigree we will often have date/year of birth. To combine the two tree sequences, we must use the same time unit in both of them.

As to Yan’s code, he does say in the comment that it’s hack

# Here's the hack to place recent pedigree ancestors individualsPlace the time of the individual at a fraction of the smallest time
t = nd1.time / max_sample_time  # 0 <= t <= 1
t *= min_freq  # 0 <= t <= min_freq

First code line rescales time to [0,1]. Second rescales it to [0,min_freq]. The two constants come from above in the code

max_sample_time = max(mut_ts.node(u).time for u in mut_ts.samples())
min_freq = 1/mut_ts.num_samples # we never build an actual ancestor at min_freq, because it would be a singleton

this is based on the assumption stated at the start: …if we can assume that the multigenerations in the pedigree are all more recent than the age of any of the MRCAs associated a derived allele, so we can probably insert the historical samples between time 0 and time 1/num_samples, ordered by their date of birth.