Add terminal site support to C ancestor builder

Author: duncanMR

_tsinfermodule.c

@@ -136,18 +136,19 @@ static PyObject *
 AncestorBuilder_add_site(AncestorBuilder *self, PyObject *args, PyObject *kwds)
 {
     int err;
-    static char *kwlist[] = {"time", "genotypes", NULL};
+    static char *kwlist[] = {"time", "genotypes", "terminal", NULL};
     PyObject *ret = NULL;
     double time;
     PyObject *genotypes = NULL;
     PyArrayObject *genotypes_array = NULL;
     npy_intp *shape;
+    int terminal = 0;
 
     if (AncestorBuilder_check_state(self) != 0) {
         goto out;
     }
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "dO", kwlist,
-            &time, &genotypes)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "dO|p", kwlist,
+            &time, &genotypes, &terminal)) {
         goto out;
     }
     genotypes_array = (PyArrayObject *) PyArray_FROM_OTF(genotypes, NPY_INT8,
@@ -166,7 +167,7 @@ AncestorBuilder_add_site(AncestorBuilder *self, PyObject *args, PyObject *kwds)
     }
     Py_BEGIN_ALLOW_THREADS
     err = ancestor_builder_add_site(self->builder, time,
-            (allele_t *) PyArray_DATA(genotypes_array));
+            (allele_t *) PyArray_DATA(genotypes_array), terminal);
     Py_END_ALLOW_THREADS
     if (err != 0) {
         handle_library_error(err);

lib/ancestor_builder.c

@@ -439,6 +439,9 @@ ancestor_builder_compute_ancestral_states(const ancestor_builder_t *self, int di
     /*         (int) min_sample_set_size); */
     for (l = focal_site + direction; l >= 0 && l < (int64_t) num_sites; l += direction) {
         /* printf("\tl = %d\n", (int) l); */
+        if (sites[l].terminal) {
+            break;
+        }
         ancestor[l] = 0;
         last_site = (tsk_id_t) l;
 
@@ -653,7 +656,8 @@ 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, bool terminal)
 {
     int ret = 0;
     site_t *site;
@@ -665,21 +669,30 @@ ancestor_builder_add_site(ancestor_builder_t *self, double time, allele_t *genot
     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);
+    time_map_t *time_map = NULL;
 
+    if (self->num_sites == self->max_sites) {
+        ret = TSI_ERR_TOO_MANY_SITES;
+        goto out;
+    }
     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;
+    site = &self->sites[site_id];
+    site->time = time;
+    site->derived_count = derived_count;
+    site->terminal = terminal;
+    if (terminal) {
+        site->encoded_genotypes = NULL;
+        self->num_sites++;
         goto out;
     }
-    if (self->num_sites == self->max_sites) {
-        ret = TSI_ERR_TOO_MANY_SITES;
+    time_map = ancestor_builder_get_time_map(self, time);
+    if (time_map == NULL) {
+        ret = TSI_ERR_NO_MEMORY;
         goto out;
     }
     ret = ancestor_builder_encode_genotypes(self, genotypes, encoded_genotypes);
@@ -688,9 +701,6 @@ ancestor_builder_add_site(ancestor_builder_t *self, double time, allele_t *genot
     }
     self->num_sites++;
     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/tests/tests.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, false);
         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, false);
     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, false);
     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, false);
     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, false);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
     ret = ancestor_builder_finalise(&ancestor_builder);
     CU_ASSERT_EQUAL_FATAL(ret, 0);

lib/tsinfer.h

@@ -104,6 +104,7 @@ typedef struct {
     double time;
     uint8_t *encoded_genotypes;
     tsk_size_t derived_count;
+    bool terminal;
 } site_t;
 
 typedef struct {
@@ -251,7 +252,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, bool terminal);
 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,