merging multiple tree sequences into a single demographic model

[mshpak76]

I recently posted this on the slim discussion board, but realized that this may be a more appropriate venue since the post-processing involves tskit:

'm modeling independent selective sweeps in three subpopulations with genealogy {1,{2,3}} and conditioning on fixation of the beneficial allele in all three. Because this is done by rerunning the simulation each time the allele is lost, if the expected number of iterations for a single population is K, then it will be K^3 when conditioning on all three.

To get around this problem, I simulate three different tree sequences which I need to piece together into a single genealogy (followed by recapitation and mapping of neutral mutations). The difficulty I'm having is with the first step. Given three tree sequences ts1, ts2, ts3, how does one create a single ts consistent a specified demographic history?

The code below is what I would like to do, given the appropriate SOME_MERGE_FUNCTION(ts1,ts2,ts3) that would give me the merged ts. It seems that union() in tskit can do something like this, but I can't find an example of the correct syntax to implement what I need using union(). Is there another method that would work?

ts1 = tskit.load("T1.trees")
ts2 = tskit.load("T2.trees")
ts3 = tskit.load("T2.trees")

ts = SOME_MERGE_FUNCTION(ts1,ts2,ts3) ???

ts_new = pyslim.SlimTreeSequence(ts)

demog = msprime.Demography()
demog.add_population(name = "1", initial_size = 20000, initially_active=True)
demog.add_population(name = "2", initial_size = 10000, initially_active=True)
demog.add_population(name = "3", initial_size = 10000)
demography.add_population_parameters_change(time=1000, population="2", initial_size=5000)
demog.add_population_split(time=1000, derived=["3"], ancestral="2")
demog.add_population_split(time=2000, derived=["2"], ancestral="1")

rts = pyslim.recapitate(ts_new, recombination_rate = 1.14e-8, gene_conversion_rate = 5.9e-8, gene_conversion_tract_length = 518, ancestral_Ne=20000, random_seed = ancestry_seed)

mts = msprime.mutate(rts, rate = 1.22e-8, keep=True, random_seed = mutation_seed, model=msprime.SLiMMutationModel(type=0))

I found what seems to be a relevant example of parallelizing slim runs and applying the union() function here, but unfortunately it's not clear to me how to apply the union() function in my specific example:

https://tskit.dev/pyslim/docs/latest/vignette_parallel_phylo.html

(i.e. namely, what is the appropriate syntax for implementating union on my ts1...ts3)? I cannot use e.g. ts = ts1.union(ts2...) because in my scenario there are no shared nodes when ts1...ts3 are generated - I wish to create the shared nodes after the fact with recapitation. Is there some way to create a ts combining these when shared_nodes = None?

[petrelharp]

Copying from the SLiM-discuss list: From [https://tskit.dev/tskit/docs/stable/python-api.html#tskit.TreeSequence.union](the documentation),
"the node_mapping argument should be an array of length equal to the number of nodes in other and whose entries are the ID of the matching node in self, or tskit.NULL if there is no matching node. ". So, in your case you have no shared nodes, i.e., all entries of the node map should be tskit.NULL. So, this works, for instance:

  ts1.union(ts2, [tskit.NULL for _ in ts2.nodes()])