individual convenience extractors

c/tests/test_trees.c

@@ -6686,6 +6686,172 @@ test_treeseq_row_access_errors(void)
     tsk_table_collection_free(&tables);
 }
 
+static void
+test_treeseq_get_individuals_population_errors(void)
+{
+    int ret;
+    tsk_id_t ret_id;
+    tsk_table_collection_t tables;
+    tsk_treeseq_t ts;
+    tsk_id_t output[2];
+
+    ret = tsk_table_collection_init(&tables, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    tables.sequence_length = 1;
+
+    ret_id = tsk_population_table_add_row(&tables.populations, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 0);
+    ret_id = tsk_individual_table_add_row(
+        &tables.individuals, 0, NULL, 0, NULL, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 0);
+    ret_id = tsk_individual_table_add_row(
+        &tables.individuals, 0, NULL, 0, NULL, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 1);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 1.25, 0, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 0);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 1.25, TSK_NULL, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 1);
+    ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    ret_id = tsk_treeseq_get_individuals_population(&ts, output);
+    CU_ASSERT_EQUAL_FATAL(ret_id, TSK_ERR_INDIVIDUAL_POPULATION_MISMATCH);
+
+    tsk_treeseq_free(&ts);
+    tsk_table_collection_free(&tables);
+}
+
+static void
+test_treeseq_get_individuals_population(void)
+{
+    int ret;
+    tsk_id_t ret_id;
+    int j;
+    tsk_table_collection_t tables;
+    tsk_treeseq_t ts;
+    tsk_id_t output[4];
+
+    ret = tsk_table_collection_init(&tables, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    tables.sequence_length = 1;
+
+    for (j = 0; j < 2; j++) {
+        ret_id = tsk_population_table_add_row(&tables.populations, NULL, 0);
+        CU_ASSERT_EQUAL_FATAL(ret_id, (tsk_id_t) j);
+    }
+    for (j = 0; j < 4; j++) {
+        ret_id = tsk_individual_table_add_row(
+            &tables.individuals, 0, NULL, 0, NULL, 0, NULL, 0);
+        CU_ASSERT_EQUAL_FATAL(ret_id, (tsk_id_t) j);
+    }
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 1.25, 0, 1, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 0);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 0.0, TSK_NULL, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 1);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 3.0, 1, 3, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 2);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 0.0, TSK_NULL, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 3);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 1.25, 0, 1, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 4);
+    ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    ret = tsk_treeseq_get_individuals_population(&ts, output);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    CU_ASSERT_EQUAL_FATAL(output[0], TSK_NULL);
+    CU_ASSERT_EQUAL_FATAL(output[1], 0);
+    CU_ASSERT_EQUAL_FATAL(output[2], TSK_NULL);
+    CU_ASSERT_EQUAL_FATAL(output[3], 1);
+
+    tsk_treeseq_free(&ts);
+    tsk_table_collection_free(&tables);
+}
+
+static void
+test_treeseq_get_individuals_time_errors(void)
+{
+    int ret;
+    tsk_id_t ret_id;
+    tsk_table_collection_t tables;
+    tsk_treeseq_t ts;
+    double output[2];
+
+    ret = tsk_table_collection_init(&tables, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    tables.sequence_length = 1;
+
+    ret_id = tsk_population_table_add_row(&tables.populations, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 0);
+    ret_id = tsk_individual_table_add_row(
+        &tables.individuals, 0, NULL, 0, NULL, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 0);
+    ret_id = tsk_individual_table_add_row(
+        &tables.individuals, 0, NULL, 0, NULL, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 1);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 1.2, 0, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 0);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 0.8, 0, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 1);
+    ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    ret = tsk_treeseq_get_individuals_time(&ts, output);
+    CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_INDIVIDUAL_TIME_MISMATCH);
+
+    tsk_treeseq_free(&ts);
+    tsk_table_collection_free(&tables);
+}
+
+static void
+test_treeseq_get_individuals_time(void)
+{
+    int ret;
+    tsk_id_t ret_id;
+    int j;
+    tsk_table_collection_t tables;
+    tsk_treeseq_t ts;
+    double output[4];
+
+    ret = tsk_table_collection_init(&tables, 0);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+    tables.sequence_length = 1;
+
+    for (j = 0; j < 2; j++) {
+        ret_id = tsk_population_table_add_row(&tables.populations, NULL, 0);
+        CU_ASSERT_EQUAL_FATAL(ret_id, j);
+    }
+    for (j = 0; j < 4; j++) {
+        ret_id = tsk_individual_table_add_row(
+            &tables.individuals, 0, NULL, 0, NULL, 0, NULL, 0);
+        CU_ASSERT_EQUAL_FATAL(ret_id, j);
+    }
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 1.25, 0, 1, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 0);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 3.25, 0, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 1);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 3.0, 1, 3, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 2);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 3.25, 0, 0, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 3);
+    ret_id = tsk_node_table_add_row(&tables.nodes, 0, 1.25, 0, 1, NULL, 0);
+    CU_ASSERT_EQUAL_FATAL(ret_id, 4);
+    ret = tsk_treeseq_init(&ts, &tables, TSK_TS_INIT_BUILD_INDEXES);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    ret = tsk_treeseq_get_individuals_time(&ts, output);
+    CU_ASSERT_EQUAL_FATAL(ret, 0);
+
+    CU_ASSERT_EQUAL_FATAL(output[0], 3.25);
+    CU_ASSERT_EQUAL_FATAL(output[1], 1.25);
+    CU_ASSERT_FATAL(tsk_is_unknown_time(output[2]));
+    CU_ASSERT_EQUAL_FATAL(output[3], 3.0);
+
+    tsk_treeseq_free(&ts);
+    tsk_table_collection_free(&tables);
+}
+
 static void
 test_tree_copy_flags(void)
 {
@@ -7637,6 +7803,13 @@ main(int argc, char **argv)
         /* Misc */
         { "test_tree_errors", test_tree_errors },
         { "test_treeseq_row_access_errors", test_treeseq_row_access_errors },
+        { "test_treeseq_get_individuals_population_errors",
+            test_treeseq_get_individuals_population_errors },
+        { "test_treeseq_get_individuals_population",
+            test_treeseq_get_individuals_population },
+        { "test_treeseq_get_individuals_time_errors",
+            test_treeseq_get_individuals_time_errors },
+        { "test_treeseq_get_individuals_time", test_treeseq_get_individuals_time },
         { "test_tree_copy_flags", test_tree_copy_flags },
         { "test_genealogical_nearest_neighbours_errors",
             test_genealogical_nearest_neighbours_errors },

c/tskit/core.c

@@ -571,6 +571,18 @@ tsk_strerror_internal(int err)
             ret = "Individuals cannot be their own ancestor. "
                   "(TSK_ERR_INDIVIDUAL_PARENT_CYCLE)";
             break;
+
+        case TSK_ERR_INDIVIDUAL_POPULATION_MISMATCH:
+            ret = "Individual populations cannot be returned "
+                  "if an individual has nodes from more than one population. "
+                  "(TSK_ERR_INDIVIDUAL_POPULATION_MISMATCH)";
+            break;
+
+        case TSK_ERR_INDIVIDUAL_TIME_MISMATCH:
+            ret = "Individual times cannot be returned "
+                  "if an individual has nodes from more than one time. "
+                  "(TSK_ERR_INDIVIDUAL_TIME_MISMATCH)";
+            break;
     }
     return ret;
 }

c/tskit/core.h

@@ -812,6 +812,16 @@ An individual was its own parent.
 An individual was its own ancestor in a cycle of references.
 */
 #define TSK_ERR_INDIVIDUAL_PARENT_CYCLE                            -1702
+/**
+An individual had nodes from more than one population
+(and only one was requested).
+*/
+#define TSK_ERR_INDIVIDUAL_POPULATION_MISMATCH                     -1703
+/**
+An individual had nodes from more than one time
+(and only one was requested).
+*/
+#define TSK_ERR_INDIVIDUAL_TIME_MISMATCH                           -1704
 /** @} */
 // clang-format on
 

c/tskit/trees.c

@@ -716,6 +716,68 @@ tsk_treeseq_has_reference_sequence(const tsk_treeseq_t *self)
     return tsk_table_collection_has_reference_sequence(self->tables);
 }
 
+int
+tsk_treeseq_get_individuals_population(const tsk_treeseq_t *self, tsk_id_t *output)
+{
+    int ret = 0;
+    tsk_size_t i, j;
+    tsk_individual_t ind;
+    tsk_id_t ind_pop;
+    const tsk_id_t *node_population = self->tables->nodes.population;
+    const tsk_size_t num_individuals = self->tables->individuals.num_rows;
+
+    tsk_memset(output, TSK_NULL, num_individuals * sizeof(*output));
+
+    for (i = 0; i < num_individuals; i++) {
+        ret = tsk_treeseq_get_individual(self, (tsk_id_t) i, &ind);
+        tsk_bug_assert(ret == 0);
+        if (ind.nodes_length > 0) {
+            ind_pop = -2;
+            for (j = 0; j < ind.nodes_length; j++) {
+                if (ind_pop == -2) {
+                    ind_pop = node_population[ind.nodes[j]];
+                } else if (ind_pop != node_population[ind.nodes[j]]) {
+                    ret = TSK_ERR_INDIVIDUAL_POPULATION_MISMATCH;
+                    goto out;
+                }
+            }
+            output[ind.id] = ind_pop;
+        }
+    }
+out:
+    return ret;
+}
+
+int
+tsk_treeseq_get_individuals_time(const tsk_treeseq_t *self, double *output)
+{
+    int ret = 0;
+    tsk_size_t i, j;
+    tsk_individual_t ind;
+    double ind_time;
+    const double *node_time = self->tables->nodes.time;
+    const tsk_size_t num_individuals = self->tables->individuals.num_rows;
+
+    for (i = 0; i < num_individuals; i++) {
+        ret = tsk_treeseq_get_individual(self, (tsk_id_t) i, &ind);
+        tsk_bug_assert(ret == 0);
+        /* the default is UNKNOWN_TIME, but nodes cannot have
+         * UNKNOWN _TIME so this is safe. */
+        ind_time = TSK_UNKNOWN_TIME;
+        for (j = 0; j < ind.nodes_length; j++) {
+            if (j == 0) {
+                ind_time = node_time[ind.nodes[j]];
+            } else if (ind_time != node_time[ind.nodes[j]]) {
+                ret = TSK_ERR_INDIVIDUAL_TIME_MISMATCH;
+                goto out;
+            }
+        }
+        output[ind.id] = ind_time;
+    }
+out:
+    return ret;
+}
+
 /* Stats functions */
 
 #define GET_2D_ROW(array, row_len, row) (array + (((size_t)(row_len)) * (size_t) row))

c/tskit/trees.h

@@ -697,7 +697,7 @@ Returns the location of each node in the list of samples or
 @endrst
 
 @param self A pointer to a tsk_treeseq_t object.
-@return Returns the pointer to the breakpoint array.
+@return Returns the pointer to the array of sample indexes.
 */
 const tsk_id_t *tsk_treeseq_get_sample_index_map(const tsk_treeseq_t *self);
 
@@ -890,6 +890,9 @@ int tsk_treeseq_split_edges(const tsk_treeseq_t *self, double time, tsk_flags_t
 
 bool tsk_treeseq_has_reference_sequence(const tsk_treeseq_t *self);
 
+int tsk_treeseq_get_individuals_population(const tsk_treeseq_t *self, tsk_id_t *output);
+int tsk_treeseq_get_individuals_time(const tsk_treeseq_t *self, double *output);
+
 int tsk_treeseq_kc_distance(const tsk_treeseq_t *self, const tsk_treeseq_t *other,
     double lambda_, double *result);
 

python/CHANGELOG.rst

@@ -26,6 +26,10 @@
   edges at a specific time.
   (:user:`jeromekelleher`, :issue:`2276`, :pr:`2296`).
 
+- Add the ``TreeSequence.individuals_time`` and ``TreeSequence.individuals_population``
+  methods to return arrays of per-individual times and populations, respectively.
+  (:user:`petrelharp`, :issue:`1481`, :pr:`2298`).
+
 **Breaking Changes**
 
 - The JSON metadata codec now interprets the empty string as an empty object. This means