Repository: tskit-dev/tsinfer
Pull request #1012: Revise ancestor generation algorithm to improve performance
Author: login: Duncan-JR
Created: 2025-03-14T18:32:58Z | Merged: 2025-03-20T02:01:37Z

A major limitation of tsinfer 0.4 has been that ancestors with high frequency focal sites are excessively long. This reduces parallelism and wastes computation time in ancestor matching. We've explored a number of solutions (e.g. #911), but we've found a simple one that seems to work really well.

It's easiest to explain with an example:
![image](https://github.com/user-attachments/assets/404500ae-7570-4efb-be58-5d1ba5b20f7c)
Constructing an ancestor with focal site 5 and moving to the left, we start with a sample set of C-F. Since the focal AC is 4, we can only use sites with an AC of 5 or higher as inference sites. In the current implementation, we would ignore sites 1-4 entirely, but we are then missing an informative signal. Since B also carries the derived allele at site 3, and is the only other non-carrier at site 2, it seems that C has recombined into a clade with B and should be excluded from the sample set.

In the new approach, we still only insert a mutation into the ancestral haplotype if it is older than the focal site, but we use sites of all frequencies for determining when to exclude samples from the sample set. However, we only count conflicts at sites where there are carriers outside of the sample set (i.e. the derived AC in the full population > AC in the sample set).

Initial validation of the new approach looks promising. I ran a small out-of-Africa sim of 1mbp with 200 samples. Using the old validation code, I can match the inferred ancestors to simulated ones and compute `max(true_left - inferred_left, inferred_right - true_right)`, which I call the max overshoot. While the old algorithm increasingly overestimates the length of the ancestors as the focal frequency increases, the new algorithm does not. We don't seem to make many ancestors shorter than their true length either.

![image](https://github.com/user-attachments/assets/db59575e-cead-409d-934a-0233e0a348a4)

I've done some testing on larger simulations with similar results, but there is still a lot more validation to do, e.g. with mispolarised alleles and sequencing errors. I did run @hyanwong's code for analysing sample-to-root edges from #903, and in simulations we don't seem to see any change. However, there are a lot more sample-to-root edges in real data with the new method, which is fixed by using a small mismatch ratio for sample matching.

The new approach does require more work from the CPU for ancestor generation, but it seems negligible in comparison to the gains in parallelism and speed of ancestor matching. Comparing the methods on chr20q of the 1000 Genomes Project, we see that fewer ancestor groups are needed and there is a 5X decrease in CPU time needed for matching. The decrease in wall time is more modest: it seems that we aren't using the 126 threads of this EPYC system as effectively in the new approach.

![image](https://github.com/user-attachments/assets/fb1d55cc-56cf-488b-9992-ac53b339e955)

I've implemented the new algorithm in C and Python: the biggest change required is that the derived allele count of a site needs to be stored in the ancestor builder.



Diff:
diff --git a/lib/ancestor_builder.c b/lib/ancestor_builder.c
index 5af66d7c..16490d31 100644
--- a/lib/ancestor_builder.c
+++ b/lib/ancestor_builder.c
@@ -419,8 +419,9 @@ ancestor_builder_compute_ancestral_states(const ancestor_builder_t *self, int di
     tsk_id_t last_site = focal_site;
     int64_t l;
     tsk_id_t u;
-    size_t j, ones, zeros, tmp_size, sample_set_size, min_sample_set_size;
+    size_t j, ones, zeros, tmp_size, sample_set_size, min_sample_set_size, derived_count;
     double focal_site_time = self->sites[focal_site].time;
+    double site_time;
     const site_t *restrict sites = self->sites;
     const size_t num_sites = self->num_sites;
     allele_t consensus;
@@ -440,73 +441,78 @@ ancestor_builder_compute_ancestral_states(const ancestor_builder_t *self, int di
         /* printf("\tl = %d\n", (int) l); */
         ancestor[l] = 0;
         last_site = (tsk_id_t) l;
-        if (sites[l].time > focal_site_time) {
 
-            /* printf("\t%d\t%d:", (int) l, (int) sample_set_size); */
-            /* for (j = 0; j < sample_set_size; j++) { */
-            /*     printf("%d, ", sample_set[j]); */
-            /* } */
-            /* printf("\n"); */
-
-            ancestor_builder_get_site_genotypes_subset(
-                self, (tsk_id_t) l, sample_set, sample_set_size, genotypes);
-            ones = 0;
-            zeros = 0;
-            for (j = 0; j < sample_set_size; j++) {
-                switch (genotypes[j]) {
-                    case 0:
-                        zeros++;
-                        break;
-                    case 1:
-                        ones++;
-                        break;
-                }
+        /* printf("\t%d\t%d:", (int) l, (int) sample_set_size); */
+        /* for (j = 0; j < sample_set_size; j++) { */
+        /*     printf("%d, ", sample_set[j]); */
+        /* } */
+        /* printf("\n"); */
+
+        ancestor_builder_get_site_genotypes_subset(
+            self, (tsk_id_t) l, sample_set, sample_set_size, genotypes);
+        ones = 0;
+        zeros = 0;
+        for (j = 0; j < sample_set_size; j++) {
+            switch (genotypes[j]) {
+                case 0:
+                    zeros++;
+                    break;
+                case 1:
+                    ones++;
+                    break;
+            }
+        }
+        if (ones >= zeros) {
+            consensus = 1;
+        } else {
+            consensus = 0;
+        }
+        /* printf("\t:ones=%d, consensus=%d\n", (int) ones, consensus); */
+        /* fflush(stdout); */
+        for (j = 0; j < sample_set_size; j++) {
+            u = sample_set[j];
+            if (disagree[u] && (genotypes[j] != consensus)
+                && (genotypes[j] != TSK_MISSING_DATA)) {
+                /* This sample has disagreed with consensus twice in a row,
+                 * so remove it */
+                /* printf("\t\tremoving %d\n", sample_set[j]); */
+                sample_set[j] = -1;
             }
+        }
+
+        site_time = sites[l].time;
+        if (site_time > focal_site_time) {
             if (ones + zeros == 0) {
                 ancestor[l] = TSK_MISSING_DATA;
             } else {
-                if (ones >= zeros) {
-                    consensus = 1;
-                } else {
-                    consensus = 0;
-                }
-                /* printf("\t:ones=%d, consensus=%d\n", (int) ones, consensus); */
-                /* fflush(stdout); */
-                for (j = 0; j < sample_set_size; j++) {
-                    u = sample_set[j];
-                    if (disagree[u] && (genotypes[j] != consensus)
-                        && (genotypes[j] != TSK_MISSING_DATA)) {
-                        /* This sample has disagreed with consensus twice in a row,
-                         * so remove it */
-                        /* printf("\t\tremoving %d\n", sample_set[j]); */
-                        sample_set[j] = -1;
-                    }
-                }
                 ancestor[l] = consensus;
-                /* For the remaining samples, set the disagree flags based
-                 * on whether they agree with the consensus for this site. */
-                for (j = 0; j < sample_set_size; j++) {
-                    u = sample_set[j];
-                    if (u != -1) {
-                        disagree[u] = ((genotypes[j] != consensus)
-                                       && (genotypes[j] != TSK_MISSING_DATA));
-                    }
-                }
-                /* Repack the sample set */
-                tmp_size = 0;
-                for (j = 0; j < sample_set_size; j++) {
-                    if (sample_set[j] != -1) {
-                        sample_set[tmp_size] = sample_set[j];
-                        tmp_size++;
-                    }
-                }
-                sample_set_size = tmp_size;
-                if (sample_set_size <= min_sample_set_size) {
-                    /* printf("BREAK\n"); */
-                    break;
+            }
+        }
+        /* For the remaining samples, set the disagree flags based
+         * on whether they agree with the consensus for this site. */
+        derived_count = sites[l].derived_count;
+        if ((site_time > focal_site_time) || (derived_count > ones)) {
+            for (j = 0; j < sample_set_size; j++) {
+                u = sample_set[j];
+                if (u != -1) {
+                    disagree[u] = ((genotypes[j] != consensus)
+                                   && (genotypes[j] != TSK_MISSING_DATA));
                 }
             }
         }
+        /* Repack the sample set */
+        tmp_size = 0;
+        for (j = 0; j < sample_set_size; j++) {
+            if (sample_set[j] != -1) {
+                sample_set[tmp_size] = sample_set[j];
+                tmp_size++;
+            }
+        }
+        sample_set_size = tmp_size;
+        if (sample_set_size <= min_sample_set_size) {
+            /* printf("BREAK\n"); */
+            break;
+        }
     }
     *last_site_ret = last_site;
     return ret;
@@ -658,8 +664,16 @@ ancestor_builder_add_site(ancestor_builder_t *self, double time, allele_t *genot
     uint8_t *stored_genotypes = NULL;
     avl_tree_t *pattern_map;
     tsk_id_t site_id = (tsk_id_t) self->num_sites;
+    size_t derived_count, j;
     time_map_t *time_map = ancestor_builder_get_time_map(self, time);
 
+    derived_count = 0;
+    for (j = 0; j < (size_t) self->num_samples; j++) {
+        if (genotypes[j] == 1) {
+            derived_count++;
+        }
+    }
+
     if (time_map == NULL) {
         ret = TSI_ERR_NO_MEMORY;
         goto out;
@@ -676,6 +690,7 @@ ancestor_builder_add_site(ancestor_builder_t *self, double time, allele_t *genot
     pattern_map = &time_map->pattern_map;
     site = &self->sites[site_id];
     site->time = time;
+    site->derived_count = derived_count;
 
     search.encoded_genotypes = encoded_genotypes;
     search.encoded_genotypes_size = self->encoded_genotypes_size;
diff --git a/lib/tsinfer.h b/lib/tsinfer.h
index 628a5519..e6efbf12 100644
--- a/lib/tsinfer.h
+++ b/lib/tsinfer.h
@@ -69,6 +69,7 @@ typedef struct _node_segment_list_node_t {
 typedef struct {
     double time;
     uint8_t *encoded_genotypes;
+    tsk_size_t derived_count;
 } site_t;
 
 typedef struct {
diff --git a/tests/test_inference.py b/tests/test_inference.py
index ae415225..419945dd 100644
--- a/tests/test_inference.py
+++ b/tests/test_inference.py
@@ -3533,16 +3533,16 @@ def test_simple_case(self):
 
     def test_simulation_with_error(self):
         ts = msprime.simulate(
-            50, mutation_rate=10, random_seed=4, recombination_rate=15
+            50, mutation_rate=10, random_seed=5, recombination_rate=15
         )
-        ts = eval_util.insert_errors(ts, 0.2, seed=32)
+        ts = eval_util.insert_errors(ts, 0.2, seed=33)
         sample_data = tsinfer.SampleData.from_tree_sequence(ts, use_sites_time=False)
         self.verify(sample_data)
 
     def test_small_random_data(self):
         n = 25
         m = 20
-        G, positions = get_random_data_example(n, m)
+        G, positions = get_random_data_example(n, m, seed=101)
         with tsinfer.SampleData(sequence_length=m) as sample_data:
             for genotypes, position in zip(G, positions):
                 sample_data.add_site(position, genotypes)
@@ -3551,7 +3551,7 @@ def test_small_random_data(self):
     def test_large_random_data(self):
         n = 100
         m = 30
-        G, positions = get_random_data_example(n, m)
+        G, positions = get_random_data_example(n, m, seed=100)
         with tsinfer.SampleData(sequence_length=m) as sample_data:
             for genotypes, position in zip(G, positions):
                 sample_data.add_site(position, genotypes)
diff --git a/tsinfer/algorithm.py b/tsinfer/algorithm.py
index 7b08c133..5ad03724 100644
--- a/tsinfer/algorithm.py
+++ b/tsinfer/algorithm.py
@@ -57,6 +57,7 @@ class Site:
 
     id = attr.ib()
     time = attr.ib()
+    derived_count = attr.ib()
 
 
 class AncestorBuilder:
@@ -141,8 +142,9 @@ def add_site(self, time, genotypes):
         Adds a new site at the specified ID to the builder.
         """
         site_id = len(self.sites)
+        derived_count = np.sum(genotypes == 1)
         self.store_site_genotypes(site_id, genotypes)
-        self.sites.append(Site(site_id, time))
+        self.sites.append(Site(site_id, time, derived_count))
         sites_at_fixed_timepoint = self.time_map[time]
         # Sites with an identical variant distribution (i.e. with the same
         # genotypes.tobytes() value) and at the same time, are put into the same ancestor
@@ -213,38 +215,46 @@ def compute_ancestral_states(self, a, focal_site, sites):
         """
         focal_time = self.sites[focal_site].time
         g = self.get_site_genotypes(focal_site)
-        S = set(np.where(g == 1)[0])
-        # Break when we've lost half of S
-        min_sample_set_size = len(S) // 2
-        remove_buffer = []
+        sample_set = np.where(g == 1)[0]
+        min_sample_set_size = len(sample_set) // 2
         last_site = focal_site
-        # print("Focal site", focal_site, "time", focal_time)
+        disagree = np.zeros(self.num_samples, dtype=bool)
+
         for site_index in sites:
             a[site_index] = 0
             last_site = site_index
-            if self.sites[site_index].time > focal_time:
-                g_l = self.get_site_genotypes(site_index)
-                ones = sum(g_l[u] == 1 for u in S)
-                zeros = sum(g_l[u] == 0 for u in S)
-                # print("pos", site_index, ". Ones:", ones, ". Zeros:", zeros)
+            g_l = self.get_site_genotypes(site_index)
+            ones = sum(g_l[u] == 1 for u in sample_set)
+            zeros = sum(g_l[u] == 0 for u in sample_set)
+            consensus = 1 if ones >= zeros else 0
+            site_time = self.sites[site_index].time
+            derived_count = self.sites[site_index].derived_count
+
+            for j, u in enumerate(sample_set):
+                if (
+                    disagree[u]
+                    and (g_l[u] != consensus)
+                    and (g_l[u] != tskit.MISSING_DATA)
+                ):
+                    sample_set[j] = -1
+
+            if site_time > focal_time:
                 if ones + zeros == 0:
                     a[site_index] = tskit.MISSING_DATA
                 else:
-                    consensus = 1 if ones >= zeros else 0
-                    # print("\tP", site_index, "\t", len(S), ":ones=", ones, consensus)
-                    for u in remove_buffer:
-                        if g_l[u] != consensus and g_l[u] != tskit.MISSING_DATA:
-                            # print("\t\tremoving", u)
-                            S.remove(u)
                     a[site_index] = consensus
-                    # print("\t", len(S), remove_buffer, consensus, sep="\t")
-                    if len(S) <= min_sample_set_size:
-                        # print("BREAKING", len(S), min_sample_set_size)
-                        break
-                    remove_buffer.clear()
-                    for u in S:
-                        if g_l[u] != consensus and g_l[u] != tskit.MISSING_DATA:
-                            remove_buffer.append(u)
+
+            if (site_time > focal_time) or (derived_count > ones):
+                for u in sample_set:
+                    if u != -1:
+                        disagree[u] = (
+                            g_l[u] != consensus and g_l[u] != tskit.MISSING_DATA
+                        )
+
+            sample_set = sample_set[sample_set != -1]
+            if len(sample_set) <= min_sample_set_size:
+                break
+
         assert a[last_site] != tskit.MISSING_DATA
         return last_site
 


Reviews:
user: login: jeromekelleher
body: Generally looks great, but not obvious that we need the API changes to pass through derived count (definitely a premature optimisation if it can be computed easily).
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-18T13:15:57Z
user: login: Duncan-JR
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T14:24:16Z
user: login: jeromekelleher
body: LGTM!
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-19T08:35:15Z

Review comments:
diff_hunk: @@ -676,6 +683,7 @@ ancestor_builder_add_site(ancestor_builder_t *self, double time, allele_t *genot
     pattern_map = &time_map->pattern_map;
     site = &self->sites[site_id];
     site->time = time;
+    site->derived_count = derived_count;
path: lib/ancestor_builder.c
user: login: jeromekelleher
body: Why are we passing in the derived count? Is this the same as the number of entries in the genotypes > 0? 

If so, we should just compute it there, and leave the signatures alone. Lets not worry about one extra iteration over the genotypes at  site.
created_at: 2025-03-18T13:13:59Z
author_association: MEMBER
diff_hunk: @@ -676,6 +683,7 @@ ancestor_builder_add_site(ancestor_builder_t *self, double time, allele_t *genot
     pattern_map = &time_map->pattern_map;
     site = &self->sites[site_id];
     site->time = time;
+    site->derived_count = derived_count;
path: lib/ancestor_builder.c
user: login: Duncan-JR
body: Fair point: I've corrected it so that the derived count calculation happens within `add_site` for each ancestor builder
created_at: 2025-03-18T14:24:16Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov[bot]
created_at: 2025-03-14T18:47:52Z
body: ## [Codecov](https://app.codecov.io/gh/tskit-dev/tsinfer/pull/1012?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 93.41%. Comparing base [(`ca2cf16`)](https://app.codecov.io/gh/tskit-dev/tsinfer/commit/ca2cf165de953a49f953662f8edba36b137389f5?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev) to head [(`40196b3`)](https://app.codecov.io/gh/tskit-dev/tsinfer/commit/40196b38e0712fab67ac9123622cd2a5c46be6d3?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev).
> Report is 2 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #1012      +/-   ##
==========================================
+ Coverage   93.39%   93.41%   +0.01%     
==========================================
  Files          18       18              
  Lines        6483     6496      +13     
  Branches     1103     1107       +4     
==========================================
+ Hits         6055     6068      +13     
  Misses        291      291              
  Partials      137      137              
```

| [Flag](https://app.codecov.io/gh/tskit-dev/tsinfer/pull/1012/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev) | Coverage Δ | |
|---|---|---|
| [C](https://app.codecov.io/gh/tskit-dev/tsinfer/pull/1012/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev) | `93.41% <100.00%> (+0.01%)` | :arrow_up: |
| [python](https://app.codecov.io/gh/tskit-dev/tsinfer/pull/1012/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev) | `95.78% <100.00%> (?)` | |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/tskit-dev/tsinfer/pull/1012?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=tskit-dev).

<details><summary>🚀 New features to boost your workflow: </summary>

- ❄ [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE
user: login: jeromekelleher
created_at: 2025-03-17T10:02:21Z
body: Not sure why you're touching _tsinfermodule.c here @duncanMR, but that needs to be dropped from the diff. Your editor probably got overzealous there.
author_association: MEMBER
user: login: Duncan-JR
created_at: 2025-03-17T10:38:44Z
body: > Not sure why you're touching _tsinfermodule.c here @duncanMR, but that needs to be dropped from the diff. Your editor probably got overzealous there.

Apologies, my clang-format was configured incorrectly and changed the formatting of the whole file. I did have to modify `_tsinfermodule.c` to add allow for the addition of derived counts when adding sites. For some reason my pre-commit isn't linting the C code correctly, just trying to fix that now. I will speak with Ben about the failing Windows tests.
author_association: CONTRIBUTOR
user: login: Duncan-JR
created_at: 2025-03-18T17:10:50Z
body: @benjeffery I've simplified the changes based on Jerome's suggestions. I think the plan is to merge this in after 0.4.1 has been released, if that's okay with you? 
author_association: CONTRIBUTOR
user: login: benjeffery
created_at: 2025-03-18T18:49:43Z
body: Yep, makes sense to me! 0.4.1 should be tomorrow.
author_association: MEMBER
user: login: hyanwong
created_at: 2025-03-19T14:46:38Z
body: A trivial point, but perhaps you want to change the plot in the [large_scale.md](https://github.com/tskit-dev/tsinfer/blob/main/docs/large_scale.md) doc file, to show the new matching profile, if it is much different. Here's what we currently have (I guess @benjeffery knows how it was generated):

<img width="492" alt="Screenshot 2025-03-19 at 14 45 04" src="https://github.com/user-attachments/assets/215459bc-5232-4ee2-9d85-ef403c691a1d" />

author_association: MEMBER
user: login: benjeffery
created_at: 2025-03-19T14:58:31Z
body: That plot was from a GeL chromosome, but any similar will do that illustrates the point.
author_association: MEMBER
user: login: jeromekelleher
created_at: 2025-03-19T20:45:57Z
body: Good to merge @benjeffery?
author_association: MEMBER
user: login: benjeffery
created_at: 2025-03-19T22:48:40Z
body: Merging! Thanks @duncanMR this is awesome!
author_association: MEMBER
user: login: mergify[bot]
created_at: 2025-03-19T22:48:41Z
body: This pull request has been removed from the queue for the following reason: <code>pull request branch update failed</code>.

The pull request can't be updated

You should update or rebase your pull request manually.

If you want to requeue this pull request, you can post a `@mergifyio requeue` comment.

author_association: CONTRIBUTOR
user: login: benjeffery
created_at: 2025-03-20T01:23:49Z
body: @mergifyio rebase
author_association: MEMBER
user: login: mergify[bot]
created_at: 2025-03-20T01:23:58Z
body: > rebase

#### ✅ Branch has been successfully rebased



<!---
DO NOT EDIT
-*- Mergify Payload -*-
{"command": "rebase", "conclusion": "success", "action_is_running": true}
-*- Mergify Payload End -*-
-->
author_association: CONTRIBUTOR