Merge pull request #1802 from nikbaya/update-climb-pedigrees

Update setup of WF pedigree model for table-based approach

lib/msprime.c

@@ -844,6 +844,7 @@ msp_free_pedigree(msp_t *self)
         for (i = 0; i < self->pedigree->num_inds; i++) {
             msp_safe_free(ind->parents);
             msp_safe_free(ind->segments);
+            msp_safe_free(ind->nodes);
             ind++;
         }
     }
@@ -2122,7 +2123,8 @@ alloc_individual(individual_t *ind, size_t ploidy)
     // Better to allocate these as a block?
     ind->parents = malloc(ploidy * sizeof(individual_t *));
     ind->segments = malloc(ploidy * sizeof(avl_tree_t));
-    if (ind->parents == NULL || ind->segments == NULL) {
+    ind->nodes = malloc(ploidy * sizeof(*ind->nodes));
+    if (ind->parents == NULL || ind->segments == NULL || ind->nodes == NULL) {
         ret = MSP_ERR_NO_MEMORY;
         goto out;
     }
@@ -2223,45 +2225,6 @@ msp_alloc_pedigree(msp_t *self, size_t num_inds, size_t ploidy)
     return ret;
 }
 
-/* TODO merge this method into where it's called - we don't really need these
- * arrays any more. */
-static int MSP_WARN_UNUSED
-msp_set_pedigree(msp_t *self, tsk_id_t *parents, double *times, tsk_flags_t *is_sample)
-{
-    size_t i, j;
-    tsk_id_t parent_ix;
-    tsk_flags_t sample_flag;
-    size_t sample_num;
-    individual_t *ind = NULL;
-
-    tsk_bug_assert(self->pedigree != NULL);
-
-    ind = self->pedigree->inds;
-    sample_num = 0;
-    for (i = 0; i < self->pedigree->num_inds; i++) {
-        ind = &self->pedigree->inds[i];
-        ind->id = (tsk_id_t) i;
-        ind->time = times[i];
-
-        // Link individuals to parents
-        for (j = 0; j < self->ploidy; j++) {
-            parent_ix = parents[i * self->ploidy + j];
-            if (parent_ix != TSK_NULL) {
-                ind->parents[j] = &self->pedigree->inds[parent_ix];
-            }
-        }
-        // Set samples
-        sample_flag = is_sample[i];
-        if (sample_flag != 0) {
-            tsk_bug_assert(sample_flag == 1);
-            self->pedigree->samples[sample_num] = ind;
-            sample_num++;
-        }
-    }
-    self->pedigree->num_samples = sample_num;
-    return 0;
-}
-
 static void
 msp_check_samples(msp_t *self)
 {
@@ -7230,13 +7193,20 @@ msp_set_simulation_model_wf_ped(msp_t *self)
 {
     int ret;
     tsk_size_t j, k;
+    tsk_treeseq_t ts;
     tsk_size_t num_individuals = self->tables->individuals.num_rows;
-    tsk_id_t *parents = NULL;
-    const tsk_id_t *ind_parents;
-    tsk_individual_t ind;
-    tsk_node_t node;
-    double *time = NULL;
-    tsk_flags_t *is_sample = NULL;
+    tsk_individual_t ind_obj;
+    individual_t *ind;
+    tsk_node_t node_obj;
+    tsk_id_t node_id;
+    tsk_size_t sample_num = 0;
+    bool sample_flag;
+
+    ret = tsk_treeseq_init(&ts, self->tables, TSK_BUILD_INDEXES);
+    if (ret != 0) {
+        ret = msp_set_tsk_error(ret);
+        goto out;
+    }
 
     /* TODO better error codes here */
     if (self->ploidy != 2) {
@@ -7253,52 +7223,47 @@ msp_set_simulation_model_wf_ped(msp_t *self)
         goto out;
     }
 
-    parents = malloc(num_individuals * self->ploidy * sizeof(*parents));
-    time = malloc(num_individuals * sizeof(*time));
-    is_sample = malloc(num_individuals * sizeof(*is_sample));
-    if (parents == NULL || time == NULL || is_sample == NULL) {
-        ret = MSP_ERR_NO_MEMORY;
-        goto out;
-    }
-
     for (j = 0; j < num_individuals; j++) {
-        ret = tsk_individual_table_get_row(
-            &self->tables->individuals, (tsk_id_t) j, &ind);
+        ret = tsk_treeseq_get_individual(&ts, (tsk_id_t) j, &ind_obj);
         tsk_bug_assert(ret == 0);
-        ind_parents = ind.parents;
-        tsk_bug_assert(ind.parents_length == self->ploidy);
+        ind = &self->pedigree->inds[j];
+        ind->id = (tsk_id_t) j;
+        tsk_bug_assert(ind_obj.parents_length == self->ploidy);
         for (k = 0; k < self->ploidy; k++) {
-            if (ind_parents[k] < TSK_NULL
-                || ind_parents[k] >= (tsk_id_t) num_individuals) {
+            if (ind_obj.parents[k] < TSK_NULL
+                || ind_obj.parents[k] >= (tsk_id_t) num_individuals) {
                 ret = TSK_ERR_INDIVIDUAL_OUT_OF_BOUNDS;
                 goto out;
             }
-            parents[j * self->ploidy + k] = ind_parents[k];
+            if (ind_obj.parents[k] != TSK_NULL) {
+                ind->parents[k] = &self->pedigree->inds[ind_obj.parents[k]];
+            }
+        }
+        sample_flag = false;
+
+        for (k = 0; k < ind_obj.nodes_length; k++) {
+            // ret = tsk_treeseq_get_node(&ts, ind_obj.nodes[k], &node_obj);
+            node_id
+                = (int) j * (int) ind_obj.nodes_length
+                  + (int)
+                        k; // ugly way to access each node associated with an individual
+            ret = tsk_treeseq_get_node(&ts, node_id, &node_obj);
+            assert(ret == 0);
+            assert(node_obj.individual == (tsk_id_t) j);
+            ind->time = node_obj.time;
+            ind->nodes[k] = node_obj.id;
+            sample_flag = node_obj.flags & TSK_NODE_IS_SAMPLE;
+        }
+
+        if (sample_flag) {
+            self->pedigree->samples[sample_num] = ind;
+            sample_num++;
         }
     }
-    /* Every individual should have ploidy nodes associated with it in the
-     * tables. We get the time from there - we're not really doing any error
-     * checking here, just getting it sort-of working for now. */
-    for (j = 0; j < self->tables->nodes.num_rows; j++) {
-        ret = tsk_node_table_get_row(&self->tables->nodes, (tsk_id_t) j, &node);
-        tsk_bug_assert(ret == 0);
-        tsk_bug_assert(node.individual != TSK_NULL);
-        is_sample[node.individual] = !!(node.flags & TSK_NODE_IS_SAMPLE);
-        time[node.individual] = node.time;
-    }
-    /* JK: This next method isn't really needed, I'm just calling into
-     * the existing code to minimise changes at this point. The pedigree
-     * code needs a full going over to take the new table-based approach
-     * into account */
-    ret = msp_set_pedigree(self, parents, time, is_sample);
-    if (ret != 0) {
-        goto out;
-    }
+    self->pedigree->num_samples = sample_num;
     ret = msp_set_simulation_model(self, MSP_MODEL_WF_PED);
 out:
-    msp_safe_free(parents);
-    msp_safe_free(time);
-    msp_safe_free(is_sample);
+    tsk_treeseq_free(&ts);
     return ret;
 }
 

lib/msprime.h

@@ -121,6 +121,7 @@ typedef struct individual_t_t {
     int sex;
     double time;
     bool queued;
+    tsk_id_t *nodes;
     // For debugging, to ensure we only merge once.
     bool merged;
 } individual_t;

lib/tests/test_pedigrees.c

@@ -176,15 +176,15 @@ test_pedigree_errors(void)
     parents[0] = -2;
     ret = build_pedigree_sim(
         &msp, &tables, rng, 100, ploidy, num_inds, parents, time, is_sample);
-    CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_INDIVIDUAL_OUT_OF_BOUNDS);
+    CU_ASSERT_EQUAL_FATAL(ret, msp_set_tsk_error(TSK_ERR_INDIVIDUAL_OUT_OF_BOUNDS));
     ret = msp_initialise(&msp);
     tsk_table_collection_free(&tables);
     msp_free(&msp);
 
     parents[0] = 100;
     ret = build_pedigree_sim(
         &msp, &tables, rng, 100, ploidy, num_inds, parents, time, is_sample);
-    CU_ASSERT_EQUAL_FATAL(ret, TSK_ERR_INDIVIDUAL_OUT_OF_BOUNDS);
+    CU_ASSERT_EQUAL_FATAL(ret, msp_set_tsk_error(TSK_ERR_INDIVIDUAL_OUT_OF_BOUNDS));
     ret = msp_initialise(&msp);
     tsk_table_collection_free(&tables);
     msp_free(&msp);