Pass num_alleles through to MatcherIndexes on construction

Add num_alleles parameter to matcher_indexes_alloc in C; if NULL,
defaults to biallelic (2).  Thread through Python/C binding and
MatcherIndexes wrapper.  Fixes TSI_ERR_BAD_HAPLOTYPE_ALLELE on
multi-allelic sites.  Add triallelic C test.

Author: jeromekelleher

_tsinfermodule.c

@@ -1683,29 +1683,42 @@ MatcherIndexes_init(MatcherIndexes *self, PyObject *args, PyObject *kwds)
     int ret = -1;
     int err;
     LightweightTableCollection *tables;
-    static char *kwlist[] = { "tables", NULL };
+    PyObject *num_alleles_obj = NULL;
+    PyArrayObject *num_alleles_array = NULL;
+    tsk_size_t *num_alleles_data = NULL;
+    static char *kwlist[] = { "tables", "num_alleles", NULL };
 
     self->matcher_indexes = NULL;
-    if (!PyArg_ParseTupleAndKeywords(
-            args, kwds, "O!", kwlist, &LightweightTableCollectionType, &tables)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "O!|O", kwlist,
+            &LightweightTableCollectionType, &tables, &num_alleles_obj)) {
         goto out;
     }
     if (LightweightTableCollection_check_state(tables) != 0) {
         goto out;
     }
+    if (num_alleles_obj != NULL && num_alleles_obj != Py_None) {
+        num_alleles_array = (PyArrayObject *) PyArray_FROMANY(
+            num_alleles_obj, NPY_UINT64, 1, 1, NPY_ARRAY_IN_ARRAY);
+        if (num_alleles_array == NULL) {
+            goto out;
+        }
+        num_alleles_data = PyArray_DATA(num_alleles_array);
+    }
 
     self->matcher_indexes = PyMem_Calloc(1, sizeof(*self->matcher_indexes));
     if (self->matcher_indexes == NULL) {
         PyErr_NoMemory();
         goto out;
     }
-    err = matcher_indexes_alloc(self->matcher_indexes, tables->tables, 0);
+    err = matcher_indexes_alloc(
+        self->matcher_indexes, tables->tables, num_alleles_data, 0);
     if (err != 0) {
         handle_library_error(err);
         goto out;
     }
     ret = 0;
 out:
+    Py_XDECREF(num_alleles_array);
     return ret;
 }
 

lib/ancestor_matcher.c

@@ -1108,8 +1108,8 @@ matcher_indexes_add_mutation(
 }
 
 static int
-matcher_indexes_copy_mutation_data(
-    matcher_indexes_t *self, const tsk_table_collection_t *tables)
+matcher_indexes_copy_mutation_data(matcher_indexes_t *self,
+    const tsk_table_collection_t *tables, const tsk_size_t *num_alleles)
 {
     int ret = 0;
     tsk_size_t j;
@@ -1124,7 +1124,7 @@ matcher_indexes_copy_mutation_data(
         /* TODO check for under/overflow */
         converted_position[j] = (coordinate_t) sites_position[j];
         self->sites.mutations[j] = NULL;
-        self->sites.num_alleles[j] = 2;
+        self->sites.num_alleles[j] = num_alleles != NULL ? num_alleles[j] : 2;
     }
     converted_position[j] = (coordinate_t) tables->sequence_length;
 
@@ -1137,8 +1137,6 @@ matcher_indexes_copy_mutation_data(
         }
 
         self->sites.mutations[site] = NULL;
-        /* FIXME */
-        self->sites.num_alleles[site] = 2;
         ret = matcher_indexes_add_mutation(self, site, mutations_node[j], 1);
         if (ret != 0) {
             goto out;
@@ -1151,8 +1149,8 @@ matcher_indexes_copy_mutation_data(
 }
 
 int
-matcher_indexes_alloc(
-    matcher_indexes_t *self, const tsk_table_collection_t *tables, tsk_flags_t flags)
+matcher_indexes_alloc(matcher_indexes_t *self, const tsk_table_collection_t *tables,
+    const tsk_size_t *num_alleles, tsk_flags_t flags)
 {
     int ret = 0;
 
@@ -1185,7 +1183,7 @@ matcher_indexes_alloc(
     if (ret != 0) {
         goto out;
     }
-    ret = matcher_indexes_copy_mutation_data(self, tables);
+    ret = matcher_indexes_copy_mutation_data(self, tables, num_alleles);
     if (ret != 0) {
         goto out;
     }

lib/tests/tests.c

@@ -1082,7 +1082,7 @@ run_match(const tsk_treeseq_t *ts, double rho, double mu, const allele_t *h,
         recombination_rate[j] = rho;
     }
 
-    ret = matcher_indexes_alloc(&mi, ts->tables, 0);
+    ret = matcher_indexes_alloc(&mi, ts->tables, NULL, 0);
     CU_ASSERT_EQUAL_FATAL(ret, 0);
     /* matcher_indexes_print_state(&mi, stdout); */
     ret = ancestor_matcher2_alloc(&am, &mi, recombination_rate, mutation_rate, 14, 0);
@@ -1441,6 +1441,44 @@ test_matching_deep_chain_ts(void)
     tsk_treeseq_free(&ts);
 }
 
+static void
+test_matching_triallelic_ts(void)
+{
+    /*
+     * Test that num_alleles > 2 is accepted.
+     * Use the star topology but tell the matcher that site 0 has 3 alleles.
+     * Query with allele 2 at site 0 should not trigger BAD_HAPLOTYPE_ALLELE.
+     */
+    tsk_treeseq_t ts;
+    ancestor_matcher2_t am;
+    matcher_indexes_t mi;
+    allele_t match[3];
+    tsk_id_t left[3], right[3], parent[3];
+    tsk_size_t path_length;
+    tsk_size_t num_alleles[] = { 3, 2, 2 };
+    double rho[] = { 1e-9, 1e-9, 1e-9 };
+    double mu[] = { 1e-20, 1e-20, 1e-20 };
+    int ret;
+
+    build_star_ts(&ts);
+
+    ret = matcher_indexes_alloc(&mi, ts.tables, num_alleles, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    ret = ancestor_matcher2_alloc(&am, &mi, rho, mu, 14, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    /* Query [2,0,0]: allele 2 at site 0 needs num_alleles >= 3 */
+    allele_t h[] = { 2, 0, 0 };
+    ret = ancestor_matcher2_find_path(
+        &am, 0, 3, h, match, &path_length, left, right, parent);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    CU_ASSERT_EQUAL(path_length, 1);
+
+    ancestor_matcher2_free(&am);
+    matcher_indexes_free(&mi);
+    tsk_treeseq_free(&ts);
+}
+
 static void
 test_strerror(void)
 {
@@ -1563,6 +1601,7 @@ main(int argc, char **argv)
         { "test_matching_binary_ts", test_matching_binary_ts },
         { "test_matching_two_tree_ts", test_matching_two_tree_ts },
         { "test_matching_deep_chain_ts", test_matching_deep_chain_ts },
+        { "test_matching_triallelic_ts", test_matching_triallelic_ts },
 
         { "test_strerror", test_strerror },
 

lib/tsinfer.h

@@ -358,8 +358,8 @@ int tree_sequence_builder_dump_mutations(tree_sequence_builder_t *self, tsk_id_t
 
 /* New impelementation */
 
-int matcher_indexes_alloc(
-    matcher_indexes_t *self, const tsk_table_collection_t *tables, tsk_flags_t options);
+int matcher_indexes_alloc(matcher_indexes_t *self, const tsk_table_collection_t *tables,
+    const tsk_size_t *num_alleles, tsk_flags_t options);
 int matcher_indexes_print_state(const matcher_indexes_t *self, FILE *out);
 int matcher_indexes_free(matcher_indexes_t *self);
 

tests/test_lshmm.py

@@ -75,7 +75,7 @@ class MatcherIndexes:
     The memory that can be shared between AncestorMatcher instances.
     """
 
-    def __init__(self, in_tables, *, vestigial_root=True):
+    def __init__(self, in_tables, *, vestigial_root=True, num_alleles=None):
         ts = in_tables.tree_sequence()
         if vestigial_root:
             ts = matching.add_vestigial_root(ts)
@@ -93,8 +93,10 @@ def __init__(self, in_tables, *, vestigial_root=True):
             tables.edges, tables.indexes.edge_removal_order
         )
 
-        # TODO fixme
-        self.num_alleles = np.zeros(self.num_sites, dtype=int) + 2
+        if num_alleles is not None:
+            self.num_alleles = np.asarray(num_alleles, dtype=int)
+        else:
+            self.num_alleles = np.zeros(self.num_sites, dtype=int) + 2
         self.sites_position = np.zeros(ts.num_sites + 1, dtype=np.uint32)
         self.sites_position[:-1] = tables.sites.position
         self.sites_position[-1] = tables.sequence_length

tsinfer/matching.py

@@ -264,7 +264,9 @@ def __init__(
         )
         t0 = time_.monotonic()
         if self._use_matcher2:
-            self._matcher_indexes = MatcherIndexes(ts, vestigial_root=False)
+            self._matcher_indexes = MatcherIndexes(
+                ts, vestigial_root=False, num_alleles=num_alleles
+            )
         else:
             self._tsb = _tsb_from_ts(
                 ts,
@@ -551,15 +553,15 @@ def mean_traceback_size(self):
 class MatcherIndexes(_tsinfer.MatcherIndexes):
     """Wrapper around the C MatcherIndexes, built from a tree sequence."""
 
-    def __init__(self, ts, *, vestigial_root=True):
+    def __init__(self, ts, *, vestigial_root=True, num_alleles=None):
         if vestigial_root:
             ts = add_vestigial_root(ts)
         tables = ts.dump_tables()
         ll_tables = _tsinfer.LightweightTableCollection(tables.sequence_length)
         ll_tables.fromdict(tables.asdict())
         self.sequence_length = ts.sequence_length
         self.num_sites = ts.num_sites
-        super().__init__(ll_tables)
+        super().__init__(ll_tables, num_alleles=num_alleles)
 
 
 class AncestorMatcher2(_tsinfer.AncestorMatcher2):