Optional one-bit encoding for generate_ancestors

_tsinfermodule.c

@@ -92,20 +92,21 @@ AncestorBuilder_init(AncestorBuilder *self, PyObject *args, PyObject *kwds)
 {
     int ret = -1;
     int err;
-    static char *kwlist[] = {"num_samples", "max_sites", NULL};
-    int num_samples, max_sites;
+    static char *kwlist[] = {"num_samples", "max_sites", "genotype_encoding", NULL};
+    int num_samples, max_sites, genotype_encoding;
     int flags = 0;
 
     self->builder = NULL;
-    if (!PyArg_ParseTupleAndKeywords(args, kwds, "ii", kwlist,
-                &num_samples, &max_sites)) {
+    if (!PyArg_ParseTupleAndKeywords(args, kwds, "ii|i", kwlist,
+                &num_samples, &max_sites, &genotype_encoding)) {
         goto out;
     }
     self->builder = PyMem_Malloc(sizeof(ancestor_builder_t));
     if (self->builder == NULL) {
         PyErr_NoMemory();
         goto out;
     }
+    flags = genotype_encoding;
     Py_BEGIN_ALLOW_THREADS
     err = ancestor_builder_alloc(self->builder, num_samples, max_sites, flags);
     Py_END_ALLOW_THREADS
@@ -274,34 +275,6 @@ AncestorBuilder_ancestor_descriptors(AncestorBuilder *self)
         }
         PyTuple_SET_ITEM(descriptors, j, py_descriptor);
     }
-    /* j = 0; */
-    /* /1* It's not great that we're breaking encapsulation here and looking */
-    /*  * directly in to the builder's data structures. However, it's quite an */
-    /*  * awkward set of data to communicate, so it seems OK. *1/ */
-    /* for (f = self->builder->num_samples; f > 0; f--) { */
-    /*     for (a = self->builder->frequency_map[f].head; a != NULL; a = a->next) { */
-    /*         map_elem = (pattern_map_t *) a->item; */
-    /*         dims = map_elem->num_sites; */
-    /*         site_array = (PyArrayObject *) PyArray_SimpleNew(1, &dims, NPY_INT32); */
-    /*         if (site_array == NULL) { */
-    /*             goto out; */
-    /*         } */
-    /*         site_array_data = (int32_t *) PyArray_DATA(site_array); */
-    /*         /1* The elements are listed backwards, so reverse them *1/ */
-    /*         k = map_elem->num_sites - 1; */
-    /*         for (s = map_elem->sites; s != NULL; s = s->next) { */
-    /*             site_array_data[k] = (int32_t) s->site; */
-    /*             k--; */
-    /*         } */
-    /*         descriptor = Py_BuildValue("kO", (unsigned long) f, site_array); */
-    /*         if (descriptor == NULL) { */
-    /*             Py_DECREF(site_array); */
-    /*             goto out; */
-    /*         } */
-    /*         PyTuple_SET_ITEM(descriptors, j, descriptor); */
-    /*         j++; */
-    /*     } */
-    /* } */
     ret = descriptors;
     descriptors = NULL;
 out:
@@ -335,13 +308,31 @@ AncestorBuilder_get_num_ancestors(AncestorBuilder *self, void *closure)
     return ret;
 }
 
+static PyObject *
+AncestorBuilder_get_memsize(AncestorBuilder *self, void *closure)
+{
+    PyObject *ret = NULL;
+
+    if (AncestorBuilder_check_state(self) != 0) {
+        goto out;
+    }
+    ret = Py_BuildValue("k", (unsigned long) ancestor_builder_get_memsize(
+                self->builder));
+out:
+    return ret;
+}
+
+
 static PyMemberDef AncestorBuilder_members[] = {
     {NULL}  /* Sentinel */
 };
 
 static PyGetSetDef AncestorBuilder_getsetters[] = {
     {"num_sites", (getter) AncestorBuilder_get_num_sites, NULL, "The number of sites."},
-    {"num_ancestors", (getter) AncestorBuilder_get_num_ancestors, NULL, "The number of ancestors."},
+    {"num_ancestors", (getter) AncestorBuilder_get_num_ancestors, NULL,
+        "The number of ancestors."},
+    {"mem_size", (getter) AncestorBuilder_get_memsize, NULL,
+        "The number of allocated bytes."},
     {NULL}  /* Sentinel */
 };
 

docs/api.rst

@@ -49,6 +49,9 @@ Running inference
 
 .. autofunction:: tsinfer.generate_ancestors
 
+.. autoclass:: tsinfer.GenotypeEncoding
+   :members:
+
 .. autofunction:: tsinfer.match_ancestors
 
 .. autofunction:: tsinfer.match_samples