Optimizing ancestry fraction calculations from SLiM tree sequences

[Xiaohe Tian]

Hi everyone,

I am working on a project involving a two-way admixture population simulation. I am dumping tree sequences at every generation using sim.treeSeqOutput() and then, in a post-processing step, calculating the ancestry fraction for individuals (i.e., the proportion of their genome that traces back to a specific ancestral population at a reference time t_ref).

I have implemented the following Python script to compute these fractions for batches of individuals, parallelized by individual ID ranges:

def ancestry_fracs_subset(ts, t_ref, ind_ids, pop1=1):
    """
    Compute ancestry fraction (fraction of genome tracing to pop1 at time t_ref)
    for a subset of individuals in a tree sequence.
    Assumes diploid individuals (2 nodes per individual).
    """
    pops  = ts.tables.nodes.population
    times = ts.tables.nodes.time

    # Build haplotype node list + owner mapping (avoid allocating node_to_idx for all nodes)
    nodes = []
    owners = []
    for j, ind_id in enumerate(ind_ids):
        ind = ts.individual(ind_id)
        for n in ind.nodes:
            nodes.append(n)
            owners.append(j)

    nodes = np.asarray(nodes, dtype=np.int32)
    owners = np.asarray(owners, dtype=np.int32)

    p1_len = np.zeros(len(ind_ids), dtype=np.float64)
    tot_len = np.zeros(len(ind_ids), dtype=np.float64)

    for tree in ts.trees():
        span = tree.span
        for k, u in enumerate(nodes):
            v = u
            p = tree.parent(v)

            # climb until we reach the first ancestor at/older than t_ref
            while p != tskit.NULL and times[p] < t_ref:
                v = p
                p = tree.parent(v)

            anc = p if p != tskit.NULL else v

            j = owners[k]
            tot_len[j] += span
            if pops[anc] == pop1:
                p1_len[j] += span

    return p1_len / tot_len

Currently, my workflow involves iterating through all trees in the tree sequence for each batch of individuals to climb the tree until t_ref is reached. Given the size of my simulations, this is becoming computationally expensive. 

I have a few questions regarding potential optimizations: 

1. Traversal Efficiency: Is the manual while loop (climbing the tree for every individual in every tree span) the most efficient way to query tskit for these ancestry fractions? Are there tskit methods or higher-level abstractions that could bypass the tree-by-tree traversal? 
2. Memory/IO: Since I have already dumped the full tree sequences, is there a better way to structure this calculation?  
3. Data Structure: Are there any known patterns for querying ancestry in admixed populations within tree sequences that I am missing?

 Any advice on optimizing this traversal or alternative approaches for calculating these ancestry fractions would be greatly appreciated. 

Thanks!

[Gregor Gorjanc]

Would this help https://tskit.dev/software/tspop.html

With regards,
  Gregor

--
SLiM forward genetic simulation: http://messerlab.org/slim/
---
You received this message because you are subscribed to the Google Groups "slim-discuss" group.
To unsubscribe from this group and stop receiving emails from it, send an email to slim-discuss...@googlegroups.com.
To view this discussion visit https://groups.google.com/d/msgid/slim-discuss/a79442b1-d030-41e7-93ab-20b9973eeab3n%40googlegroups.com.

[Hanbin Lee]

I also used tspop in my paper. https://pubmed.ncbi.nlm.nih.gov/40646547/