Revise ancestor generation algorithm to improve performance

Author: duncanMR

_tsinfermodule.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;
@@ -647,7 +653,7 @@ ancestor_builder_allocate_genotypes(ancestor_builder_t *self)
 }
 
 int WARN_UNUSED
-ancestor_builder_add_site(ancestor_builder_t *self, double time, allele_t *genotypes)
+ancestor_builder_add_site(ancestor_builder_t *self, double time, allele_t *genotypes, tsk_size_t derived_count)
 {
     int ret = 0;
     site_t *site;
@@ -676,6 +682,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;

lib/ancestor_builder.c

@@ -375,7 +375,7 @@ run_random_data(size_t num_samples, size_t num_sites, int seed,
             genotypes[k] = samples[k][j];
             time += genotypes[k];
         }
-        ret = ancestor_builder_add_site(&ancestor_builder, time, genotypes);
+        ret = ancestor_builder_add_site(&ancestor_builder, time, genotypes, 1);
         CU_ASSERT_EQUAL_FATAL(ret, 0);
     }
     /* ancestor_builder_print_state(&ancestor_builder, stdout); */
@@ -478,15 +478,15 @@ test_ancestor_builder_errors(void)
     ret = ancestor_builder_alloc(&ancestor_builder, 2, 0, -1, 0);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
     CU_ASSERT_EQUAL_FATAL(ancestor_builder.num_sites, 0);
-    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes_ones);
+    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes_ones, 1);
     CU_ASSERT_EQUAL_FATAL(ret, TSI_ERR_TOO_MANY_SITES);
     ancestor_builder_free(&ancestor_builder);
 
     ret = ancestor_builder_alloc(&ancestor_builder, 4, 2, -1, 0);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
-    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes_zeros);
+    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes_zeros, 1);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
-    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes_ones);
+    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes_ones, 1);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
     CU_ASSERT_EQUAL_FATAL(ancestor_builder.num_sites, 2);
     ret = ancestor_builder_finalise(&ancestor_builder);
@@ -509,7 +509,7 @@ test_ancestor_builder_one_site(void)
 
     ret = ancestor_builder_alloc(&ancestor_builder, 4, 1, -1, 0);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
-    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes);
+    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes, 1);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
     ret = ancestor_builder_finalise(&ancestor_builder);
     CU_ASSERT_EQUAL_FATAL(ret, 0);

lib/tests/tests.c

@@ -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 {
@@ -214,7 +215,7 @@ int ancestor_builder_alloc(ancestor_builder_t *self, size_t num_samples,
 int ancestor_builder_free(ancestor_builder_t *self);
 int ancestor_builder_print_state(ancestor_builder_t *self, FILE *out);
 int ancestor_builder_add_site(
-    ancestor_builder_t *self, double time, allele_t *genotypes);
+    ancestor_builder_t *self, double time, allele_t *genotypes, tsk_size_t derived_count);
 int ancestor_builder_finalise(ancestor_builder_t *self);
 int ancestor_builder_make_ancestor(const ancestor_builder_t *self,
     size_t num_focal_sites, const tsk_id_t *focal_sites, tsk_id_t *start, tsk_id_t *end,

lib/tsinfer.h

@@ -1702,7 +1702,7 @@ def test_array_args(self, tmp_path, tmpdir):
         ts.tables.assert_equals(ts_batch.tables, ignore_provenance=True)
 
 
-class TestAncestorGeneratorsEquivalant:
+class TestAncestorGeneratorsEquivalent:
     """
     Tests for the ancestor generation process.
     """
@@ -1920,7 +1920,7 @@ def test_bad_focal_sites(self):
             g = np.zeros(2, dtype=np.int8)
             h = np.zeros(1, dtype=np.int8)
             generator = tsinfer.AncestorsGenerator(sample_data, None, {}, engine=engine)
-            generator.ancestor_builder.add_site(1, g)
+            generator.ancestor_builder.add_site(1, g, derived_count=0)
             with pytest.raises(error):
                 generator.ancestor_builder.make_ancestor([0], h)
 
@@ -2725,7 +2725,10 @@ def test_ancestor_builder_print_state(self):
         sample_data = self.sample_example(n_samples, n_sites)
         ancestor_builder = tsinfer.algorithm.AncestorBuilder(n_samples, n_sites)
         for variant in sample_data.variants():
-            ancestor_builder.add_site(variant.site.time, variant.genotypes)
+            derived_count = np.sum(variant.genotypes)
+            ancestor_builder.add_site(
+                variant.site.time, variant.genotypes, derived_count
+            )
         with mock.patch("sys.stdout", new=io.StringIO()) as mock_output:
             ancestor_builder.print_state()
             # Simply check some text is output
@@ -3533,16 +3536,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 +3554,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)

tests/test_inference.py

@@ -134,19 +134,22 @@ def test_add_site(self):
         ab = _tsinfer.AncestorBuilder(num_samples=2, max_sites=10)
         for bad_type in ["sdf", {}, None]:
             with pytest.raises(TypeError):
-                ab.add_site(time=bad_type, genotypes=[0, 0])
+                ab.add_site(time=bad_type, genotypes=[0, 0], derived_count=0)
         for bad_genotypes in ["asdf", [[], []], [0, 1, 2]]:
             with pytest.raises(ValueError):
-                ab.add_site(time=0, genotypes=bad_genotypes)
+                ab.add_site(time=0, genotypes=bad_genotypes, derived_count=0)
+        for bad_derived_count in ["asdf", 1.2, [0, 1]]:
+            with pytest.raises(TypeError):
+                ab.add_site(time=0, genotypes=[0, 0], derived_count=bad_derived_count)
 
     def test_add_too_many_sites(self):
         for max_sites in range(10):
             ab = _tsinfer.AncestorBuilder(num_samples=2, max_sites=max_sites)
             for _ in range(max_sites):
-                ab.add_site(time=1, genotypes=[0, 1])
+                ab.add_site(time=1, genotypes=[0, 1], derived_count=0)
             for _ in range(2 * max_sites):
                 with pytest.raises(_tsinfer.LibraryError) as record:
-                    ab.add_site(time=1, genotypes=[0, 1])
+                    ab.add_site(time=1, genotypes=[0, 1], derived_count=0)
                 msg = "Cannot add more sites than the specified maximum."
                 assert str(record.value) == msg