Simplify handling of derived counts

Author: duncanMR

_tsinfermodule.c

@@ -474,12 +474,12 @@ ancestor_builder_compute_ancestral_states(const ancestor_builder_t *self, int di
             if (disagree[u] && (genotypes[j] != consensus)
                 && (genotypes[j] != TSK_MISSING_DATA)) {
                 /* This sample has disagreed with consensus twice in a row,
-                    * so remove it */
+                 * 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) {
@@ -489,14 +489,14 @@ ancestor_builder_compute_ancestral_states(const ancestor_builder_t *self, int di
             }
         }
         /* For the remaining samples, set the disagree flags based
-            * on whether they agree with the consensus for this site. */
+         * 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));
+                                   && (genotypes[j] != TSK_MISSING_DATA));
                 }
             }
         }
@@ -653,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, tsk_size_t derived_count)
+ancestor_builder_add_site(ancestor_builder_t *self, double time, allele_t *genotypes)
 {
     int ret = 0;
     site_t *site;
@@ -664,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;

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, 1);
+        ret = ancestor_builder_add_site(&ancestor_builder, time, genotypes);
         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, 1);
+    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes_ones);
     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, 1);
+    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes_zeros);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
-    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes_ones, 1);
+    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes_ones);
     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, 1);
+    ret = ancestor_builder_add_site(&ancestor_builder, 4, genotypes);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
     ret = ancestor_builder_finalise(&ancestor_builder);
     CU_ASSERT_EQUAL_FATAL(ret, 0);

lib/tests/tests.c

@@ -215,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, tsk_size_t derived_count);
+    ancestor_builder_t *self, double time, allele_t *genotypes);
 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 TestAncestorGeneratorsEquivalent:
+class TestAncestorGeneratorsEquivalant:
     """
     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, derived_count=0)
+            generator.ancestor_builder.add_site(1, g)
             with pytest.raises(error):
                 generator.ancestor_builder.make_ancestor([0], h)
 
@@ -2725,10 +2725,7 @@ 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():
-            derived_count = np.sum(variant.genotypes)
-            ancestor_builder.add_site(
-                variant.site.time, variant.genotypes, derived_count
-            )
+            ancestor_builder.add_site(variant.site.time, variant.genotypes)
         with mock.patch("sys.stdout", new=io.StringIO()) as mock_output:
             ancestor_builder.print_state()
             # Simply check some text is output

tests/test_inference.py

@@ -134,22 +134,19 @@ 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], derived_count=0)
+                ab.add_site(time=bad_type, genotypes=[0, 0])
         for bad_genotypes in ["asdf", [[], []], [0, 1, 2]]:
             with pytest.raises(ValueError):
-                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)
+                ab.add_site(time=0, genotypes=bad_genotypes)
 
     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], derived_count=0)
+                ab.add_site(time=1, genotypes=[0, 1])
             for _ in range(2 * max_sites):
                 with pytest.raises(_tsinfer.LibraryError) as record:
-                    ab.add_site(time=1, genotypes=[0, 1], derived_count=0)
+                    ab.add_site(time=1, genotypes=[0, 1])
                 msg = "Cannot add more sites than the specified maximum."
                 assert str(record.value) == msg
 

tests/test_low_level.py

@@ -137,11 +137,12 @@ def store_site_genotypes(self, site_id, genotypes):
         stop = start + self.encoded_genotypes_size
         self.genotype_store[start:stop] = genotypes
 
-    def add_site(self, time, genotypes, derived_count):
+    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, derived_count))
         sites_at_fixed_timepoint = self.time_map[time]
@@ -203,6 +204,15 @@ def ancestor_descriptors(self):
         return ret
 
     def compute_ancestral_states(self, a, focal_site, sites):
+        """
+        For a given focal site, and set of sites to fill in (usually all the ones
+        leftwards or rightwards), augment the haplotype array a with the inferred sites
+        Together with `make_ancestor`, which calls this function, these describe the main
+        algorithm as implemented in Fig S2 of the preprint, with the buffer.
+
+        At the moment we assume that the derived state is 1. We should alter this so
+        that we allow the derived state to be a different non-zero integer.
+        """
         focal_time = self.sites[focal_site].time
         g = self.get_site_genotypes(focal_site)
         sample_set = np.where(g == 1)[0]

tsinfer/algorithm.py

@@ -1881,9 +1881,7 @@ def add_sites(self, exclude_positions=None):
                 if np.isnan(time):
                     use_site = False  # Site with meaningless time value: skip inference
             if use_site:
-                self.ancestor_builder.add_site(
-                    time, variant.genotypes, int(counts.derived)
-                )
+                self.ancestor_builder.add_site(time, variant.genotypes)
                 inference_site_id.append(site.id)
                 self.num_sites += 1
             progress.update()